@@ -118,6 +118,7 @@ pollTimeoutMs : Int<Ms>
118118pollTimeoutMs = 25000 -- long-poll heartbeat interval (how long a poll blocks before returning empty)
119119maxPollsPerKey = 3 -- max concurrent long-poll connections per pubkey
120120maxActivePolls = 5000
121+maxWatchKeys = 50
121122presenceRefreshInterval : Int<Ms>
122123presenceRefreshInterval = 300000 -- 5 minutes
123124presenceRefreshThreshold : Int<Ms>
@@ -223,7 +224,7 @@ type Message = {messageId: MessageIdHex, toKey: PubkeyHex, encryptedBlob: BlobHe
223224-- Presence: tracks that a pubkey is online at a given server, with expiry
224225type PresenceEntry = {pubkey: PubkeyHex, server: AnyServerName, expiresAt: Timestamp, timestamp: Timestamp}
225226-- Queued federation forward: message waiting to be delivered to a remote server
226−type ForwardEntry = {messageId: MessageIdHex, toServer: ServerName, createdAt: Timestamp, toKey: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp}
227+type ForwardEntry = {messageId: MessageIdHex, toServer: ServerName, createdAt: Timestamp, toKey: PubkeyHex, encryptedBlob: BlobHex, timestamp: Timestamp, originServer: Text}
227228-- Tracks message IDs that were acked by a polling client, pending federation ack back to origin
228229type AckedEntry = {messageId: MessageIdHex, toKey: PubkeyHex, ackedAt: Timestamp, originServer: ServerName}
229230-- Terminal state: message delivered and ack confirmed (used for dedup after ack overflow)
@@ -424,8 +425,8 @@ route AuthApi where
424425-- Client-facing API (requires session token in Authorization header)
425426route ClientApi where
426427 -- Long-poll for new messages; cursor tracks last-seen receivedAt timestamp
427− POST {cursor: Timestamp, generation: Int} /poll headers {authorization: Text, clientIp: Text}
428− -> {events: [PollEvent], cursor: Timestamp, authorized: Bool, rateLimited: Bool, draining: Bool, generation: Int}
428+ POST {cursor: Timestamp, generation: Int, watchKeys: [PubkeyHex]} /poll headers {authorization: Text, clientIp: Text}
429+ -> {events: [PollEvent], cursor: Timestamp, authorized: Bool, rateLimited: Bool, draining: Bool, generation: Int, presenceUpdates: [{pubkey: PubkeyHex, online: Bool}]}
429430 = Poll
430431 /messages
431432 -- Send an E2E encrypted message to a recipient pubkey
@@ -956,6 +957,34 @@ transformAllMessageBuckets = \f -> do
956957 b14 <- *messageBucket14; set *messageBucket14 = f b14
957958 b15 <- *messageBucket15; set *messageBucket15 = f b15
958959
960+-- Read only the message buckets at the given indices (avoids reading all 16)
961+readSelectedBuckets : [Int] -> IO [Message]
962+readSelectedBuckets = \buckets ->
963+ case headOf buckets of
964+ Nothing {} -> yield []
965+ Just {value: b} -> do
966+ let rest = do x <- buckets; where x != b; yield x
967+ msgs <- readMessageBucket b
968+ remaining <- readSelectedBuckets rest
969+ yield (union msgs remaining)
970+
971+-- Transform only the message buckets at the given indices
972+transformSelectedBuckets : [Int] -> ([Message] -> [Message]) -> IO ()
973+transformSelectedBuckets = \buckets f ->
974+ case headOf buckets of
975+ Nothing {} -> yield {}
976+ Just {value: b} -> do
977+ let rest = do x <- buckets; where x != b; yield x
978+ msgs <- readMessageBucket b
979+ writeMessageBucket b (f msgs)
980+ transformSelectedBuckets rest f
981+
982+-- Collect unique bucket indices for a set of pubkeys
983+uniqueBucketsFor = \items getKey ->
984+ fold (\acc x ->
985+ let b = pubkeyBucket (getKey x) in
986+ if any acc (\i -> i == b) then acc else union acc [b]) [] items
987+
959988-- Prune a single message bucket and return kept count + per-recipient counts.
960989-- Since messages are sharded by recipient pubkey, each bucket's recipient counts
961990-- are disjoint from other buckets — simple union merges them.
@@ -982,13 +1011,22 @@ adjustRecipientCount = \counts pubkey delta ->
9821011 else if delta > 0 then union counts [{pubkey: pubkey, count: delta}]
9831012 else counts
9841013
1014+-- Count items per toKey in a single pass: collect unique keys, then count each.
1015+-- Avoids rebuilding the accumulator list on every matched item.
1016+countByRecipient = \items ->
1017+ let uniqueKeys = fold (\acc x ->
1018+ if any acc (\k -> k == x.toKey) then acc
1019+ else union acc [x.toKey]) [] items
1020+ in do k <- uniqueKeys; yield {pubkey: k, count: count (do x <- items; where x.toKey == k; yield x)}
1021+
9851022recomputeRecipientCounts : [Message] -> [RecipientCount]
986−recomputeRecipientCounts = \messages ->
987− fold (\acc m ->
988− if any acc (\c -> c.pubkey == m.toKey)
989− then do c <- acc; yield (if c.pubkey == m.toKey then {pubkey: c.pubkey, count: c.count + 1} else c)
990− else union acc [{pubkey: m.toKey, count: 1}]
991− ) [] messages
1023+recomputeRecipientCounts = countByRecipient
1024+
1025+-- Batch-decrement recipient counts: subtract per-key removal counts in a single pass
1026+-- instead of calling adjustRecipientCount in a loop.
1027+batchDecrementRecipientCounts : [RecipientCount] -> [RecipientCount] -> [RecipientCount]
1028+batchDecrementRecipientCounts = \counts removedCounts ->
1029+ do c <- counts; yield {pubkey: c.pubkey, count: max 0 (c.count - getRecipientCount c.pubkey removedCounts)}
9921030
9931031-- Check message dedup only (without counting — count comes from recipientMsgCounts)
9941032checkMessageDup : [Message] -> Text -> PubkeyHex -> Bool
@@ -1521,7 +1559,7 @@ retryForwardsForServer = \entries remaining ->
15211559 let rest = do x <- entries; where x.messageId != f.messageId; yield x
15221560 result <- fetchFederation f.toServer (ForwardMessage {
15231561 toKey: f.toKey, encryptedBlob: f.encryptedBlob,
1524− timestamp: f.timestamp, originServer: serverHostN, messageId: f.messageId,
1562+ timestamp: f.timestamp, originServer: f.originServer, messageId: f.messageId,
15251563 fromServer: serverHostN, authorization: federationAuthForServer f.toServer})
15261564 case result of
15271565 Err e -> do
@@ -1663,6 +1701,19 @@ retryOverflowAcks = \overflow -> do
16631701 ackOverflow <- *ackOverflow
16641702 set *ackOverflow = do a <- ackOverflow; where not (a.originServer == server && hasMessage a.messageId a.toKey batch); yield a))
16651703
1704+-- Poll federationInFlight until all in-flight requests complete or deadline passes.
1705+-- Used during drain to wait for outstanding federation work before signaling completion.
1706+waitForInFlightDrain = \deadline -> do
1707+ t <- now
1708+ if t >= deadline then println "drain: timed out waiting for in-flight federation tasks"
1709+ else do
1710+ inFlight <- *federationInFlight
1711+ let activeCount = fold (\acc r -> acc + r.count) 0 inFlight
1712+ if activeCount <= 0 then yield {}
1713+ else do
1714+ sleep 1000
1715+ waitForInFlightDrain deadline
1716+
16661717-- Full maintenance cycle: prune expired state, retry pending forwards and acks,
16671718-- retry overflow acks, broadcast offline gossip, and refresh presence for local keys
16681719-- nearing expiry. Called by both the background task and the /maintenance/prune endpoint.
@@ -1929,17 +1980,21 @@ federatedAckEntries = \ackedMessages t ->
19291980forwardsExcluding = \forwards server messageIds ->
19301981 do f <- forwards; where not (f.toServer == server && any messageIds (\mid -> f.messageId == mid)); yield f
19311982
1932−-- Find locally-originated messages that haven't been forwarded yet to each target's server
1983+-- Find messages that haven't been forwarded yet to each target's server.
1984+-- Includes relay: messages received via federation (originServer != "") can be
1985+-- re-forwarded when the recipient moves to a different server. Cascade is bounded
1986+-- by dedup (same messageId won't be forwarded to the same server twice).
19331987forwardCandidates = \targets messages forwards t ->
19341988 fold (\acc target ->
19351989 let serverFwdIds = do f <- forwards; where f.toServer == target.server; yield f.messageId in
19361990 union acc (do
19371991 m <- messages
19381992 where m.toKey == target.pubkey
1939− where m.originServer == "" -- only forward locally-originated messages
19401993 where none serverFwdIds (\fid -> fid == m.messageId)
1994+ let origin = if m.originServer == "" then serverHostN else m.originServer
19411995 yield {messageId: m.messageId, toServer: target.server, createdAt: t,
1942− toKey: m.toKey, encryptedBlob: m.encryptedBlob, timestamp: m.timestamp})) [] targets
1996+ toKey: m.toKey, encryptedBlob: m.encryptedBlob, timestamp: m.timestamp,
1997+ originServer: origin})) [] targets
19431998
19441999
19452000challengeMatch = \challenges pubkey challenge ip t ->
@@ -2074,13 +2129,16 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
20742129 validation <- atomic do
20752130 challenges <- *challenges
20762131 let hasMatch = challengeMatch challenges pubkey challenge clientIp t
2077− -- Consume only the matched challenge on IP match (rate limiter handles brute-force from same IP;
2078− -- not consuming on IP mismatch prevents remote attackers from burning challenges).
2079− -- Other challenges for the same pubkey are preserved (concurrent auth from multiple devices).
2080− when hasMatch (set *challenges = challengesExcluding challenges pubkey challenge)
2132+ -- Only consume the challenge on full success (match AND valid signature).
2133+ -- This prevents an attacker who intercepts the challenge string from burning it
2134+ -- by submitting a bad signature from the same IP. Rate limiter still constrains
2135+ -- brute-force attempts; not consuming on IP mismatch prevents remote attackers
2136+ -- from burning challenges.
20812137 if not hasMatch || not sigValid
20822138 then yield {valid: False {}}
2083− else yield {valid: True {}}
2139+ else do
2140+ set *challenges = challengesExcluding challenges pubkey challenge
2141+ yield {valid: True {}}
20842142 if validation.valid
20852143 then do
20862144 token <- genToken
@@ -2092,36 +2150,52 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
20922150 -- Read all session buckets for per-key enforcement (revokeOthers, capPerKey)
20932151 allSessions <- readAllSessions
20942152 presence <- *presence
2095− let revokedCount = if revokeOthers then count (do s <- allSessions; where s.pubkey == pubkey; yield s) else 0
2096− let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions
2097− let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions
2153+ let sessionsForKey = count (do s <- allSessions; where s.pubkey == pubkey; yield s)
2154+ let totalSessions = count allSessions
20982155 let newSession = {token: token, pubkey: pubkey, expiresAt: expiry, ip: clientIp}
2099− -- Redistribute capped sessions + new session across 16 buckets
2100− let sessionsWithNew = union capped [newSession]
2101− let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s
2102− writeSessionBucket 0 (bFor 0)
2103− writeSessionBucket 1 (bFor 1)
2104− writeSessionBucket 2 (bFor 2)
2105− writeSessionBucket 3 (bFor 3)
2106− writeSessionBucket 4 (bFor 4)
2107− writeSessionBucket 5 (bFor 5)
2108− writeSessionBucket 6 (bFor 6)
2109− writeSessionBucket 7 (bFor 7)
2110− writeSessionBucket 8 (bFor 8)
2111− writeSessionBucket 9 (bFor 9)
2112− writeSessionBucket 10 (bFor 10)
2113− writeSessionBucket 11 (bFor 11)
2114− writeSessionBucket 12 (bFor 12)
2115− writeSessionBucket 13 (bFor 13)
2116− writeSessionBucket 14 (bFor 14)
2117− writeSessionBucket 15 (bFor 15)
2118− let presenceWithout = presenceWithoutLocalKey presence pubkey
2119− let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
2120− let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
2121− set *presence = updatedPresence
2122− yield (Just {value: {peerServers: if federationEnabled
2123− then collectPeerServers updatedPresence (\p -> p.server)
2124− else [], revokedCount: revokedCount}})
2156+ -- Fast path: no revocation, under per-key and global limits — only write one bucket
2157+ if not revokeOthers && sessionsForKey < maxSessionsPerKey && totalSessions < maxSessions
2158+ then do
2159+ let nb = tokenBucket token
2160+ currentBucket <- readSessionBucket nb
2161+ writeSessionBucket nb (union currentBucket [newSession])
2162+ let presenceWithout = presenceWithoutLocalKey presence pubkey
2163+ let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
2164+ let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
2165+ set *presence = updatedPresence
2166+ yield (Just {value: {peerServers: if federationEnabled
2167+ then collectPeerServers updatedPresence (\p -> p.server)
2168+ else [], revokedCount: 0}})
2169+ else do
2170+ -- Slow path: revocation or capacity pressure — redistribute across all buckets
2171+ let revokedCount = if revokeOthers then sessionsForKey else 0
2172+ let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions
2173+ let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions
2174+ let sessionsWithNew = union capped [newSession]
2175+ let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s
2176+ writeSessionBucket 0 (bFor 0)
2177+ writeSessionBucket 1 (bFor 1)
2178+ writeSessionBucket 2 (bFor 2)
2179+ writeSessionBucket 3 (bFor 3)
2180+ writeSessionBucket 4 (bFor 4)
2181+ writeSessionBucket 5 (bFor 5)
2182+ writeSessionBucket 6 (bFor 6)
2183+ writeSessionBucket 7 (bFor 7)
2184+ writeSessionBucket 8 (bFor 8)
2185+ writeSessionBucket 9 (bFor 9)
2186+ writeSessionBucket 10 (bFor 10)
2187+ writeSessionBucket 11 (bFor 11)
2188+ writeSessionBucket 12 (bFor 12)
2189+ writeSessionBucket 13 (bFor 13)
2190+ writeSessionBucket 14 (bFor 14)
2191+ writeSessionBucket 15 (bFor 15)
2192+ let presenceWithout = presenceWithoutLocalKey presence pubkey
2193+ let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
2194+ let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
2195+ set *presence = updatedPresence
2196+ yield (Just {value: {peerServers: if federationEnabled
2197+ then collectPeerServers updatedPresence (\p -> p.server)
2198+ else [], revokedCount: revokedCount}})
21252199 case authResult of
21262200 Nothing {} -> yield (respond {token: "", expiresAt: 0})
21272201 Just {value: result} -> do
@@ -2135,23 +2209,31 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
21352209-- POST /poll: Long-poll for new messages. Cursor-based: messages at or before cursor
21362210-- are acked (removed from mailbox). Blocks until new messages arrive or heartbeat fires.
21372211-- Acks for federated messages are sent to origin servers in the background.
2138−handlePoll = \cursor generation authorization clientIp respond -> do
2212+handlePoll = \cursor generation watchKeys authorization clientIp respond -> do
21392213 t <- now
21402214 serverGen <- *serverGeneration
21412215 let effectiveCursor = if generation == serverGen then cursor else 0
2216+ let cappedWatchKeys = takeRelation maxWatchKeys watchKeys
21422217 draining <- *draining
2143− if draining then yield (respond {events: [], cursor: 0, authorized: True {}, rateLimited: False {}, draining: True {}, generation: serverGen})
2218+ if draining then yield (respond {events: [], cursor: 0, authorized: True {}, rateLimited: False {}, draining: True {}, generation: serverGen, presenceUpdates: []})
21442219 else do
21452220 globalRateLimited <- checkGlobalRate t
2146− if globalRateLimited then yield (respond {events: [], cursor: 0, authorized: False {}, rateLimited: True {}, draining: False {}, generation: serverGen})
2221+ if globalRateLimited then yield (respond {events: [], cursor: 0, authorized: False {}, rateLimited: True {}, draining: False {}, generation: serverGen, presenceUpdates: []})
21472222 else do
21482223 result <- case extractToken authorization of
21492224 Nothing {} -> yield {authorized: False {}, events: [], cursor: 0, pendingAcks: [], evictedEntries: []}
21502225 Just {value: token} -> pollOnce token effectiveCursor clientIp
2226+ -- Check presence for watched keys (piggybacks on poll cycle; max latency = pollTimeoutMs)
2227+ presenceUpdates <- if result.authorized && count cappedWatchKeys > 0 then do
2228+ t2 <- now
2229+ atomic do
2230+ presence <- *presence
2231+ yield (do key <- cappedWatchKeys; yield {pubkey: key, online: count (presenceForKey presence key t2) > 0})
2232+ else yield []
21512233 fork (when federationEnabled (do
21522234 sendAcksAndClean result.pendingAcks
21532235 retryEvictedAcks result.evictedEntries result.pendingAcks))
2154− yield (respond {events: result.events, cursor: result.cursor, authorized: result.authorized, rateLimited: False {}, draining: False {}, generation: serverGen})
2236+ yield (respond {events: result.events, cursor: result.cursor, authorized: result.authorized, rateLimited: False {}, draining: False {}, generation: serverGen, presenceUpdates: presenceUpdates})
21552237
21562238-- POST /messages: Send an encrypted message to a recipient.
21572239-- Deduplicates by messageId across messages, ackedIds, and deliveredIds.
@@ -2228,7 +2310,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp
22282310 -- Create forward entries atomically with message storage (prevents race with ack)
22292311 when (federationEnabled && count remoteServers > 0) (do
22302312 forwards <- *forwards
2231− let newFwds = do s <- remoteServers; where none forwards (\f -> f.messageId == effectiveMsgId && f.toServer == s); yield {messageId: effectiveMsgId, toServer: s, createdAt: t, toKey: to, encryptedBlob: encryptedBlob, timestamp: clampedTs}
2313+ let newFwds = do s <- remoteServers; where none forwards (\f -> f.messageId == effectiveMsgId && f.toServer == s); yield {messageId: effectiveMsgId, toServer: s, createdAt: t, toKey: to, encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: serverHostN}
22322314 let cappedFwds = takeRelation (maxForwardsEntries - count forwards) newFwds
22332315 set *forwards = union forwards cappedFwds)
22342316 yield {send: SendDone {delivery: delivery, remoteServers: remoteServers}, senderRateLimited: False {}, globalCapReached: False {}})
@@ -2368,11 +2450,15 @@ handleRecvGossip = \events hopsRemaining fromServer authorization respond -> do
23682450 yield (respond {ok: True {}, error: ""}))
23692451
23702452-- POST /federation/forward: Receive a message forwarded from a peer server.
2371−-- Only accepts messages where fromServer == originServer (no cascade forwarding).
2372−-- Deduplicates and stores; rejects if recipient mailbox is full or recv rate exceeded.
2453+-- Accepts both direct forwards (fromServer == originServer) and relay forwards
2454+-- (fromServer != originServer) where an intermediate server re-forwards because
2455+-- the recipient moved. Cascade is bounded by dedup (same messageId won't be stored
2456+-- twice) and forwardCandidates (won't re-forward to the same server).
23732457handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fromServer authorization respond -> do
23742458 t <- now
2375− withFederationAuth authorization t respond fromServer (not (isLocalServer originServer) && fromServer == originServer) (do
2459+ withFederationAuth authorization t respond fromServer (not (isLocalServer originServer) && not (isLocalServer fromServer)) (do
2460+ when (fromServer != originServer)
2461+ (println ("forward: relay from " ++ fromServer ++ " (origin: " ++ originServer ++ ") for " ++ toKey))
23762462 -- Replay protection: reject forwards with timestamps older than messageMaxAge.
23772463 -- Legitimate queued messages have timestamps within this window. Replayed captures
23782464 -- become stale once dedup state (messages + ackedIds + deliveredIds) is pruned (~30 days).
@@ -2441,8 +2527,9 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
24412527 let ackedForwards = do f <- forwards; where f.toServer == fromServer && any dedupedIds (\mid -> f.messageId == mid); yield f
24422528 let updatedForwards = forwardsExcluding forwards fromServer dedupedIds
24432529 set *forwards = updatedForwards
2444− -- Read all message buckets (cold path: federation ack handles at most maxMessageIds entries)
2445− allMessages <- readAllMessages
2530+ -- Read only message buckets containing acked recipients (not all 16)
2531+ let neededBuckets = uniqueBucketsFor ackedForwards (\f -> f.toKey)
2532+ relevantMessages <- readSelectedBuckets neededBuckets
24462533 presence <- *presence
24472534 let removable = do
24482535 f <- ackedForwards
@@ -2450,18 +2537,18 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
24502537 where none presence (\p -> p.pubkey == f.toKey && isLocalServer p.server && p.expiresAt > t)
24512538 yield f
24522539 -- Only remove messages stored long enough for local clients to have polled them (grace period)
2453− let pastGrace = do f <- removable; where any allMessages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f
2540+ let pastGrace = do f <- removable; where any relevantMessages (\m -> m.messageId == f.messageId && m.toKey == f.toKey && t - m.receivedAt >= ackGracePeriod); yield f
24542541 let removed = count pastGrace
2455− -- Remove matched messages from their buckets (pastGrace is small, so scanning all 16 buckets is fast)
2542+ -- Remove matched messages from only the affected buckets
24562543 when (removed > 0) (do
2457− transformAllMessageBuckets (\bucket ->
2544+ transformSelectedBuckets neededBuckets (\bucket ->
24582545 do m <- bucket; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m)
24592546 totalMsgCount <- *totalMessageCount
24602547 set *totalMessageCount = totalMsgCount - removed
2461− -- Adjust per-recipient counts for removed messages
2548+ -- Batch-adjust per-recipient counts (single pass instead of per-entry fold)
24622549 recipientCounts <- *recipientMsgCounts
2463− let adjustedCounts = fold (\acc f -> adjustRecipientCount acc f.toKey (0 - 1)) recipientCounts pastGrace
2464− set *recipientMsgCounts = adjustedCounts)
2550+ let removedCounts = countByRecipient pastGrace
2551+ set *recipientMsgCounts = batchDecrementRecipientCounts recipientCounts removedCounts)
24652552 yield (respond {ok: True {}, error: ""}))
24662553
24672554-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
@@ -2485,12 +2572,16 @@ handleDrain = \authorization respond -> do
24852572 let pendingForwardCount = count forwards
24862573 let pendingAckCount = count ackedIds
24872574 -- Fork gossip so drain responds immediately (hung peers can't block the response).
2488− -- Operator can poll /stats for drainCompletedAt > 0 to know when it's safe to stop.
2575+ -- After gossip, wait for any in-flight federation requests (from prior forks) to
2576+ -- complete before signaling drainCompletedAt. Operator polls /stats for drainCompletedAt > 0
2577+ -- to know when it's truly safe to stop the process.
24892578 atomic do set *drainCompletedAt = 0
24902579 fork (do
24912580 when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do
24922581 let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}
24932582 sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops)
2583+ -- Wait for in-flight federation requests to complete (timeout: 2x fetch timeout)
2584+ when federationEnabled (waitForInFlightDrain (t + federationFetchTimeoutMs * 2))
24942585 t2 <- now
24952586 atomic do set *drainCompletedAt = t2)
24962587 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")
@@ -2643,7 +2734,7 @@ serve = \req -> case req of
26432734 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond
26442735 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond
26452736 RevokeSession {targetToken, authorization, clientIp, respond} -> handleRevoke targetToken authorization clientIp respond
2646− Poll {cursor, generation, authorization, clientIp, respond} -> handlePoll cursor generation authorization clientIp respond
2737+ Poll {cursor, generation, watchKeys, authorization, clientIp, respond} -> handlePoll cursor generation watchKeys authorization clientIp respond
26472738 SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond
26482739 LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond
26492740 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond
@@ -2682,8 +2773,8 @@ backgroundPrune = do
26822773 set *lastHeartbeatAt = t)
26832774 status <- atomic do
26842775 running <- *pruneRunning
2685− if running > 0 && t - running < backgroundPruneInterval * 3
2686− then yield {canRun: False {}, recovered: False {}, gen: 0}
2776+ if running > 0 && t - running < backgroundPruneInterval * 2
2777+ then yield {canRun: False {}, recovered: False {}, gen: 0, stuckSince: 0}
26872778 else do
26882779 set *pruneRunning = t
26892780 gen <- *pruneGeneration
@@ -2691,9 +2782,9 @@ backgroundPrune = do
26912782 set *pruneGeneration = newGen
26922783 pc <- *pruneCount
26932784 set *pruneCount = pc + 1
2694− yield {canRun: True {}, recovered: running > 0, gen: newGen}
2785+ yield {canRun: True {}, recovered: running > 0, gen: newGen, stuckSince: running}
26952786 when status.recovered
2696− (println "backgroundPrune: previous cycle appears stuck, forcing new cycle")
2787+ (println ("backgroundPrune: previous cycle stuck for " ++ show (t - status.stuckSince) ++ "ms, forcing recovery"))
26972788 if status.canRun
26982789 then fork (do
26992790 pruneAndRetry t