@@ -3543,10 +3543,11 @@ adaptiveRateScale = \msgCount sessCount presCount ->
35433543-- and scales the rate proportionally. Returns at least 1 to prevent division-by-zero in callers.
35443544adaptiveRate : Int -> IO Int
35453545adaptiveRate = \baseRate -> do
3546− tmc <- *totalMessageCount
3547− sc <- *totalSessionCount
3548− pc <- *totalPresenceCount
3549− let scale = adaptiveRateScale tmc sc pc
3546+ scale <- atomic do
3547+ tmc <- *totalMessageCount
3548+ sc <- *totalSessionCount
3549+ pc <- *totalPresenceCount
3550+ yield (adaptiveRateScale tmc sc pc)
35503551 yield (max 1 (baseRate * scale / 100))
35513552
35523553-- Effective rate limit readers: return runtime override if set (> 0), else compile-time default.
@@ -3760,12 +3761,15 @@ withSessionAuth = \authorization t clientIp onUnauth onAuth -> do
37603761 if not (matchIp clientIp ip)
37613762 then onUnauth
37623763 else do
3763− -- Sliding window: refresh session expiry on every authenticated operation.
3764− -- Without this, users who send/lookup/ack but don't poll would have their
3765− -- session expire after sessionExpiry even though they're actively using it.
3766− -- Pre-filter by pubkey before safeCompare to avoid O(bucket_size) constant-time
3767− -- comparisons (same optimization as reauthenticateSession).
3768− writeSessionBucket bucket (do s <- sessions; yield (if s.pubkey == pubkey && safeCompare s.token token then {token: s.token, pubkey: s.pubkey, expiresAt: t + sessionExpiry, ip: s.ip, createdAt: s.createdAt} else s))
3764+ -- Sliding window: refresh session expiry when within half of sessionExpiry.
3765+ -- Avoids rewriting the session bucket on every authenticated request (reduces
3766+ -- write contention on hot paths like rapid ack/delete sequences). Sessions
3767+ -- still stay alive as long as the client is active within any sessionExpiry window.
3768+ -- Uses pubkey pre-filter to find the matching session's expiry without safeCompare,
3769+ -- since pubkey was already verified by authenticateSession above.
3770+ let needsRefresh = any sessions (\s -> s.pubkey == pubkey && s.expiresAt - t < sessionExpiry / 2)
3771+ when needsRefresh
3772+ (writeSessionBucket bucket (do s <- sessions; yield (if s.pubkey == pubkey && safeCompare s.token token then {token: s.token, pubkey: s.pubkey, expiresAt: t + sessionExpiry, ip: s.ip, createdAt: s.createdAt} else s)))
37693773 onAuth pubkey
37703774
37713775-- Check if a pubkey has no remaining sessions and, if so, clean up local presence.
@@ -4931,17 +4935,20 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
49314935 if count forKey > 0 then union others (evictOneOldest forKey (\s -> s.expiresAt))
49324936 else bucket))
49334937 -- Global cap: if over maxSessions, evict oldest globally.
4934− -- Must re-count after revocation/per-key cap to get accurate total.
4935− currentTotal <- countAllSessions
4936− when (currentTotal >= maxSessions)
4938+ -- Derive count from cached tsc: after revokeOthers, we removed revokedCount sessions;
4939+ -- after per-key cap, we removed at most 1. Avoids a redundant 16-bucket scan.
4940+ let perKeyEvicted = if not revokeOthers && sessionsForKey >= maxSessionsPerKey then 1 else 0
4941+ let estimatedTotal = tsc - revokedCount - perKeyEvicted
4942+ when (estimatedTotal >= maxSessions)
49374943 (do -- Find bucket with oldest session and evict from it
49384944 allSess <- readAllSessions
49394945 let globalCapped = capGlobal allSess maxSessions (\s -> s.expiresAt)
49404946 writeAllSessions globalCapped)
49414947 currentBucket <- readSessionBucket nb
49424948 writeSessionBucket nb (union currentBucket [newSession])
4943− finalCount <- countAllSessions
4944− set *totalSessionCount = finalCount
4949+ -- Final count: estimatedTotal + 1 (new session) - globalEviction (0 or 1)
4950+ let globalEvicted = if estimatedTotal >= maxSessions then 1 else 0
4951+ set *totalSessionCount = estimatedTotal - globalEvicted + 1
49454952 yield (Just {value: {peerServers: peerServers, revokedCount: revokedCount}})
49464953 case authResult of
49474954 Nothing {} -> yield (respond {token: "", expiresAt: 0, error: "draining"})
@@ -5185,6 +5192,9 @@ sendBatchForwardsSequential = \msgs server ->
51855192handleSendMessage = \to encryptedBlob timestamp messageId expiresInMs priority authorization clientIp respond -> do
51865193 let clampedPriority = clamp 0 maxMessagePriority priority
51875194 t <- now
5195+ dp <- *drainPhase
5196+ if dp >= 2 then yield (respond {status: MsDraining {}, messageId: "", retryAfterMs: 0, error: "draining"})
5197+ else do
51885198 globalRateLimited <- checkGlobalRate t
51895199 if globalRateLimited then yield (respond {status: MsRateLimited {}, messageId: "", retryAfterMs: 0, error: "rate_limited"})
51905200 else do
@@ -5198,9 +5208,7 @@ handleSendMessage = \to encryptedBlob timestamp messageId expiresInMs priority a
51985208 (logWarn ("sendMessage: stale timestamp (" ++ show ((t - timestamp) / 86400000) ++ " days old) for " ++ effectiveMsgId))
51995209 let clampedTs = clampTimestamp timestamp t
52005210 result <- atomic do
5201− dp <- *drainPhase
5202− if dp >= 2 then yield {status: MsDraining {}, remoteServers: [], retryAfterMs: 0}
5203− else withSessionAuth authorization t clientIp
5211+ withSessionAuth authorization t clientIp
52045212 (yield {status: MsUnauthorized {}, remoteServers: [], retryAfterMs: 0})
52055213 (\senderKey -> sendMessageCore senderKey to encryptedBlob clampedTs effectiveMsgId expiresInMs clampedPriority t)
52065214 case result.status of
@@ -5749,10 +5757,12 @@ handleFederationAck = \messageIds originServer fromServer requestTimestamp corre
57495757 if ackedCount == 0 && count dedupedIds > 0
57505758 then do
57515759 logWarn ("federation/ack: rejecting ack from " ++ fromServer ++ " — all " ++ show (count dedupedIds) ++ " messageIds unmatched (no forward entries)")
5760+ logEvent "warn" "ack_all_unmatched" (fromServer ++ " " ++ show (count dedupedIds))
57525761 yield (respond {ok: False {}, error: "no_matching_forwards", ackedCount: 0, results: do mid <- dedupedIds; yield {messageId: mid, found: False {}}})
57535762 else do
5754− when (unmatched > 0)
5755− (logWarn ("federation/ack: " ++ fromServer ++ " acked " ++ show unmatched ++ " messageIds with no matching forward entries — skipping cleanup for those"))
5763+ when (unmatched > 0) (do
5764+ logWarn ("federation/ack: " ++ fromServer ++ " acked " ++ show unmatched ++ " messageIds with no matching forward entries — skipping cleanup for those")
5765+ logEvent "warn" "ack_partial_unmatched" (fromServer ++ " " ++ show unmatched ++ "/" ++ show (count dedupedIds)))
57565766 -- Only clear failure record when ALL acked IDs matched forward entries.
57575767 -- Prevents a malicious peer from resetting its backoff by padding ack requests
57585768 -- with known messageIds that no longer have forward entries.
@@ -5921,16 +5931,28 @@ handleFederationBatchForward = \forwards fromServer requestTimestamp correlation
59215931 (logInfo ("federation/forward-batch: stored " ++ show storedCount ++ "/" ++ show (count capped) ++ " messages from " ++ fromServer))
59225932 yield (respond {ok: True {}, error: "", results: results}))
59235933
5924−-- Process each forward entry in a batch sequentially. Each entry runs the same
5925−-- validation and storage logic as handleForwardMessage (presence check, dedup,
5926−-- rate limiting, capacity check, nonce recording).
5934+-- Process forward entries grouped by bucket: all entries sharing a pubkeyBucket(toKey) are
5935+-- processed in a single atomic block, reducing bucket reads from O(entries) to O(buckets).
5936+-- Pre-validation (timestamp, origin) runs outside the atomic block since it reads global state.
59275937forwardBatchInner = \entries fromServer t acc ->
59285938 case headOf entries of
59295939 Nothing {} -> yield acc
5940+ Just {value: first} -> do
5941+ -- Pre-validate all entries: reject stale timestamps and untrusted relay origins.
5942+ -- Reads global state once (knownPeers, blocked, allowed) outside atomic blocks.
5943+ validated <- prevalidateForwardBatch entries fromServer t [] []
5944+ -- Group validated entries by bucket
5945+ forwardBatchByBucket validated.valid fromServer t (union acc validated.rejected)
5946+
5947+-- Pre-validate forward entries: reject stale timestamps and untrusted relay origins.
5948+-- Returns {valid: entries to store, rejected: per-message results for failed entries}.
5949+prevalidateForwardBatch = \entries fromServer t validAcc rejectedAcc ->
5950+ case headOf entries of
5951+ Nothing {} -> yield {valid: validAcc, rejected: rejectedAcc}
59305952 Just {value: entry} -> do
59315953 let rest = do e <- entries; where e.messageId != entry.messageId || e.toKey != entry.toKey; yield e
59325954 if isLocalServer entry.originServer || t - entry.timestamp > messageMaxAge
5933− then forwardBatchInner rest fromServer t (union acc [{messageId: entry.messageId, ok: False {}}])
5955+ then prevalidateForwardBatch rest fromServer t validAcc (union rejectedAcc [{messageId: entry.messageId, ok: False {}}])
59345956 else do
59355957 -- Relay origin verification (same as handleForwardMessage)
59365958 originOk <- if fromServer == entry.originServer then yield (True {})
@@ -5942,24 +5964,43 @@ forwardBatchInner = \entries fromServer t acc ->
59425964 && not (isBlockedPeer entry.originServer blocked)
59435965 && isAllowedByAllowlist entry.originServer allowed)
59445966 if not originOk
5945− then forwardBatchInner rest fromServer t (union acc [{messageId: entry.messageId, ok: False {}}])
5946− else do
5967+ then prevalidateForwardBatch rest fromServer t validAcc (union rejectedAcc [{messageId: entry.messageId, ok: False {}}])
5968+ else prevalidateForwardBatch rest fromServer t (union validAcc [entry]) rejectedAcc
5969+
5970+-- Group validated entries by pubkeyBucket and process each group in one atomic block.
5971+forwardBatchByBucket = \entries fromServer t acc ->
5972+ case headOf entries of
5973+ Nothing {} -> yield acc
5974+ Just {value: first} -> do
5975+ let b = pubkeyBucket first.toKey
5976+ let sameBucket = do e <- entries; where pubkeyBucket e.toKey == b; yield e
5977+ let otherBuckets = do e <- entries; where pubkeyBucket e.toKey != b; yield e
5978+ -- Process the entire bucket group in one atomic block
5979+ results <- atomic do forwardBucketGroup sameBucket fromServer b t []
5980+ forwardBatchByBucket otherBuckets fromServer t (union acc results)
5981+
5982+-- Process all forward entries for a single bucket within one atomic block.
5983+-- Reads presenceBucket/messageBucket/ackedBucket/deliveredBucket/nonceBucket once
5984+-- and applies all forwards, avoiding redundant per-entry bucket reads.
5985+forwardBucketGroup = \entries fromServer b t acc ->
5986+ case headOf entries of
5987+ Nothing {} -> yield acc
5988+ Just {value: entry} -> do
5989+ let rest = do e <- entries; where e.messageId != entry.messageId || e.toKey != entry.toKey; yield e
59475990 let clampedTs = clampTimestamp entry.timestamp t
5948− result <- atomic do
5949− let bkt = pubkeyBucket entry.toKey
5950− presenceBkt <- readPresenceBucket bkt
5951− let hasPresence = any presenceBkt (\p -> p.pubkey == entry.toKey && p.expiresAt > t)
5952− if not hasPresence then yield (False {})
5991+ presenceBkt <- readPresenceBucket b
5992+ let hasPresence = any presenceBkt (\p -> p.pubkey == entry.toKey && p.expiresAt > t)
5993+ result <- if not hasPresence then yield (False {})
59535994 else do
5954− bucket <- readMessageBucket bkt
5955− ackedBkt <- readAckedBucket bkt
5956− delivBkt <- readDelivBucket bkt
5957− nonceBkt <- readNonceBucket bkt
5995+ bucket <- readMessageBucket b
5996+ ackedBkt <- readAckedBucket b
5997+ delivBkt <- readDelivBucket b
5998+ nonceBkt <- readNonceBucket b
59585999 let isDup = checkMessageDup bucket entry.messageId entry.toKey || hasMessage entry.messageId entry.toKey ackedBkt || hasMessage entry.messageId entry.toKey delivBkt || any nonceBkt (\n -> n.messageId == entry.messageId && n.toKey == entry.toKey)
59596000 if isDup then yield (True {})
59606001 else do
59616002 totalMsgCount <- *totalMessageCount
5962− recipientCounts <- readRcBucket bkt
6003+ recipientCounts <- readRcBucket b
59636004 let recipientCount = getRecipientCount entry.toKey recipientCounts
59646005 if totalMsgCount >= maxMessages || recipientCount >= maxMessagesPerRecipient
59656006 then yield (False {})
@@ -5981,19 +6022,19 @@ forwardBatchInner = \entries fromServer t acc ->
59816022 set *seqCounter = seqTs
59826023 let fwdExpiresAt = if entry.expiresAt > 0 && entry.expiresAt > t then min entry.expiresAt (seqTs + unackedMessageMaxAge) else seqTs + unackedMessageMaxAge
59836024 let batchClampedPriority = clamp 0 maxMessagePriority entry.priority
5984− writeMessageBucket bkt (union bucket
6025+ writeMessageBucket b (union bucket
59856026 [{messageId: entry.messageId, toKey: entry.toKey,
59866027 encryptedBlob: entry.encryptedBlob, timestamp: clampedTs, originServer: entry.originServer, receivedAt: seqTs, expiresAt: fwdExpiresAt, senderKey: "", priority: batchClampedPriority}])
59876028 set *totalMessageCount = totalMsgCount + 1
5988− writeRcBucket bkt (adjustRecipientCount recipientCounts entry.toKey 1)
6029+ writeRcBucket b (adjustRecipientCount recipientCounts entry.toKey 1)
59896030 signalMessageForKey entry.toKey seqTs
5990− fwdNonces <- readNonceBucket bkt
6031+ fwdNonces <- readNonceBucket b
59916032 let updatedNonces = union fwdNonces [{messageId: entry.messageId, toKey: entry.toKey, seenAt: t}]
5992− writeNonceBucket bkt (if count updatedNonces > maxNoncesPerBucket
6033+ writeNonceBucket b (if count updatedNonces > maxNoncesPerBucket
59936034 then evictNonces updatedNonces maxNoncesPerBucket forwardNonceRetention t
59946035 else updatedNonces)
59956036 yield (True {})
5996− forwardBatchInner rest fromServer t (union acc [{messageId: entry.messageId, ok: result}])
6037+ forwardBucketGroup rest fromServer b t (union acc [{messageId: entry.messageId, ok: result}])
59976038
59986039-- POST /messages/cancel-batch: Cancel multiple sent messages in one request.
59996040-- Auth checked once; each cancel processed independently. Partial success possible.
@@ -6722,6 +6763,9 @@ handleRestore = \version checksumIn presenceData ackedIdsData ackOverflowData de
67226763-- Rate-limited via global rate limiter; requires session auth.
67236764handleAckMessages = \messageIds authorization clientIp respond -> do
67246765 t <- now
6766+ dp <- *drainPhase
6767+ if dp >= 2 then yield (respond {ok: False {}, authorized: False {}, ackedCount: 0, error: "draining"})
6768+ else do
67256769 globalRateLimited <- checkGlobalRate t
67266770 if globalRateLimited then yield (respond {ok: False {}, authorized: False {}, ackedCount: 0, error: "rate_limited"})
67276771 else
@@ -6967,6 +7011,9 @@ handleCancelMessage = \messageId to authorization clientIp respond -> do
69677011-- Partial success is possible — check per-message status in results.
69687012handleSendBatch = \messages authorization clientIp respond -> do
69697013 t <- now
7014+ dp <- *drainPhase
7015+ if dp >= 2 then yield (respond {results: [], authorized: False {}, retryAfterMs: 0})
7016+ else do
69707017 globalRateLimited <- checkGlobalRate t
69717018 if globalRateLimited then yield (respond {results: [], authorized: False {}, retryAfterMs: 0})
69727019 else
@@ -7650,12 +7697,12 @@ serveOverloaded = \req -> case req of
76507697 CancelBatch {respond} -> respond {results: [], authorized: False {}, error: "overloaded"}
76517698 LookupPresence {respond} -> respond {server: "", authorized: False {}, error: "overloaded"}
76527699 BatchLookupPresence {respond} -> respond {results: [], authorized: False {}, error: "overloaded"}
7653− RecvGossip {respond} -> respond {ok: False {}, error: "overloaded"}
7654− ForwardMessage {respond} -> respond {ok: False {}, error: "overloaded"}
7655− FederationAck {respond} -> respond {ok: False {}, error: "overloaded", ackedCount: 0, results: []}
7656− FederationCancel {respond} -> respond {ok: False {}, error: "overloaded", cancelled: False {}}
7657− FederationBatchCancel {respond} -> respond {ok: False {}, error: "overloaded", results: []}
7658− FederationBatchForward {respond} -> respond {ok: False {}, error: "overloaded", results: []}
7700+ RecvGossip {respond} -> respond {ok: False {}, error: "overloaded:1000ms"}
7701+ ForwardMessage {respond} -> respond {ok: False {}, error: "overloaded:1000ms"}
7702+ FederationAck {respond} -> respond {ok: False {}, error: "overloaded:1000ms", ackedCount: 0, results: []}
7703+ FederationCancel {respond} -> respond {ok: False {}, error: "overloaded:1000ms", cancelled: False {}}
7704+ FederationBatchCancel {respond} -> respond {ok: False {}, error: "overloaded:1000ms", results: []}
7705+ FederationBatchForward {respond} -> respond {ok: False {}, error: "overloaded:1000ms", results: []}
76597706 VerifyToken {respond} -> respond {ok: False {}}
76607707 TopRecipients {respond} -> respond {recipients: [], ok: False {}, total: 0}
76617708 BlockPeer {respond} -> respond {ok: False {}, blocked: 0}
@@ -7881,10 +7928,12 @@ backgroundSnapshot = do
78817928 -- sidecar process can persist it. Only when there's meaningful state to save.
78827929 when (autoSnapshotExportInterval > 0) (do
78837930 lastAutoExport <- *lastAutoSnapshotExportAt
7884− let hasState = result.presence > 0 || result.forwards > 0 || result.ackedIds > 0 || result.deliveredIds > 0
7931+ let hasState = result.presence > 0 || result.forwards > 0 || result.ackedIds > 0 || result.deliveredIds > 0 || result.ackOverflow > 0
78857932 -- Fingerprint: skip export when key counters haven't changed since last export.
78867933 -- Uses multiplicative hash to distinguish {messages:1,sessions:0} from {messages:0,sessions:1}.
7887− let fingerprint = result.messages + result.sessions * 100003 + result.presence * 100019 + result.forwards * 100043 + result.ackedIds * 100049 + result.deliveredIds * 100057
7934+ -- Includes ackOverflow and forwardFailures so state changes in those collections
7935+ -- (e.g., acks moving to overflow during a federation outage) trigger an export.
7936+ let fingerprint = result.messages + result.sessions * 100003 + result.presence * 100019 + result.forwards * 100043 + result.ackedIds * 100049 + result.deliveredIds * 100057 + result.ackOverflow * 100069 + result.forwardFailures * 100081
78887937 lastFp <- *lastSnapshotFingerprint
78897938 let stateChanged = lastFp == 0 || fingerprint != lastFp
78907939 when (hasState && stateChanged && (lastAutoExport == 0 || t - lastAutoExport >= autoSnapshotExportInterval)) (do
@@ -8018,7 +8067,7 @@ startupProbeNext = \peers probed ->
80188067 result <- fetchFederation server (RecvGossip {
80198068 authorization: "",
80208069 events: [], hopsRemaining: 1, fromServer: serverHostN,
8021− requestTimestamp: startTime, correlationId: ""})
8070+ requestTimestamp: startTime, correlationId: "", protocolVersion: federationProtocolVersion})
80228071 endTime <- now
80238072 let latency = endTime - startTime
80248073 case result of