@@ -505,6 +505,8 @@ type SenderExpiration = {senderKey: Text, messageId: MessageIdHex, toKey: Pubkey
505505*fwdBucket14 : [ForwardEntry]
506506*fwdBucket15 : [ForwardEntry]
507507
508+*totalForwardCount : Int
509+
508510-- Messages acked by local recipients, pending federation ack to origin server.
509511-- Sharded into 16 buckets by pubkeyBucket(toKey) for O(n/16) dedup in sendMessageCore.
510512*ackedBucket0 : [AckedEntry]
@@ -2339,8 +2341,10 @@ removeForwardsByServer = \server messageIds b acc ->
23392341 removed <- atomic do
23402342 bucket <- readFwdBucket b
23412343 let matched = do f <- bucket; where f.toServer == server && any messageIds (\mid -> f.messageId == mid); yield f
2342− when (count matched > 0)
2343− (writeFwdBucket b (do f <- bucket; where not (f.toServer == server && any messageIds (\mid -> f.messageId == mid)); yield f))
2344+ when (count matched > 0) (do
2345+ writeFwdBucket b (do f <- bucket; where not (f.toServer == server && any messageIds (\mid -> f.messageId == mid)); yield f)
2346+ tfc <- *totalForwardCount
2347+ set *totalForwardCount = max 0 (tfc - count matched))
23442348 yield matched
23452349 removeForwardsByServer server messageIds (b + 1) (if count removed > 0 then union acc removed else acc)
23462350
@@ -2622,11 +2626,11 @@ dedupMessageIds = \messageIds ->
26222626 else union [mid] acc) [] capped
26232627
26242628
2625−-- Constant-time comparison iteration count. Derived from tokenHexLen (48) + 16
2626−-- headroom. Must be >= the longest secret compared (tokens = tokenHexLen hex
2627−-- chars, maintenance secret = variable). If tokenChunks changes, this adjusts
2628−-- automatically.
2629−safeCompareLen = tokenHexLen + 16
2629+-- Constant-time comparison iteration count. Must be >= the longest secret
2630+-- compared: session tokens are tokenHexLen (48) hex chars, but maintenance
2631+-- secrets can be up to maxTokenLen (256) chars (validated by extractToken).
2632+-- Using maxTokenLen + 16 ensures full comparison for all accepted inputs.
2633+safeCompareLen = maxTokenLen + 16
26302634
26312635-- Constant-time string comparison to prevent timing side-channels.
26322636-- Always iterates safeCompareLen characters regardless of input length;
@@ -2985,8 +2989,10 @@ pruneExpired = \t -> do
29852989 atomic do
29862990 forwards <- readAllForwards
29872991 -- Clear all forwards when federation is disabled (they can never be retried)
2988− writeAllForwards (if not federationEnabled then []
2989− else do f <- forwards; where f.createdAt > retainCutoff; yield f)
2992+ let prunedForwards = if not federationEnabled then []
2993+ else do f <- forwards; where f.createdAt > retainCutoff; yield f
2994+ writeAllForwards prunedForwards
2995+ set *totalForwardCount = count prunedForwards
29902996 -- AckedIds: per-bucket pruning (each bucket in its own atomic block to reduce lock
29912997 -- contention with poll-ack paths that write to ackedIds concurrently).
29922998 pruneAllAckedBuckets 0 retainCutoff
@@ -3342,7 +3348,14 @@ withSessionAuth = \authorization t clientIp onUnauth onAuth -> do
33423348 Authorized {pubkey, ip} ->
33433349 if not (matchIp clientIp ip)
33443350 then onUnauth
3345− else onAuth pubkey
3351+ else do
3352+ -- Sliding window: refresh session expiry on every authenticated operation.
3353+ -- Without this, users who send/lookup/ack but don't poll would have their
3354+ -- session expire after sessionExpiry even though they're actively using it.
3355+ -- Pre-filter by pubkey before safeCompare to avoid O(bucket_size) constant-time
3356+ -- comparisons (same optimization as reauthenticateSession).
3357+ 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))
3358+ onAuth pubkey
33463359
33473360-- Check if a pubkey has no remaining sessions and, if so, clean up local presence.
33483361-- Must be called inside an atomic block after the session has already been removed.
@@ -3381,7 +3394,12 @@ gossipOfflineIfLastSession = \r label -> do
33813394removeForward = \mid srv toKey -> atomic do
33823395 let b = pubkeyBucket toKey
33833396 fwdBkt <- readFwdBucket b
3384− writeFwdBucket b (forwardsExcluding fwdBkt srv [mid])
3397+ let kept = forwardsExcluding fwdBkt srv [mid]
3398+ writeFwdBucket b kept
3399+ let removed = count fwdBkt - count kept
3400+ when (removed > 0) (do
3401+ tfc <- *totalForwardCount
3402+ set *totalForwardCount = max 0 (tfc - removed))
33853403
33863404-- Reclaim origin messages after federation ack confirms delivery.
33873405-- Batched by message bucket: all entries for the same bucket are processed in one
@@ -3476,6 +3494,8 @@ incrementForwardRetries = \messageId toServer toKey -> atomic do
34763494 if shouldDrop
34773495 then do
34783496 writeFwdBucket b (do f <- fwdBkt; where not (f.messageId == messageId && f.toServer == toServer); yield f)
3497+ tfc <- *totalForwardCount
3498+ set *totalForwardCount = max 0 (tfc - 1)
34793499 logWarn ("federation: dropping forward " ++ messageId ++ " to " ++ toServer ++ " after " ++ show maxForwardRetries ++ " retries")
34803500 else writeFwdBucket b (do f <- fwdBkt; yield (if f.messageId == messageId && f.toServer == toServer then {messageId: f.messageId, toServer: f.toServer, createdAt: f.createdAt, toKey: f.toKey, encryptedBlob: f.encryptedBlob, timestamp: f.timestamp, originServer: f.originServer, retries: f.retries + 1} else f))
34813501
@@ -3553,7 +3573,7 @@ forwardQueuedMessages = \targets t -> do
35533573 let neededBuckets = uniqueBucketsFor remoteTargets (\target -> target.pubkey)
35543574 relevantMessages <- readSelectedBuckets neededBuckets
35553575 let candidates = forwardCandidates remoteTargets relevantMessages forwards t
3556− let totalFwdCount <- countAllForwards 0
3576+ totalFwdCount <- *totalForwardCount
35573577 -- Per-key cap: filter candidates to respect maxForwardsPerKey per recipient
35583578 let perKeyCapped = fold (\acc c ->
35593579 let existingForKey = count (do f <- forwards; where f.toKey == c.toKey; yield f)
@@ -3563,6 +3583,7 @@ forwardQueuedMessages = \targets t -> do
35633583 -- Write back only the affected forward buckets (not all 16)
35643584 let combined = union forwards capped
35653585 writeSelectedFwdBuckets neededFwdBuckets combined
3586+ set *totalForwardCount = totalFwdCount + count capped
35663587 yield capped
35673588 let servers = collectPeerServers newForwards (\x -> x.toServer)
35683589 failures <- *forwardFailures
@@ -3585,6 +3606,41 @@ fetchFederationAck = \origin messageIds -> do
35853606 originServer: origin, fromServer: serverHostN,
35863607 requestTimestamp: t, correlationId: corrId})
35873608
3609+-- Per-bucket ack retry increment: scans only the buckets containing entries from the batch.
3610+-- Each bucket is processed in its own atomic block to reduce lock contention.
3611+-- Drops entries exceeding maxAckRetries to deliveredIds. Returns total dropped count.
3612+incrementAckRetriesPerBucket = \bucketIds server batch droppedAcc ->
3613+ case headOf bucketIds of
3614+ Nothing {} -> yield droppedAcc
3615+ Just {value: b} -> do
3616+ let remainingBuckets = do x <- bucketIds; where x != b; yield x
3617+ bucketDropped <- atomic do
3618+ bucket <- readAckedBucket b
3619+ let isTarget = \a -> a.originServer == server && hasMessage a.messageId a.toKey batch
3620+ let toDrop = do a <- bucket; where isTarget a && a.retries + 1 >= maxAckRetries; yield a
3621+ let dropCount = count toDrop
3622+ writeAckedBucket b (do a <- bucket
3623+ where not (isTarget a && a.retries + 1 >= maxAckRetries)
3624+ yield (if isTarget a then {messageId: a.messageId, toKey: a.toKey, ackedAt: a.ackedAt, originServer: a.originServer, retries: a.retries + 1} else a))
3625+ when (dropCount > 0) (do
3626+ delivBkt <- readDelivBucket b
3627+ let newDelivered = do d <- toDrop; yield {messageId: d.messageId, toKey: d.toKey, deliveredAt: d.ackedAt}
3628+ let combined = union delivBkt newDelivered
3629+ writeDelivBucket b (if count combined > maxDeliveredIds then (evictSmallest combined (count combined - maxDeliveredIds) (\d -> d.deliveredAt)).kept else combined))
3630+ yield dropCount
3631+ incrementAckRetriesPerBucket remainingBuckets server batch (droppedAcc + bucketDropped)
3632+
3633+-- Per-bucket ack removal on success: only rewrites buckets containing matched entries.
3634+removeAckedPerBucket = \bucketIds server batch ->
3635+ case headOf bucketIds of
3636+ Nothing {} -> yield {}
3637+ Just {value: b} -> do
3638+ let remainingBuckets = do x <- bucketIds; where x != b; yield x
3639+ atomic do
3640+ bucket <- readAckedBucket b
3641+ writeAckedBucket b (do a <- bucket; where not (a.originServer == server && hasMessage a.messageId a.toKey batch); yield a)
3642+ removeAckedPerBucket remainingBuckets server batch
3643+
35883644sendBatchedAcksToServer : [{messageId: MessageIdHex, toKey: PubkeyHex, originServer: ServerName, | a}] -> ServerName -> IO {network} {}
35893645sendBatchedAcksToServer = \serverAcks server ->
35903646 forEachBatch serverAcks maxMessageIds (\batch -> do
@@ -3593,29 +3649,17 @@ sendBatchedAcksToServer = \serverAcks server ->
35933649 Err e -> do
35943650 logWarn ("federation: ack to " ++ server ++ " failed: " ++ show e)
35953651 recordFailure server
3596− -- Increment retry counters; drop entries exceeding maxAckRetries to deliveredIds
3597− dropped <- atomic do
3598− ackedIds <- readAllAcked
3599− let isTarget = \a -> a.originServer == server && hasMessage a.messageId a.toKey batch
3600− let toDrop = do a <- ackedIds; where isTarget a && a.retries + 1 >= maxAckRetries; yield a
3601− let dropCount = count toDrop
3602− writeAllAcked (do a <- ackedIds
3603− where not (isTarget a && a.retries + 1 >= maxAckRetries)
3604− yield (if isTarget a then {messageId: a.messageId, toKey: a.toKey, ackedAt: a.ackedAt, originServer: a.originServer, retries: a.retries + 1} else a))
3605− when (dropCount > 0) (do
3606− deliveredIds <- readAllDelivered
3607− let newDelivered = do d <- toDrop; yield {messageId: d.messageId, toKey: d.toKey, deliveredAt: d.ackedAt}
3608− let combined = union deliveredIds newDelivered
3609− writeAllDelivered (if count combined > maxDeliveredIds then (evictSmallest combined (count combined - maxDeliveredIds) (\d -> d.deliveredAt)).kept else combined))
3610− yield dropCount
3652+ -- Increment retry counters per-bucket; drop entries exceeding maxAckRetries to deliveredIds.
3653+ -- Only touches buckets containing entries from this batch (vs readAllAcked scanning all 16).
3654+ let affectedBuckets = uniqueBucketsFor batch (\a -> a.toKey)
3655+ dropped <- incrementAckRetriesPerBucket affectedBuckets server batch 0
36113656 when (dropped > 0)
36123657 (logWarn ("federation: dropped " ++ show dropped ++ " ack entries to " ++ server ++ " after " ++ show maxAckRetries ++ " retries"))
36133658 Ok _ -> do
36143659 clearFailure server
3615− atomic do
3616− ackedIds <- readAllAcked
3617− let updatedAckedIds = ackedIdsExcludingBatch ackedIds server batch
3618− writeAllAcked updatedAckedIds)
3660+ -- Remove acked entries per-bucket (vs readAllAcked/writeAllAcked repartitioning all 16)
3661+ let affectedBuckets = uniqueBucketsFor batch (\a -> a.toKey)
3662+ removeAckedPerBucket affectedBuckets server batch)
36193663
36203664-- Send delivery acknowledgments back to origin servers and remove acked entries on success
36213665sendAcksAndClean : [{messageId: MessageIdHex, toKey: PubkeyHex, originServer: ServerName, | a}] -> IO {network} {}
@@ -3771,8 +3815,13 @@ pruneAndRetry = \t -> do
37713815 stored <- *totalSessionCount
37723816 set *totalSessionCount = actual
37733817 yield (actual - stored)
3774− when (presDrift != 0 || sessDrift != 0)
3775− (logInfo ("reconciliation: corrected drift — presence: " ++ show presDrift ++ ", sessions: " ++ show sessDrift))
3818+ fwdDrift <- atomic do
3819+ actual <- countAllForwards 0
3820+ stored <- *totalForwardCount
3821+ set *totalForwardCount = actual
3822+ yield (actual - stored)
3823+ when (presDrift != 0 || sessDrift != 0 || fwdDrift != 0)
3824+ (logInfo ("reconciliation: corrected drift — presence: " ++ show presDrift ++ ", sessions: " ++ show sessDrift ++ ", forwards: " ++ show fwdDrift))
37763825 -- Message reconciliation: cheap check first (sum rcBucket .count fields vs totalMessageCount).
37773826 -- Full scan only when drift detected or every Nth cycle (reconciliationFullCycleInterval).
37783827 pc <- *pruneCount
@@ -4489,13 +4538,14 @@ sendMessageCore = \senderKey to encryptedBlob clampedTs effectiveMsgId t ->
44894538 let fbi = pubkeyBucket to
44904539 fwdBkt <- readFwdBucket fbi
44914540 let newFwds = do s <- remoteServers; where none fwdBkt (\f -> f.messageId == effectiveMsgId && f.toServer == s); yield {messageId: effectiveMsgId, toServer: s, createdAt: t, toKey: to, encryptedBlob: encryptedBlob, timestamp: clampedTs, originServer: serverHostN, retries: 0}
4492− totalFwdCount <- countAllForwards 0
4541+ totalFwdCount <- *totalForwardCount
44934542 let globalCap = maxForwardsEntries - totalFwdCount
44944543 -- Per-key cap: prevents one key from consuming disproportionate forward slots
44954544 let keyFwdCount = count (do f <- fwdBkt; where f.toKey == to; yield f)
44964545 let keyCap = maxForwardsPerKey - keyFwdCount
44974546 let cappedFwds = takeRelation (min globalCap keyCap) newFwds
44984547 writeFwdBucket fbi (union fwdBkt cappedFwds)
4548+ set *totalForwardCount = totalFwdCount + count cappedFwds
44994549 -- Return whether any forwards were actually enqueued (vs dropped by capacity cap)
45004550 yield (count cappedFwds > 0)
45014551 else yield (False {})
@@ -5300,9 +5350,14 @@ handleReconcile = \authorization respond -> do
53005350 stored <- *totalSessionCount
53015351 set *totalSessionCount = actual
53025352 yield (actual - stored)
5353+ fwdDrift <- atomic do
5354+ actual <- countAllForwards 0
5355+ stored <- *totalForwardCount
5356+ set *totalForwardCount = actual
5357+ yield (actual - stored)
53035358 msgDrift <- reconcileMessageCounts 0 0
5304− when (presDrift != 0 || sessDrift != 0 || msgDrift != 0)
5305− (logInfo ("reconcile: corrected drift — messages: " ++ show msgDrift ++ ", presence: " ++ show presDrift ++ ", sessions: " ++ show sessDrift))
5359+ when (presDrift != 0 || sessDrift != 0 || fwdDrift != 0 || msgDrift != 0)
5360+ (logInfo ("reconcile: corrected drift — messages: " ++ show msgDrift ++ ", presence: " ++ show presDrift ++ ", sessions: " ++ show sessDrift ++ ", forwards: " ++ show fwdDrift))
53065361 yield (respond {ok: True {}, messageDrift: msgDrift, presenceDrift: presDrift, sessionDrift: sessDrift})
53075362
53085363-- Compute health state and update the cached health values.
@@ -5667,6 +5722,7 @@ handleRestore = \version checksumIn presenceData ackedIdsData ackOverflowData de
56675722 let mergedForwards = mergeByKey (\a b -> a.messageId == b.messageId && a.toServer == b.toServer) (\x -> x.createdAt) existingForwards validForwards
56685723 let cappedForwards = takeRelation maxForwardsEntries mergedForwards
56695724 writeAllForwards cappedForwards
5725+ set *totalForwardCount = count cappedForwards
56705726 -- Restore forward nonces (dedup by messageId + toKey + originServer, keep fresher on conflict)
56715727 existingNonces <- readAllNonces
56725728 let validNonces = do n <- forwardNoncesData; where n.seenAt <= t && t - n.seenAt < forwardNonceRetention; yield n