@@ -118,6 +118,7 @@ pollTimeoutMs : Int<Ms>
118pollTimeoutMs = 25000 -- long-poll heartbeat interval (how long a poll blocks before returning empty)
118pollTimeoutMs = 25000 -- long-poll heartbeat interval (how long a poll blocks before returning empty)
119maxPollsPerKey = 3 -- max concurrent long-poll connections per pubkey
119maxPollsPerKey = 3 -- max concurrent long-poll connections per pubkey
120maxActivePolls = 5000
120maxActivePolls = 5000
121presenceRefreshInterval : Int<Ms>
122presenceRefreshInterval : Int<Ms>
122presenceRefreshInterval = 300000 -- 5 minutes
123presenceRefreshInterval = 300000 -- 5 minutes
123presenceRefreshThreshold : Int<Ms>
124presenceRefreshThreshold : Int<Ms>
@@ -223,7 +224,7 @@ type Message = {messageId: MessageIdHex, toKey: PubkeyHex, encryptedBlob: BlobHe
223-- Presence: tracks that a pubkey is online at a given server, with expiry
224-- Presence: tracks that a pubkey is online at a given server, with expiry
224type PresenceEntry = {pubkey: PubkeyHex, server: AnyServerName, expiresAt: Timestamp, timestamp: Timestamp}
225type PresenceEntry = {pubkey: PubkeyHex, server: AnyServerName, expiresAt: Timestamp, timestamp: Timestamp}
225-- Queued federation forward: message waiting to be delivered to a remote server
226-- 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}
227-- Tracks message IDs that were acked by a polling client, pending federation ack back to origin
228-- Tracks message IDs that were acked by a polling client, pending federation ack back to origin
228type AckedEntry = {messageId: MessageIdHex, toKey: PubkeyHex, ackedAt: Timestamp, originServer: ServerName}
229type AckedEntry = {messageId: MessageIdHex, toKey: PubkeyHex, ackedAt: Timestamp, originServer: ServerName}
229-- Terminal state: message delivered and ack confirmed (used for dedup after ack overflow)
230-- Terminal state: message delivered and ack confirmed (used for dedup after ack overflow)
@@ -424,8 +425,8 @@ route AuthApi where
424-- Client-facing API (requires session token in Authorization header)
425-- Client-facing API (requires session token in Authorization header)
425route ClientApi where
426route ClientApi where
426 -- Long-poll for new messages; cursor tracks last-seen receivedAt timestamp
427 -- Long-poll for new messages; cursor tracks last-seen receivedAt timestamp
427− POST {cursor: Timestamp, generation: Int} /poll headers {authorization: Text, clientIp: Text}
428+ POST {cursor: Timestamp, generation: Int, watchKeys: [PubkeyHex]} /poll headers {authorization: Text, clientIp: Text}
428− -> {events: [PollEvent], cursor: Timestamp, authorized: Bool, rateLimited: Bool, draining: Bool, generation: Int}
429+ -> {events: [PollEvent], cursor: Timestamp, authorized: Bool, rateLimited: Bool, draining: Bool, generation: Int, presenceUpdates: [{pubkey: PubkeyHex, online: Bool}]}
430 /messages
431 /messages
431 -- Send an E2E encrypted message to a recipient pubkey
432 -- Send an E2E encrypted message to a recipient pubkey
@@ -956,6 +957,34 @@ transformAllMessageBuckets = \f -> do
956 b14 <- *messageBucket14; set *messageBucket14 = f b14
957 b14 <- *messageBucket14; set *messageBucket14 = f b14
957 b15 <- *messageBucket15; set *messageBucket15 = f b15
958 b15 <- *messageBucket15; set *messageBucket15 = f b15
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)
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
982+-- Collect unique bucket indices for a set of pubkeys
983+uniqueBucketsFor = \items getKey ->
985+ let b = pubkeyBucket (getKey x) in
986+ if any acc (\i -> i == b) then acc else union acc [b]) [] items
959-- Prune a single message bucket and return kept count + per-recipient counts.
988-- Prune a single message bucket and return kept count + per-recipient counts.
960-- Since messages are sharded by recipient pubkey, each bucket's recipient counts
989-- Since messages are sharded by recipient pubkey, each bucket's recipient counts
961-- are disjoint from other buckets — simple union merges them.
990-- are disjoint from other buckets — simple union merges them.
@@ -982,13 +1011,22 @@ adjustRecipientCount = \counts pubkey delta ->
982 else if delta > 0 then union counts [{pubkey: pubkey, count: delta}]
1011 else if delta > 0 then union counts [{pubkey: pubkey, count: delta}]
983 else counts
1012 else counts
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)}
985recomputeRecipientCounts : [Message] -> [RecipientCount]
1022recomputeRecipientCounts : [Message] -> [RecipientCount]
986−recomputeRecipientCounts = \messages ->
1023+recomputeRecipientCounts = countByRecipient
987− fold (\acc m ->
1024+
988− if any acc (\c -> c.pubkey == m.toKey)
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 ->
989− then do c <- acc; yield (if c.pubkey == m.toKey then {pubkey: c.pubkey, count: c.count + 1} else c)
1029+ do c <- counts; yield {pubkey: c.pubkey, count: max 0 (c.count - getRecipientCount c.pubkey removedCounts)}
990− else union acc [{pubkey: m.toKey, count: 1}]
993-- Check message dedup only (without counting — count comes from recipientMsgCounts)
1031-- Check message dedup only (without counting — count comes from recipientMsgCounts)
994checkMessageDup : [Message] -> Text -> PubkeyHex -> Bool
1032checkMessageDup : [Message] -> Text -> PubkeyHex -> Bool
@@ -1521,7 +1559,7 @@ retryForwardsForServer = \entries remaining ->
1521 let rest = do x <- entries; where x.messageId != f.messageId; yield x
1559 let rest = do x <- entries; where x.messageId != f.messageId; yield x
1522 result <- fetchFederation f.toServer (ForwardMessage {
1560 result <- fetchFederation f.toServer (ForwardMessage {
1523 toKey: f.toKey, encryptedBlob: f.encryptedBlob,
1561 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,
1525 fromServer: serverHostN, authorization: federationAuthForServer f.toServer})
1563 fromServer: serverHostN, authorization: federationAuthForServer f.toServer})
1526 case result of
1564 case result of
1527 Err e -> do
1565 Err e -> do
@@ -1663,6 +1701,19 @@ retryOverflowAcks = \overflow -> do
1663 ackOverflow <- *ackOverflow
1701 ackOverflow <- *ackOverflow
1664 set *ackOverflow = do a <- ackOverflow; where not (a.originServer == server && hasMessage a.messageId a.toKey batch); yield a))
1702 set *ackOverflow = do a <- ackOverflow; where not (a.originServer == server && hasMessage a.messageId a.toKey batch); yield a))
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
1708+ if t >= deadline then println "drain: timed out waiting for in-flight federation tasks"
1710+ inFlight <- *federationInFlight
1711+ let activeCount = fold (\acc r -> acc + r.count) 0 inFlight
1712+ if activeCount <= 0 then yield {}
1715+ waitForInFlightDrain deadline
1666-- Full maintenance cycle: prune expired state, retry pending forwards and acks,
1717-- Full maintenance cycle: prune expired state, retry pending forwards and acks,
1667-- retry overflow acks, broadcast offline gossip, and refresh presence for local keys
1718-- retry overflow acks, broadcast offline gossip, and refresh presence for local keys
1668-- nearing expiry. Called by both the background task and the /maintenance/prune endpoint.
1719-- nearing expiry. Called by both the background task and the /maintenance/prune endpoint.
@@ -1929,17 +1980,21 @@ federatedAckEntries = \ackedMessages t ->
1929forwardsExcluding = \forwards server messageIds ->
1980forwardsExcluding = \forwards server messageIds ->
1930 do f <- forwards; where not (f.toServer == server && any messageIds (\mid -> f.messageId == mid)); yield f
1981 do f <- forwards; where not (f.toServer == server && any messageIds (\mid -> f.messageId == mid)); yield f
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).
1933forwardCandidates = \targets messages forwards t ->
1987forwardCandidates = \targets messages forwards t ->
1934 fold (\acc target ->
1988 fold (\acc target ->
1935 let serverFwdIds = do f <- forwards; where f.toServer == target.server; yield f.messageId in
1989 let serverFwdIds = do f <- forwards; where f.toServer == target.server; yield f.messageId in
1936 union acc (do
1990 union acc (do
1937 m <- messages
1991 m <- messages
1938 where m.toKey == target.pubkey
1992 where m.toKey == target.pubkey
1939− where m.originServer == "" -- only forward locally-originated messages
1940 where none serverFwdIds (\fid -> fid == m.messageId)
1993 where none serverFwdIds (\fid -> fid == m.messageId)
1994+ let origin = if m.originServer == "" then serverHostN else m.originServer
1941 yield {messageId: m.messageId, toServer: target.server, createdAt: t,
1995 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
1945challengeMatch = \challenges pubkey challenge ip t ->
2000challengeMatch = \challenges pubkey challenge ip t ->
@@ -2074,13 +2129,16 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
2074 validation <- atomic do
2129 validation <- atomic do
2075 challenges <- *challenges
2130 challenges <- *challenges
2076 let hasMatch = challengeMatch challenges pubkey challenge clientIp t
2131 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;
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
2078− -- not consuming on IP mismatch prevents remote attackers from burning challenges).
2135+ -- brute-force attempts; not consuming on IP mismatch prevents remote attackers
2079− -- Other challenges for the same pubkey are preserved (concurrent auth from multiple devices).
2136+ -- from burning challenges.
2080− when hasMatch (set *challenges = challengesExcluding challenges pubkey challenge)
2081 if not hasMatch || not sigValid
2137 if not hasMatch || not sigValid
2082 then yield {valid: False {}}
2138 then yield {valid: False {}}
2140+ set *challenges = challengesExcluding challenges pubkey challenge
2083− else yield {valid: True {}}
2141+ yield {valid: True {}}
2084 if validation.valid
2142 if validation.valid
2086 token <- genToken
2144 token <- genToken
@@ -2092,36 +2150,52 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
2092 -- Read all session buckets for per-key enforcement (revokeOthers, capPerKey)
2150 -- Read all session buckets for per-key enforcement (revokeOthers, capPerKey)
2093 allSessions <- readAllSessions
2151 allSessions <- readAllSessions
2094 presence <- *presence
2152 presence <- *presence
2095− let revokedCount = if revokeOthers then count (do s <- allSessions; where s.pubkey == pubkey; yield s) else 0
2153+ let sessionsForKey = count (do s <- allSessions; where s.pubkey == pubkey; yield s)
2096− let afterRevoke = if revokeOthers then do s <- allSessions; where s.pubkey != pubkey; yield s else allSessions
2154+ let totalSessions = count allSessions
2097− let capped = capPerKeyAndGlobal afterRevoke (\s -> s.pubkey == pubkey) maxSessionsPerKey (\s -> s.expiresAt) t maxSessions
2098 let newSession = {token: token, pubkey: pubkey, expiresAt: expiry, ip: clientIp}
2155 let newSession = {token: token, pubkey: pubkey, expiresAt: expiry, ip: clientIp}
2099− -- Redistribute capped sessions + new session across 16 buckets
2156+ -- Fast path: no revocation, under per-key and global limits — only write one bucket
2157+ if not revokeOthers && sessionsForKey < maxSessionsPerKey && totalSessions < maxSessions
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}})
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
2100− let sessionsWithNew = union capped [newSession]
2174+ let sessionsWithNew = union capped [newSession]
2101− let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s
2175+ let bFor = \b -> do s <- sessionsWithNew; where tokenBucket s.token == b; yield s
2102− writeSessionBucket 0 (bFor 0)
2176+ writeSessionBucket 0 (bFor 0)
2103− writeSessionBucket 1 (bFor 1)
2177+ writeSessionBucket 1 (bFor 1)
2104− writeSessionBucket 2 (bFor 2)
2178+ writeSessionBucket 2 (bFor 2)
2105− writeSessionBucket 3 (bFor 3)
2179+ writeSessionBucket 3 (bFor 3)
2106− writeSessionBucket 4 (bFor 4)
2180+ writeSessionBucket 4 (bFor 4)
2107− writeSessionBucket 5 (bFor 5)
2181+ writeSessionBucket 5 (bFor 5)
2108− writeSessionBucket 6 (bFor 6)
2182+ writeSessionBucket 6 (bFor 6)
2109− writeSessionBucket 7 (bFor 7)
2183+ writeSessionBucket 7 (bFor 7)
2110− writeSessionBucket 8 (bFor 8)
2184+ writeSessionBucket 8 (bFor 8)
2111− writeSessionBucket 9 (bFor 9)
2185+ writeSessionBucket 9 (bFor 9)
2112− writeSessionBucket 10 (bFor 10)
2186+ writeSessionBucket 10 (bFor 10)
2113− writeSessionBucket 11 (bFor 11)
2187+ writeSessionBucket 11 (bFor 11)
2114− writeSessionBucket 12 (bFor 12)
2188+ writeSessionBucket 12 (bFor 12)
2115− writeSessionBucket 13 (bFor 13)
2189+ writeSessionBucket 13 (bFor 13)
2116− writeSessionBucket 14 (bFor 14)
2190+ writeSessionBucket 14 (bFor 14)
2117− writeSessionBucket 15 (bFor 15)
2191+ writeSessionBucket 15 (bFor 15)
2118− let presenceWithout = presenceWithoutLocalKey presence pubkey
2192+ let presenceWithout = presenceWithoutLocalKey presence pubkey
2119− let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
2193+ let presenceEvicted = capGlobal presenceWithout maxPresenceEntries (\p -> p.expiresAt)
2120− let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
2194+ let updatedPresence = union presenceEvicted [{pubkey: pubkey, server: serverHostN, expiresAt: expiry, timestamp: t}]
2121− set *presence = updatedPresence
2195+ set *presence = updatedPresence
2122− yield (Just {value: {peerServers: if federationEnabled
2196+ yield (Just {value: {peerServers: if federationEnabled
2123− then collectPeerServers updatedPresence (\p -> p.server)
2197+ then collectPeerServers updatedPresence (\p -> p.server)
2124− else [], revokedCount: revokedCount}})
2198+ else [], revokedCount: revokedCount}})
2125 case authResult of
2199 case authResult of
2126 Nothing {} -> yield (respond {token: "", expiresAt: 0})
2200 Nothing {} -> yield (respond {token: "", expiresAt: 0})
2127 Just {value: result} -> do
2201 Just {value: result} -> do
@@ -2135,23 +2209,31 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
2135-- POST /poll: Long-poll for new messages. Cursor-based: messages at or before cursor
2209-- POST /poll: Long-poll for new messages. Cursor-based: messages at or before cursor
2136-- are acked (removed from mailbox). Blocks until new messages arrive or heartbeat fires.
2210-- are acked (removed from mailbox). Blocks until new messages arrive or heartbeat fires.
2137-- Acks for federated messages are sent to origin servers in the background.
2211-- 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
2139 t <- now
2213 t <- now
2140 serverGen <- *serverGeneration
2214 serverGen <- *serverGeneration
2141 let effectiveCursor = if generation == serverGen then cursor else 0
2215 let effectiveCursor = if generation == serverGen then cursor else 0
2216+ let cappedWatchKeys = takeRelation maxWatchKeys watchKeys
2142 draining <- *draining
2217 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: []})
2145 globalRateLimited <- checkGlobalRate t
2220 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: []})
2148 result <- case extractToken authorization of
2223 result <- case extractToken authorization of
2149 Nothing {} -> yield {authorized: False {}, events: [], cursor: 0, pendingAcks: [], evictedEntries: []}
2224 Nothing {} -> yield {authorized: False {}, events: [], cursor: 0, pendingAcks: [], evictedEntries: []}
2150 Just {value: token} -> pollOnce token effectiveCursor clientIp
2225 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
2230+ presence <- *presence
2231+ yield (do key <- cappedWatchKeys; yield {pubkey: key, online: count (presenceForKey presence key t2) > 0})
2151 fork (when federationEnabled (do
2233 fork (when federationEnabled (do
2152 sendAcksAndClean result.pendingAcks
2234 sendAcksAndClean result.pendingAcks
2153 retryEvictedAcks result.evictedEntries result.pendingAcks))
2235 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})
2156-- POST /messages: Send an encrypted message to a recipient.
2238-- POST /messages: Send an encrypted message to a recipient.
2157-- Deduplicates by messageId across messages, ackedIds, and deliveredIds.
2239-- Deduplicates by messageId across messages, ackedIds, and deliveredIds.
@@ -2228,7 +2310,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp
2228 -- Create forward entries atomically with message storage (prevents race with ack)
2310 -- Create forward entries atomically with message storage (prevents race with ack)
2229 when (federationEnabled && count remoteServers > 0) (do
2311 when (federationEnabled && count remoteServers > 0) (do
2230 forwards <- *forwards
2312 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}
2232 let cappedFwds = takeRelation (maxForwardsEntries - count forwards) newFwds
2314 let cappedFwds = takeRelation (maxForwardsEntries - count forwards) newFwds
2233 set *forwards = union forwards cappedFwds)
2315 set *forwards = union forwards cappedFwds)
2234 yield {send: SendDone {delivery: delivery, remoteServers: remoteServers}, senderRateLimited: False {}, globalCapReached: False {}})
2316 yield {send: SendDone {delivery: delivery, remoteServers: remoteServers}, senderRateLimited: False {}, globalCapReached: False {}})
@@ -2368,11 +2450,15 @@ handleRecvGossip = \events hopsRemaining fromServer authorization respond -> do
2368 yield (respond {ok: True {}, error: ""}))
2450 yield (respond {ok: True {}, error: ""}))
2370-- POST /federation/forward: Receive a message forwarded from a peer server.
2452-- POST /federation/forward: Receive a message forwarded from a peer server.
2371−-- Only accepts messages where fromServer == originServer (no cascade forwarding).
2453+-- Accepts both direct forwards (fromServer == originServer) and relay forwards
2372−-- Deduplicates and stores; rejects if recipient mailbox is full or recv rate exceeded.
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).
2373handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fromServer authorization respond -> do
2457handleForwardMessage = \toKey encryptedBlob timestamp originServer messageId fromServer authorization respond -> do
2374 t <- now
2458 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))
2376 -- Replay protection: reject forwards with timestamps older than messageMaxAge.
2462 -- Replay protection: reject forwards with timestamps older than messageMaxAge.
2377 -- Legitimate queued messages have timestamps within this window. Replayed captures
2463 -- Legitimate queued messages have timestamps within this window. Replayed captures
2378 -- become stale once dedup state (messages + ackedIds + deliveredIds) is pruned (~30 days).
2464 -- become stale once dedup state (messages + ackedIds + deliveredIds) is pruned (~30 days).
@@ -2441,8 +2527,9 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
2441 let ackedForwards = do f <- forwards; where f.toServer == fromServer && any dedupedIds (\mid -> f.messageId == mid); yield f
2527 let ackedForwards = do f <- forwards; where f.toServer == fromServer && any dedupedIds (\mid -> f.messageId == mid); yield f
2442 let updatedForwards = forwardsExcluding forwards fromServer dedupedIds
2528 let updatedForwards = forwardsExcluding forwards fromServer dedupedIds
2443 set *forwards = updatedForwards
2529 set *forwards = updatedForwards
2444− -- Read all message buckets (cold path: federation ack handles at most maxMessageIds entries)
2530+ -- Read only message buckets containing acked recipients (not all 16)
2445− allMessages <- readAllMessages
2531+ let neededBuckets = uniqueBucketsFor ackedForwards (\f -> f.toKey)
2532+ relevantMessages <- readSelectedBuckets neededBuckets
2446 presence <- *presence
2533 presence <- *presence
2447 let removable = do
2534 let removable = do
2448 f <- ackedForwards
2535 f <- ackedForwards
@@ -2450,18 +2537,18 @@ handleFederationAck = \messageIds originServer fromServer authorization respond
2450 where none presence (\p -> p.pubkey == f.toKey && isLocalServer p.server && p.expiresAt > t)
2537 where none presence (\p -> p.pubkey == f.toKey && isLocalServer p.server && p.expiresAt > t)
2452 -- Only remove messages stored long enough for local clients to have polled them (grace period)
2539 -- 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
2454 let removed = count pastGrace
2541 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
2456 when (removed > 0) (do
2543 when (removed > 0) (do
2457− transformAllMessageBuckets (\bucket ->
2544+ transformSelectedBuckets neededBuckets (\bucket ->
2458 do m <- bucket; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m)
2545 do m <- bucket; where none pastGrace (\f -> m.messageId == f.messageId && m.toKey == f.toKey); yield m)
2459 totalMsgCount <- *totalMessageCount
2546 totalMsgCount <- *totalMessageCount
2460 set *totalMessageCount = totalMsgCount - removed
2547 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)
2462 recipientCounts <- *recipientMsgCounts
2549 recipientCounts <- *recipientMsgCounts
2463− let adjustedCounts = fold (\acc f -> adjustRecipientCount acc f.toKey (0 - 1)) recipientCounts pastGrace
2550+ let removedCounts = countByRecipient pastGrace
2464− set *recipientMsgCounts = adjustedCounts)
2551+ set *recipientMsgCounts = batchDecrementRecipientCounts recipientCounts removedCounts)
2465 yield (respond {ok: True {}, error: ""}))
2552 yield (respond {ok: True {}, error: ""}))
2467-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
2554-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
@@ -2485,12 +2572,16 @@ handleDrain = \authorization respond -> do
2485 let pendingForwardCount = count forwards
2572 let pendingForwardCount = count forwards
2486 let pendingAckCount = count ackedIds
2573 let pendingAckCount = count ackedIds
2487 -- Fork gossip so drain responds immediately (hung peers can't block the response).
2574 -- Fork gossip so drain responds immediately (hung peers can't block the response).
2575+ -- After gossip, wait for any in-flight federation requests (from prior forks) to
2488− -- Operator can poll /stats for drainCompletedAt > 0 to know when it's safe to stop.
2576+ -- complete before signaling drainCompletedAt. Operator polls /stats for drainCompletedAt > 0
2577+ -- to know when it's truly safe to stop the process.
2489 atomic do set *drainCompletedAt = 0
2578 atomic do set *drainCompletedAt = 0
2490 fork (do
2579 fork (do
2491 when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do
2580 when (federationEnabled && drainedCount > 0 && count drainResult.peerServers > 0) (do
2492 let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}
2581 let gossipEvents = do p <- drainResult.localEntries; yield {eventType: GossipOffline {}, pubkey: p.pubkey, server: serverHostN, timestamp: t, ttl: offlineGossipMinTtl}
2493 sendBatchedGossip drainResult.peerServers gossipEvents gossipOriginHops)
2582 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))
2494 t2 <- now
2585 t2 <- now
2495 atomic do set *drainCompletedAt = t2)
2586 atomic do set *drainCompletedAt = t2)
2496 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")
2587 println ("drain: broadcast offline for " ++ show drainedCount ++ " local keys")
@@ -2643,7 +2734,7 @@ serve = \req -> case req of
2643 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond
2734 ReqChallenge {pubkey, clientIp, respond} -> handleChallenge pubkey clientIp respond
2644 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond
2735 VerifyAuth {pubkey, challenge, signature, revokeOthers, clientIp, respond} -> handleVerifyAuth pubkey challenge signature revokeOthers clientIp respond
2645 RevokeSession {targetToken, authorization, clientIp, respond} -> handleRevoke targetToken authorization clientIp respond
2736 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
2647 SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond
2738 SendMessage {to, encryptedBlob, timestamp, messageId, authorization, clientIp, respond} -> handleSendMessage to encryptedBlob timestamp messageId authorization clientIp respond
2648 LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond
2739 LookupPresence {pubkey, authorization, clientIp, respond} -> handleLookupPresence pubkey authorization clientIp respond
2649 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond
2740 RecvGossip {events, hopsRemaining, fromServer, authorization, respond} -> handleRecvGossip events hopsRemaining fromServer authorization respond
@@ -2682,8 +2773,8 @@ backgroundPrune = do
2682 set *lastHeartbeatAt = t)
2773 set *lastHeartbeatAt = t)
2683 status <- atomic do
2774 status <- atomic do
2684 running <- *pruneRunning
2775 running <- *pruneRunning
2685− if running > 0 && t - running < backgroundPruneInterval * 3
2776+ if running > 0 && t - running < backgroundPruneInterval * 2
2686− then yield {canRun: False {}, recovered: False {}, gen: 0}
2777+ then yield {canRun: False {}, recovered: False {}, gen: 0, stuckSince: 0}
2688 set *pruneRunning = t
2779 set *pruneRunning = t
2689 gen <- *pruneGeneration
2780 gen <- *pruneGeneration
@@ -2691,9 +2782,9 @@ backgroundPrune = do
2691 set *pruneGeneration = newGen
2782 set *pruneGeneration = newGen
2692 pc <- *pruneCount
2783 pc <- *pruneCount
2693 set *pruneCount = pc + 1
2784 set *pruneCount = pc + 1
2694− yield {canRun: True {}, recovered: running > 0, gen: newGen}
2785+ yield {canRun: True {}, recovered: running > 0, gen: newGen, stuckSince: running}
2695 when status.recovered
2786 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"))
2697 if status.canRun
2788 if status.canRun
2698 then fork (do
2789 then fork (do
2699 pruneAndRetry t
2790 pruneAndRetry t