@@ -30,6 +30,7 @@ maxPresenceTtl = 3600 -- max presence TTL from gossip
30maxPubkeyLen = 256
30maxPubkeyLen = 256
31maxBlobLen = 65536
31maxBlobLen = 65536
32maxSignatureLen = 256
32maxSignatureLen = 256
33+maxIpLen = 45 -- max length of an IP address string (IPv6 = 45 chars); used in matchIp instead of maxServerNameLen
33maxMessageIds = 100
34maxMessageIds = 100
34maxBatchSize = 100 -- max messages per batch request (bounds recursion depth in sendBatchInner)
35maxBatchSize = 100 -- max messages per batch request (bounds recursion depth in sendBatchInner)
35maxGossipEvents = 100
36maxGossipEvents = 100
@@ -211,6 +212,11 @@ maxLookupRatePerKey = 5 -- max lookups per key per window
211lookupRateWindowMs : Int<Ms>
212lookupRateWindowMs : Int<Ms>
212lookupRateWindowMs = 1000
213lookupRateWindowMs = 1000
213maxLookupRateEntries = 1000
214maxLookupRateEntries = 1000
215+-- Client write operation rate limiting (per-key, on ack/delete — separate from read-only lookups)
216+maxClientWriteRatePerKey = 10 -- max ack/delete requests per key per window
217+clientWriteRateWindowMs : Int<Ms>
218+clientWriteRateWindowMs = 1000
219+maxClientWriteRateEntries = 1000
214-- Federation verify rate limiting (per-IP, tighter than general IP rate)
220-- Federation verify rate limiting (per-IP, tighter than general IP rate)
215maxVerifyRatePerIp = 5 -- max verify requests per IP per window
221maxVerifyRatePerIp = 5 -- max verify requests per IP per window
216verifyRateWindowMs : Int<Ms>
222verifyRateWindowMs : Int<Ms>
@@ -552,7 +558,14 @@ type BatchResult = {messageId: OptionalMessageId, status: MessageStatus}
553*lastHeartbeatAt : Timestamp
559*lastHeartbeatAt : Timestamp
555−*draining : Bool
561+-- Drain phase: 0=normal (accepting), 1=soft drain (block new auth, serve existing sessions),
562+-- 2=full drain (reject all client requests). Allows staged shutdown:
563+-- phase 1 lets active clients finish while preventing new connections.
566+-- Shutdown requested: when True, background loops exit after drain completes.
567+-- Set by /maintenance/shutdown; cannot be undone (requires process restart).
568+*shutdownRequested : Bool
557*pruneRunning : Timestamp
570*pruneRunning : Timestamp
@@ -562,6 +575,9 @@ type BatchResult = {messageId: OptionalMessageId, status: MessageStatus}
563*lookupRates : [LookupRateEntry]
576*lookupRates : [LookupRateEntry]
578+-- Separate rate limiter for client write operations (ack, delete) — prevents ack bursts from starving lookups
579+*clientWriteRates : [LookupRateEntry]
565*recvRates : [RecvRateEntry]
581*recvRates : [RecvRateEntry]
567*pairRates : [PairRateEntry]
583*pairRates : [PairRateEntry]
@@ -687,6 +703,10 @@ type BatchResult = {messageId: OptionalMessageId, status: MessageStatus}
687-- Last time the background watchdog loop ran (for health check liveness detection)
703-- Last time the background watchdog loop ran (for health check liveness detection)
688*lastWatchdogAt : Timestamp
704*lastWatchdogAt : Timestamp
706+-- Fingerprint of key state counters at last auto-snapshot export.
707+-- backgroundSnapshot compares current counters against this to skip redundant exports.
708+*lastSnapshotFingerprint : Int
690-- Cached health result: returned when /health is rate-limited instead of false all-healthy.
710-- Cached health result: returned when /health is rate-limited instead of false all-healthy.
691-- Updated on every successful (non-rate-limited) health check.
711-- Updated on every successful (non-rate-limited) health check.
692*cachedHealthOk : Bool
712*cachedHealthOk : Bool
@@ -785,12 +805,18 @@ route MaintenanceApi where
785 POST {} /maintenance/prune headers {authorization: Text}
805 POST {} /maintenance/prune headers {authorization: Text}
786 -> {ok: Bool}
806 -> {ok: Bool}
808+ -- phase: 0 or 2 = full drain (backward compat), 1 = soft drain (auth-only block)
788− POST {} /maintenance/drain headers {authorization: Text}
809+ POST {phase: Int} /maintenance/drain headers {authorization: Text}
789− -> {ok: Bool, drained: Int}
810+ -> {ok: Bool, drained: Int, phase: Int}
791 POST {} /maintenance/undrain headers {authorization: Text}
812 POST {} /maintenance/undrain headers {authorization: Text}
792 -> {ok: Bool}
813 -> {ok: Bool}
793 = Undrain
814 = Undrain
815+ -- Graceful shutdown: triggers full drain + sets shutdown flag so background loops exit.
816+ -- Call this instead of drain when the process should stop after flush completes.
817+ POST {} /maintenance/shutdown headers {authorization: Text}
818+ -> {ok: Bool, drained: Int}
794 -- Snapshot: export federation metadata for persistence across restarts.
820 -- Snapshot: export federation metadata for persistence across restarts.
795 -- Messages are excluded (ephemeral by design; clients handle re-delivery).
821 -- Messages are excluded (ephemeral by design; clients handle re-delivery).
796 -- Includes presence, acked/delivered IDs (dedup state), and peer failure records.
822 -- Includes presence, acked/delivered IDs (dedup state), and peer failure records.
@@ -979,6 +1005,14 @@ isOctalStart = \s -> take 1 s == "0" && isDigitChar (take 1 (drop 1 s)) && isAll
979hasOctalOctet : Text -> Bool
1005hasOctalOctet : Text -> Bool
980hasOctalOctet = \s -> isOctalStart s || hasPatternAfterDot isOctalStart s
1006hasOctalOctet = \s -> isOctalStart s || hasPatternAfterDot isOctalStart s
1008+-- Detect percent-encoded characters (e.g., %31%32%37.0.0.1 = 127.0.0.1)
1009+-- Blocks URL-encoded hostnames that could bypass isBadServerName SSRF protection.
1010+containsPercentEncoding : Text -> Bool
1011+containsPercentEncoding = \s ->
1012+ if s == "" then False {}
1013+ else if take 1 s == "%" && isHexChar (take 1 (drop 1 s)) && isHexChar (take 1 (drop 2 s)) then True {}
1014+ else containsPercentEncoding (drop 1 s)
982isHexStart : Text -> Bool
1016isHexStart : Text -> Bool
983isHexStart = \s -> take 2 s == "0x" || take 2 s == "0X"
1017isHexStart = \s -> take 2 s == "0x" || take 2 s == "0X"
984hasHexPrefix : Text -> Bool
1018hasHexPrefix : Text -> Bool
@@ -1102,6 +1136,7 @@ isBadServerName = \sn ->
1102 || containsSuffix ".test" sn
1136 || containsSuffix ".test" sn
1103 || containsSuffix ".example" sn
1137 || containsSuffix ".example" sn
1104 || containsSuffix ".invalid" sn
1138 || containsSuffix ".invalid" sn
1139+ || containsPercentEncoding sn
1106-- Look up a non-expired session matching the given token (constant-time comparison).
1141-- Look up a non-expired session matching the given token (constant-time comparison).
1107-- Caller passes the appropriate session bucket (sharded by token prefix).
1142-- Caller passes the appropriate session bucket (sharded by token prefix).
@@ -1298,6 +1333,21 @@ countSessionsInBucket = \pubkey b ->
1298 rest <- countSessionsInBucket pubkey (b + 1)
1333 rest <- countSessionsInBucket pubkey (b + 1)
1299 yield (n + rest)
1334 yield (n + rest)
1336+-- Apply a filter/transform to every session bucket and return the total count.
1337+-- Combines transform + count in one pass, avoiding a redundant 16-bucket traversal.
1338+pruneAndCountAllSessions : ([Session] -> [Session]) -> IO Int
1339+pruneAndCountAllSessions = \f -> pruneAndCountSessionBucket f 0
1341+pruneAndCountSessionBucket : ([Session] -> [Session]) -> Int -> IO Int
1342+pruneAndCountSessionBucket = \f b ->
1343+ if b > 15 then yield 0
1345+ bucket <- readSessionBucket b
1346+ let pruned = f bucket
1347+ writeSessionBucket b pruned
1348+ rest <- pruneAndCountSessionBucket f (b + 1)
1349+ yield (count pruned + rest)
1301-- Apply a filter/transform to every session bucket
1351-- Apply a filter/transform to every session bucket
1302transformAllSessionBuckets : ([Session] -> [Session]) -> IO ()
1352transformAllSessionBuckets : ([Session] -> [Session]) -> IO ()
1303transformAllSessionBuckets = \f -> do
1353transformAllSessionBuckets = \f -> do
@@ -1997,6 +2047,48 @@ countAllForwards = \b ->
1997 rest <- countAllForwards (b + 1)
2047 rest <- countAllForwards (b + 1)
1998 yield (count bucket + rest)
2048 yield (count bucket + rest)
2050+-- Compute forward stats per-bucket without materializing all forwards into one list.
2051+-- Returns total count, oldest age, unique server count, and max per-server queue depth.
2052+computeForwardStats : Int -> Timestamp -> IO {totalCount: Int, oldestAge: Int, serverCount: Int, maxDepth: Int, serverCounts: [{server: Text, count: Int}]}
2053+computeForwardStats = \b t -> computeForwardStatsInner b t {totalCount: 0, oldestAge: 0, serverCount: 0, maxDepth: 0, serverCounts: []}
2055+computeForwardStatsInner = \b t acc ->
2056+ if b > 15 then yield acc
2058+ bucket <- readFwdBucket b
2059+ let bucketCount = count bucket
2060+ let bucketOldest = case minBy bucket (\x -> x.createdAt) of Nothing {} -> 0; Just {value: x} -> t - x.createdAt
2061+ -- Aggregate per-server counts: merge this bucket's servers into the running totals
2062+ let updatedServerCounts = fold (\sc f ->
2063+ if any sc (\s -> s.server == f.toServer)
2064+ then do s <- sc; yield (if s.server == f.toServer then {server: s.server, count: s.count + 1} else s)
2065+ else union sc [{server: f.toServer, count: 1}]) acc.serverCounts bucket
2066+ let newMaxDepth = fold (\mx s -> max mx s.count) acc.maxDepth updatedServerCounts
2067+ computeForwardStatsInner (b + 1) t {totalCount: acc.totalCount + bucketCount, oldestAge: max acc.oldestAge bucketOldest, serverCount: count updatedServerCounts, maxDepth: newMaxDepth, serverCounts: updatedServerCounts}
2069+-- Read only the forward buckets at the given indices (avoids reading all 16)
2070+readSelectedFwdBuckets : [Int] -> IO [ForwardEntry]
2071+readSelectedFwdBuckets = \buckets ->
2072+ case headOf buckets of
2073+ Nothing {} -> yield []
2074+ Just {value: b} -> do
2075+ let rest = do x <- buckets; where x != b; yield x
2076+ entries <- readFwdBucket b
2077+ remaining <- readSelectedFwdBuckets rest
2078+ yield (union entries remaining)
2080+-- Write forward entries to only the specified bucket indices. Entries are filtered
2081+-- by pubkeyBucket so each bucket gets only its entries. Unspecified buckets are untouched.
2082+writeSelectedFwdBuckets : [Int] -> [ForwardEntry] -> IO ()
2083+writeSelectedFwdBuckets = \buckets entries ->
2084+ case headOf buckets of
2085+ Nothing {} -> yield {}
2086+ Just {value: b} -> do
2087+ let rest = do x <- buckets; where x != b; yield x
2088+ let bucketEntries = do e <- entries; where pubkeyBucket e.toKey == b; yield e
2089+ writeFwdBucket b bucketEntries
2090+ writeSelectedFwdBuckets rest entries
2000writeAllForwards : [ForwardEntry] -> IO ()
2092writeAllForwards : [ForwardEntry] -> IO ()
2001writeAllForwards = \fwds -> do
2093writeAllForwards = \fwds -> do
2002 let partitioned = fold (\acc f ->
2094 let partitioned = fold (\acc f ->
@@ -2249,11 +2341,12 @@ safeCompare = \a b ->
2249-- When a session was created with an IP (behind a proxy), requests without an IP header
2341-- When a session was created with an IP (behind a proxy), requests without an IP header
2250-- are rejected — prevents bypass if the proxy is removed mid-session.
2342-- are rejected — prevents bypass if the proxy is removed mid-session.
2251-- Uses safeCompareLoop to prevent timing side-channels on IP comparison.
2343-- Uses safeCompareLoop to prevent timing side-channels on IP comparison.
2344+-- Uses maxIpLen (45) instead of maxServerNameLen (128) to avoid wasted iterations.
2252matchIp : Text -> Text -> Bool
2345matchIp : Text -> Text -> Bool
2253matchIp = \clientIp sessionIp ->
2346matchIp = \clientIp sessionIp ->
2254 if sessionIp == "" then True {} -- no binding configured for this session
2347 if sessionIp == "" then True {} -- no binding configured for this session
2255 else if clientIp == "" then False {} -- session has IP but request doesn't — proxy disappeared
2348 else if clientIp == "" then False {} -- session has IP but request doesn't — proxy disappeared
2256− else safeCompareLoop clientIp sessionIp maxServerNameLen 0 0 == 0
2349+ else safeCompareLoop clientIp sessionIp maxIpLen 0 0 == 0
2258-- Generic sliding-window rate limiter. Returns whether the request is rate-limited,
2351-- Generic sliding-window rate limiter. Returns whether the request is rate-limited,
2259-- the updated entries list, and retryAfterMs (time until the current window resets, 0 if not limited).
2352-- the updated entries list, and retryAfterMs (time until the current window resets, 0 if not limited).
@@ -2542,28 +2635,35 @@ pruneExpired = \t -> do
2542 set *knownPeerServers = presenceResults.keptPeerServers
2635 set *knownPeerServers = presenceResults.keptPeerServers
2543 when (presenceResults.totalRemoved > 0) (set *presenceVersion = t)
2636 when (presenceResults.totalRemoved > 0) (set *presenceVersion = t)
2544 let presenceResult = {expiredLocal: presenceResults.expiredLocal, peerServers: presenceResults.allPeerServers}
2637 let presenceResult = {expiredLocal: presenceResults.expiredLocal, peerServers: presenceResults.allPeerServers}
2545− -- Prune auth state
2638+ -- Prune auth state (challenges in their own block; sessions counted during transform)
2546 atomic do
2639 atomic do
2547 challenges <- *challenges
2640 challenges <- *challenges
2548 set *challenges = do x <- challenges; where x.expiresAt > t; yield x
2641 set *challenges = do x <- challenges; where x.expiresAt > t; yield x
2642+ -- Prune sessions and count in a single pass (avoids redundant 16-bucket traversal).
2643+ -- pruneAndCountAllSessions applies the filter to each bucket and sums the survivors.
2549− transformAllSessionBuckets (\bucket -> do x <- bucket; where x.expiresAt > t && t - x.createdAt < maxSessionLifetime; yield x)
2644+ sessionCount <- pruneAndCountAllSessions (\bucket -> do x <- bucket; where x.expiresAt > t && t - x.createdAt < maxSessionLifetime; yield x)
2550− sessionCount <- countAllSessions
2551− set *totalSessionCount = sessionCount
2645+ atomic do set *totalSessionCount = sessionCount
2552 -- Prune messages across all 16 buckets concurrently. Each bucket prune is atomic
2646 -- Prune messages across all 16 buckets concurrently. Each bucket prune is atomic
2553 -- and touches only its own message bucket + rcBucket. The only shared write is
2647 -- and touches only its own message bucket + rcBucket. The only shared write is
2554 -- *totalMessageCount (a single Int), so STM contention is minimal.
2648 -- *totalMessageCount (a single Int), so STM contention is minimal.
2555 -- Fork first 15 buckets; run the last one in the foreground to bound concurrency.
2649 -- Fork first 15 buckets; run the last one in the foreground to bound concurrency.
2556 let messageCutoff = t - unackedMessageMaxAge
2650 let messageCutoff = t - unackedMessageMaxAge
2557 pruneMessageBucketsConcurrent 0 messageCutoff
2651 pruneMessageBucketsConcurrent 0 messageCutoff
2558− -- Prune federation state
2652+ -- Prune federation state (split into per-collection blocks to reduce lock contention
2653+ -- with poll-ack paths that write to ackedIds/deliveredIds concurrently).
2654+ let retainCutoff = t - messageMaxAge
2655+ -- Forwards: independent from acked/delivered, safe to prune separately
2559 atomic do
2656 atomic do
2560− let retainCutoff = t - messageMaxAge
2561 forwards <- readAllForwards
2657 forwards <- readAllForwards
2562 -- Clear all forwards when federation is disabled (they can never be retried)
2658 -- Clear all forwards when federation is disabled (they can never be retried)
2563 writeAllForwards (if not federationEnabled then []
2659 writeAllForwards (if not federationEnabled then []
2564 else do f <- forwards; where f.createdAt > retainCutoff; yield f)
2660 else do f <- forwards; where f.createdAt > retainCutoff; yield f)
2661+ -- AckedIds: contends with poll-ack (ackPolledMessages writes here), narrow the lock
2565 ackedIds <- readAllAcked
2663 ackedIds <- readAllAcked
2566 writeAllAcked (do a <- ackedIds; where a.ackedAt > retainCutoff; yield a)
2664 writeAllAcked (do a <- ackedIds; where a.ackedAt > retainCutoff; yield a)
2665+ -- DeliveredIds + overflow + failures: smaller collections, grouped together
2567 forwardFailures <- *forwardFailures
2667 forwardFailures <- *forwardFailures
2568 set *forwardFailures = do f <- forwardFailures; where t - f.failedAt < maxServerDowntime; yield f
2668 set *forwardFailures = do f <- forwardFailures; where t - f.failedAt < maxServerDowntime; yield f
2569 ackOverflow <- *ackOverflow
2669 ackOverflow <- *ackOverflow
@@ -2601,6 +2701,8 @@ pruneExpired = \t -> do
2601 set *ipRates = do r <- ipRates; where t - r.windowStart < ipRateWindowMs; yield r
2701 set *ipRates = do r <- ipRates; where t - r.windowStart < ipRateWindowMs; yield r
2602 lookupRates <- *lookupRates
2702 lookupRates <- *lookupRates
2603 set *lookupRates = do r <- lookupRates; where t - r.windowStart < lookupRateWindowMs; yield r
2703 set *lookupRates = do r <- lookupRates; where t - r.windowStart < lookupRateWindowMs; yield r
2704+ clientWriteRates <- *clientWriteRates
2705+ set *clientWriteRates = do r <- clientWriteRates; where t - r.windowStart < clientWriteRateWindowMs; yield r
2604 verifyIpRates <- *verifyIpRates
2706 verifyIpRates <- *verifyIpRates
2605 set *verifyIpRates = do r <- verifyIpRates; where t - r.windowStart < verifyRateWindowMs; yield r
2707 set *verifyIpRates = do r <- verifyIpRates; where t - r.windowStart < verifyRateWindowMs; yield r
2606 uniqueKeyIpRates <- *uniqueKeyIpRates
2708 uniqueKeyIpRates <- *uniqueKeyIpRates
@@ -2794,9 +2896,10 @@ checkClientIpRate = \ip t ->
2794 yield result.limited
2896 yield result.limited
2796-- Combined pre-auth gate: drain check, global rate, IP rate. Returns "" if all pass.
2898-- Combined pre-auth gate: drain check, global rate, IP rate. Returns "" if all pass.
2899+-- Blocks at drainPhase >= 1 (soft drain blocks new auth; existing sessions still work).
2797checkPreAuth = \clientIp t -> do
2900checkPreAuth = \clientIp t -> do
2798− draining <- *draining
2901+ dp <- *drainPhase
2799− if draining then yield "draining"
2902+ if dp >= 1 then yield "draining"
2801 globalLimited <- checkGlobalRate t
2904 globalLimited <- checkGlobalRate t
2802 if globalLimited then yield "global_rate_limited"
2905 if globalLimited then yield "global_rate_limited"
@@ -2896,8 +2999,8 @@ gossipOfflineIfLastSession : {pubkey: PubkeyHex, gossipOffline: Bool, peerServer
2896gossipOfflineIfLastSession = \r label -> do
2999gossipOfflineIfLastSession = \r label -> do
2897 when r.gossipOffline (do
3000 when r.gossipOffline (do
2898 t <- now
3001 t <- now
2899− isDraining <- *draining
3002+ dp <- *drainPhase
2900− if isDraining
3003+ if dp >= 2
2901 then logInfo (label ++ ": last session for " ++ r.pubkey ++ " (drain active, skipping redundant offline gossip)")
3004 then logInfo (label ++ ": last session for " ++ r.pubkey ++ " (drain active, skipping redundant offline gossip)")
2903 logInfo (label ++ ": last session for " ++ r.pubkey ++ ", gossiping offline")
3006 logInfo (label ++ ": last session for " ++ r.pubkey ++ ", gossiping offline")
@@ -3055,7 +3158,10 @@ retryPendingForwards = \pending -> do
3055-- When a remote key comes online, create forward entries for queued messages and send them
3158-- When a remote key comes online, create forward entries for queued messages and send them
3056forwardQueuedMessages = \targets t -> do
3159forwardQueuedMessages = \targets t -> do
3057 newForwards <- atomic do
3160 newForwards <- atomic do
3058− forwards <- readAllForwards
3161+ -- Read only the forward buckets needed for target pubkeys (not all 16).
3162+ -- All entries for a pubkey hash to the same bucket, so dedup is per-bucket complete.
3163+ let neededFwdBuckets = uniqueBucketsFor targets (\target -> target.pubkey)
3164+ forwards <- readSelectedFwdBuckets neededFwdBuckets
3059 -- Read only the presence buckets needed for target pubkeys (not all 16)
3165 -- Read only the presence buckets needed for target pubkeys (not all 16)
3060 let targetPresenceBuckets = uniqueBucketsFor targets (\target -> target.pubkey)
3166 let targetPresenceBuckets = uniqueBucketsFor targets (\target -> target.pubkey)
3061 relevantPresence <- readSelectedPresenceBuckets targetPresenceBuckets
3167 relevantPresence <- readSelectedPresenceBuckets targetPresenceBuckets
@@ -3067,7 +3173,9 @@ forwardQueuedMessages = \targets t -> do
3067 let candidates = forwardCandidates remoteTargets relevantMessages forwards t
3173 let candidates = forwardCandidates remoteTargets relevantMessages forwards t
3068 let totalFwdCount <- countAllForwards 0
3174 let totalFwdCount <- countAllForwards 0
3069 let capped = takeRelation (maxForwardsEntries - totalFwdCount) candidates
3175 let capped = takeRelation (maxForwardsEntries - totalFwdCount) candidates
3176+ -- Write back only the affected forward buckets (not all 16)
3070− writeAllForwards (union forwards capped)
3177+ let combined = union forwards capped
3178+ writeSelectedFwdBuckets neededFwdBuckets combined
3071 yield capped
3179 yield capped
3072 let servers = collectPeerServers newForwards (\x -> x.toServer)
3180 let servers = collectPeerServers newForwards (\x -> x.toServer)
3073 failures <- *forwardFailures
3181 failures <- *forwardFailures
@@ -3524,8 +3632,8 @@ pollOnce = \token cursor clientIp -> do
3524 -- authenticateSession to skip safeCompare on unrelated sessions.
3632 -- authenticateSession to skip safeCompare on unrelated sessions.
3525 t2 <- now
3633 t2 <- now
3526 result <- atomic do
3634 result <- atomic do
3527− isDraining <- *draining
3635+ dp <- *drainPhase
3528− if isDraining then yield {authorized: True {}, events: [], cursor: s.clampedCursor}
3636+ if dp >= 2 then yield {authorized: True {}, events: [], cursor: s.clampedCursor}
3530 sessions2 <- readSessionBucket bucket
3638 sessions2 <- readSessionBucket bucket
3531 if not (reauthenticateSession sessions2 token s.pubkey t2)
3639 if not (reauthenticateSession sessions2 token s.pubkey t2)
@@ -3784,8 +3892,8 @@ verifyAuthInner = \pubkey challenge signature revokeOthers clientIp respond t ->
3784 token <- genToken
3892 token <- genToken
3785 let newBucket = tokenBucket token
3893 let newBucket = tokenBucket token
3786 authResult <- atomic do
3894 authResult <- atomic do
3787− isDraining <- *draining
3895+ dp <- *drainPhase
3788− if isDraining then yield (Nothing {})
3896+ if dp >= 1 then yield (Nothing {})
3790 tpc <- *totalPresenceCount
3898 tpc <- *totalPresenceCount
3791 knownPeers <- *knownPeerServers
3899 knownPeers <- *knownPeerServers
@@ -3845,8 +3953,8 @@ handlePoll = \cursor generation watchKeys authorization clientIp respond -> do
3845 -- Tells the client their cursor was invalidated so they can handle re-delivery.
3953 -- Tells the client their cursor was invalidated so they can handle re-delivery.
3846 let cursorWasReset = generationMismatch && cursor > 0
3954 let cursorWasReset = generationMismatch && cursor > 0
3847 let cappedWatchKeys = takeRelation maxWatchKeys watchKeys
3955 let cappedWatchKeys = takeRelation maxWatchKeys watchKeys
3848− draining <- *draining
3956+ dp <- *drainPhase
3849− if draining then yield (respond {events: [], cursor: 0, authorized: True {}, rateLimited: False {}, draining: True {}, generation: serverGen, presenceUpdates: [], heartbeat: False {}, retryAfterMs: 0, cursorReset: cursorWasReset})
3957+ if dp >= 2 then yield (respond {events: [], cursor: 0, authorized: True {}, rateLimited: False {}, draining: True {}, generation: serverGen, presenceUpdates: [], heartbeat: False {}, retryAfterMs: 0, cursorReset: cursorWasReset})
3851 globalRateLimited <- checkGlobalRate t
3959 globalRateLimited <- checkGlobalRate t
3852 if globalRateLimited then yield (respond {events: [], cursor: 0, authorized: False {}, rateLimited: True {}, draining: False {}, generation: serverGen, presenceUpdates: [], heartbeat: False {}, retryAfterMs: globalRateWindowMs, cursorReset: cursorWasReset})
3960 if globalRateLimited then yield (respond {events: [], cursor: 0, authorized: False {}, rateLimited: True {}, draining: False {}, generation: serverGen, presenceUpdates: [], heartbeat: False {}, retryAfterMs: globalRateWindowMs, cursorReset: cursorWasReset})
@@ -4030,8 +4138,8 @@ handleSendMessage = \to encryptedBlob timestamp messageId authorization clientIp
4030 (logWarn ("sendMessage: stale timestamp (" ++ show ((t - timestamp) / 86400000) ++ " days old) for " ++ effectiveMsgId))
4138 (logWarn ("sendMessage: stale timestamp (" ++ show ((t - timestamp) / 86400000) ++ " days old) for " ++ effectiveMsgId))
4031 let clampedTs = clampTimestamp timestamp t
4139 let clampedTs = clampTimestamp timestamp t
4032 result <- atomic do
4140 result <- atomic do
4033− draining <- *draining
4141+ dp <- *drainPhase
4034− if draining then yield {status: MsDraining {}, remoteServers: [], retryAfterMs: 0}
4142+ if dp >= 2 then yield {status: MsDraining {}, remoteServers: [], retryAfterMs: 0}
4035 else withSessionAuth authorization t clientIp
4143 else withSessionAuth authorization t clientIp
4036 (yield {status: MsUnauthorized {}, remoteServers: [], retryAfterMs: 0})
4144 (yield {status: MsUnauthorized {}, remoteServers: [], retryAfterMs: 0})
4037 (\senderKey -> sendMessageCore senderKey to encryptedBlob clampedTs effectiveMsgId t)
4145 (\senderKey -> sendMessageCore senderKey to encryptedBlob clampedTs effectiveMsgId t)
@@ -4055,8 +4163,8 @@ handleLookupPresence = \pubkey authorization clientIp respond -> do
4055 if globalRateLimited then yield (respond {server: "", authorized: False {}})
4163 if globalRateLimited then yield (respond {server: "", authorized: False {}})
4057 result <- atomic do
4165 result <- atomic do
4058− draining <- *draining
4166+ dp <- *drainPhase
4059− if draining then yield {server: "", authorized: False {}}
4167+ if dp >= 2 then yield {server: "", authorized: False {}}
4060 else withSessionAuth authorization t clientIp
4168 else withSessionAuth authorization t clientIp
4061 (yield {server: "", authorized: False {}})
4169 (yield {server: "", authorized: False {}})
4062 (\senderKey -> do
4170 (\senderKey -> do
@@ -4081,8 +4189,8 @@ handleLookupPresence = \pubkey authorization clientIp respond -> do
4081-- "delivered" (acked by recipient), "unknown" (expired or never existed).
4189-- "delivered" (acked by recipient), "unknown" (expired or never existed).
4082handleMessageStatus = \to messageId authorization clientIp respond -> do
4190handleMessageStatus = \to messageId authorization clientIp respond -> do
4083 t <- now
4191 t <- now
4084− draining <- *draining
4192+ dp <- *drainPhase
4085− if draining then yield (respond {status: "", authorized: False {}})
4193+ if dp >= 2 then yield (respond {status: "", authorized: False {}})
4087 globalRateLimited <- checkGlobalRate t
4195 globalRateLimited <- checkGlobalRate t
4088 if globalRateLimited then yield (respond {status: "", authorized: False {}})
4196 if globalRateLimited then yield (respond {status: "", authorized: False {}})
@@ -4225,8 +4333,13 @@ handleVerifyToken = \token scope clientIp respond -> do
4226 found <- atomic do
4334 found <- atomic do
4227 pending <- *pendingVerifications
4335 pending <- *pendingVerifications
4228− -- Scope is required: callers must specify the operation type (gossip/forward/ack)
4336+ -- Pre-filter by scope, expiry, and first hex char before running expensive safeCompare.
4337+ -- Reduces scan from O(n) safeCompare calls to O(n/16) by filtering on token prefix.
4338+ -- The prefix filter is not constant-time, but single-use tokens with 24 bytes of randomness
4339+ -- make first-char leakage irrelevant (16^47 remaining possibilities in 30s window).
4340+ let firstChar = take 1 token
4229− let match = any pending (\p -> safeCompare p.token token && p.scope == scope && t - p.createdAt < verificationTokenExpiry)
4341+ let candidates = do p <- pending; where p.scope == scope && t - p.createdAt < verificationTokenExpiry && take 1 p.token == firstChar; yield p
4342+ let match = any candidates (\p -> safeCompare p.token token)
4230 when match (set *pendingVerifications = do p <- pending; where not (safeCompare p.token token); yield p)
4343 when match (set *pendingVerifications = do p <- pending; where not (safeCompare p.token token); yield p)
4231 yield match
4344 yield match
4232 -- Omit scope from log to prevent scope enumeration by attackers probing the endpoint
4345 -- Omit scope from log to prevent scope enumeration by attackers probing the endpoint
@@ -4480,15 +4593,16 @@ handleFederationAck = \messageIds originServer fromServer requestTimestamp corre
4480 writeAllForwards (forwardsExcluding forwards fromServer dedupedIds)
4593 writeAllForwards (forwardsExcluding forwards fromServer dedupedIds)
4481 yield removed
4594 yield removed
4482 let ackedCount = count removedForwards
4595 let ackedCount = count removedForwards
4483− -- Only clear failure record when actual forwards were matched (proves the peer
4484− -- received and processed messages we sent). Unmatched acks don't prove forward delivery.
4485− when (ackedCount > 0) (clearFailure fromServer)
4486 -- Guard: warn if peer acks messages we had no forward entries for. This can indicate
4596 -- Guard: warn if peer acks messages we had no forward entries for. This can indicate
4487 -- a peer acking messages it never received (e.g., the forward response was lost but
4597 -- a peer acking messages it never received (e.g., the forward response was lost but
4488 -- the message was never stored), which would cause premature origin message deletion.
4598 -- the message was never stored), which would cause premature origin message deletion.
4489 let unmatched = count dedupedIds - ackedCount
4599 let unmatched = count dedupedIds - ackedCount
4490 when (unmatched > 0)
4600 when (unmatched > 0)
4491 (logWarn ("federation/ack: " ++ fromServer ++ " acked " ++ show unmatched ++ " messageIds with no matching forward entries — skipping cleanup for those"))
4601 (logWarn ("federation/ack: " ++ fromServer ++ " acked " ++ show unmatched ++ " messageIds with no matching forward entries — skipping cleanup for those"))
4602+ -- Only clear failure record when ALL acked IDs matched forward entries.
4603+ -- Prevents a malicious peer from resetting its backoff by padding ack requests
4604+ -- with known messageIds that no longer have forward entries.
4605+ when (ackedCount > 0 && unmatched == 0) (clearFailure fromServer)
4492 -- Deferred origin message cleanup: only queue cleanup for messages that actually had
4606 -- Deferred origin message cleanup: only queue cleanup for messages that actually had
4493 -- matching forward entries. Messages acked without a forward entry are NOT cleaned up,
4607 -- matching forward entries. Messages acked without a forward entry are NOT cleaned up,
4494 -- preventing premature deletion when a peer acks messages it never actually received.
4608 -- preventing premature deletion when a peer acks messages it never actually received.
@@ -4506,20 +4620,28 @@ handleFederationAck = \messageIds originServer fromServer requestTimestamp corre
4506-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
4620-- POST /maintenance/drain: Broadcast offline gossip for all local keys (graceful shutdown).
4507-- Call before stopping the server so peers learn keys went offline immediately.
4621-- Call before stopping the server so peers learn keys went offline immediately.
4508-- Gossip is sent via fork; allow a few seconds after response before process exit.
4622-- Gossip is sent via fork; allow a few seconds after response before process exit.
4623+-- phase: 0 or 2 = full drain (backward compat default), 1 = soft drain (block new auth only)
4509−handleDrain = \authorization respond -> do
4624+handleDrain = \phase authorization respond -> do
4510 t <- now
4625 t <- now
4511 let authorized = verifyMaintenanceAuth authorization
4626 let authorized = verifyMaintenanceAuth authorization
4512− if not authorized then yield (respond {ok: False {}, drained: 0})
4627+ if not authorized then yield (respond {ok: False {}, drained: 0, phase: 0})
4629+ -- Effective phase: 0 or >=2 means full drain (backward compat for clients sending empty body)
4630+ let effectivePhase = if phase == 1 then 1 else 2
4514− -- Idempotency: if already draining, return success without re-running drain logic.
4631+ -- Idempotency: if already at or above requested phase, return success.
4515− -- Prevents duplicate gossip broadcasts, snapshot exports, and flush forks.
4632+ currentPhase <- *drainPhase
4516− alreadyDraining <- *draining
4633+ if currentPhase >= effectivePhase then do
4517− if alreadyDraining then do
4518− logInfo "drain: already draining, skipping"
4634+ logInfo ("drain: already at phase " ++ show currentPhase ++ ", skipping")
4519− yield (respond {ok: True {}, drained: 0})
4635+ yield (respond {ok: True {}, drained: 0, phase: currentPhase})
4636+ -- Soft drain (phase 1): block new auth but keep serving existing sessions.
4637+ -- No gossip or flush needed since sessions are still active.
4638+ else if effectivePhase == 1 then do
4639+ atomic do set *drainPhase = 1
4640+ logInfo "drain: soft drain active (new auth blocked, existing sessions served)"
4641+ yield (respond {ok: True {}, drained: 0, phase: 1})
4521 drainResult <- atomic do
4643 drainResult <- atomic do
4522− set *draining = True {}
4644+ set *drainPhase = 2
4523 allPresence <- readAllPresence
4645 allPresence <- readAllPresence
4524 let localEntries = do p <- allPresence; where isLocalServer p.server && p.expiresAt > t; yield p
4646 let localEntries = do p <- allPresence; where isLocalServer p.server && p.expiresAt > t; yield p
4525 let remainingPresence = do p <- allPresence; where not (isLocalServer p.server) || p.expiresAt <= t; yield p
4647 let remainingPresence = do p <- allPresence; where not (isLocalServer p.server) || p.expiresAt <= t; yield p
@@ -4527,7 +4649,10 @@ handleDrain = \authorization respond -> do
4527 set *totalPresenceCount = count remainingPresence
4649 set *totalPresenceCount = count remainingPresence
4528 set *presenceVersion = t
4650 set *presenceVersion = t
4529 set *pollWakeup = t
4651 set *pollWakeup = t
4530− yield {localEntries: localEntries, peerServers: collectPeerServers allPresence (\p -> p.server)}
4652+ -- Use cached knownPeerServers instead of recomputing from allPresence
4653+ -- (knownPeerServers is updated every prune cycle, equivalent for gossip targets)
4654+ peerServers <- *knownPeerServers
4655+ yield {localEntries: localEntries, peerServers: peerServers}
4531 let drainedCount = count drainResult.localEntries
4656 let drainedCount = count drainResult.localEntries
4532 -- Read pending federation state atomically (fix: was read outside atomic)
4657 -- Read pending federation state atomically (fix: was read outside atomic)
4533 flushState <- atomic do
4658 flushState <- atomic do
@@ -4585,7 +4710,7 @@ handleDrain = \authorization respond -> do
4585 (logInfo ("drain: flushing " ++ show pendingForwardCount ++ " pending forwards"))
4710 (logInfo ("drain: flushing " ++ show pendingForwardCount ++ " pending forwards"))
4586 when (pendingAckCount > 0)
4711 when (pendingAckCount > 0)
4587 (logInfo ("drain: flushing " ++ show pendingAckCount ++ " pending acks"))
4712 (logInfo ("drain: flushing " ++ show pendingAckCount ++ " pending acks"))
4588− yield (respond {ok: True {}, drained: drainedCount})
4713+ yield (respond {ok: True {}, drained: drainedCount, phase: 2})
4590-- POST /maintenance/undrain: Resume accepting requests after a drain.
4715-- POST /maintenance/undrain: Resume accepting requests after a drain.
4591-- Background loops (prune, heartbeat) automatically resume on next iteration.
4716-- Background loops (prune, heartbeat) automatically resume on next iteration.
@@ -4594,20 +4719,40 @@ handleUndrain = \authorization respond -> do
4594 if not authorized then yield (respond {ok: False {}})
4719 if not authorized then yield (respond {ok: False {}})
4596 wasDraining <- atomic do
4721 wasDraining <- atomic do
4597− d <- *draining
4722+ d <- *drainPhase
4598− set *draining = False {}
4723+ set *drainPhase = 0
4599 set *drainCompletedAt = 0
4724 set *drainCompletedAt = 0
4725+ set *shutdownRequested = False {}
4600− yield d
4726+ yield (d > 0)
4601 when wasDraining (logInfo "undrain: server resumed, accepting new requests")
4727 when wasDraining (logInfo "undrain: server resumed, accepting new requests")
4602 when (not wasDraining) (logInfo "undrain: server was not draining")
4728 when (not wasDraining) (logInfo "undrain: server was not draining")
4603 yield (respond {ok: True {}})
4729 yield (respond {ok: True {}})
4731+-- POST /maintenance/shutdown: Trigger full drain + set shutdown flag.
4732+-- Background loops will exit after drain flush completes, allowing the process to stop.
4733+-- Unlike drain, shutdown is not undoable (requires process restart).
4734+handleShutdown = \authorization respond -> do
4736+ let authorized = verifyMaintenanceAuth authorization
4737+ if not authorized then yield (respond {ok: False {}, drained: 0})
4741+ -- Already draining, just set the shutdown flag
4742+ atomic do set *shutdownRequested = True {}
4743+ logInfo "shutdown: flag set (drain already active)"
4744+ yield (respond {ok: True {}, drained: 0})
4746+ atomic do set *shutdownRequested = True {}
4747+ -- Delegate to handleDrain with phase=2 (full drain)
4748+ handleDrain 2 authorization (\drainResult -> respond {ok: drainResult.ok, drained: drainResult.drained})
4605-- POST /maintenance/prune: Trigger a full maintenance cycle (prune + retry all pending ops).
4750-- POST /maintenance/prune: Trigger a full maintenance cycle (prune + retry all pending ops).
4606-- Accepts either federation or maintenance secret.
4751-- Accepts either federation or maintenance secret.
4607handlePrune = \authorization respond -> do
4752handlePrune = \authorization respond -> do
4608 t <- now
4753 t <- now
4609− draining <- *draining
4754+ dp <- *drainPhase
4610− let authorized = verifyMaintenanceAuth authorization && not draining
4755+ let authorized = verifyMaintenanceAuth authorization && dp < 2
4611 when authorized (fork (pruneAndRetry t))
4756 when authorized (fork (pruneAndRetry t))
4612 yield (respond {ok: authorized})
4757 yield (respond {ok: authorized})
@@ -4626,8 +4771,8 @@ handleHealth = \authorization respond -> do
4626 cF <- *cachedHealthFederation
4771 cF <- *cachedHealthFederation
4627 cS <- *cachedHealthSnapshot
4772 cS <- *cachedHealthSnapshot
4628 cW <- *cachedHealthWatchdog
4773 cW <- *cachedHealthWatchdog
4629− isDraining <- *draining
4774+ dp <- *drainPhase
4630− yield {ok: cOk && not isDraining, pruneHealthy: cP, heartbeatHealthy: cH, capacityHealthy: cC, federationHealthy: cF, snapshotHealthy: cS, watchdogHealthy: cW, draining: isDraining}
4775+ yield {ok: cOk && dp == 0, pruneHealthy: cP, heartbeatHealthy: cH, capacityHealthy: cC, federationHealthy: cF, snapshotHealthy: cS, watchdogHealthy: cW, draining: dp > 0}
4631 yield (respond cached)
4776 yield (respond cached)
4633 let isAuthed = verifyMaintenanceAuth authorization
4778 let isAuthed = verifyMaintenanceAuth authorization
@@ -4645,7 +4790,7 @@ handleHealth = \authorization respond -> do
4645 forwardFailures <- *forwardFailures
4790 forwardFailures <- *forwardFailures
4646 pruneRun <- *pruneRunning
4791 pruneRun <- *pruneRunning
4647 lastSnap <- *lastSnapshotExportAt
4792 lastSnap <- *lastSnapshotExportAt
4648− isDraining <- *draining
4793+ dp <- *drainPhase
4649 let warmup = t - sa < healthPruneMaxAge
4794 let warmup = t - sa < healthPruneMaxAge
4650 let pruneStuck = pruneRun > 0 && t - pruneRun > backgroundPruneInterval * 2
4795 let pruneStuck = pruneRun > 0 && t - pruneRun > backgroundPruneInterval * 2
4651 -- Capacity pressure: flag when any capped collection exceeds 90% full
4796 -- Capacity pressure: flag when any capped collection exceeds 90% full
@@ -4671,7 +4816,7 @@ handleHealth = \authorization respond -> do
4671 federationOk: fedOk,
4816 federationOk: fedOk,
4672 snapshotOk: snapOk,
4817 snapshotOk: snapOk,
4673 watchdogOk: wdOk,
4818 watchdogOk: wdOk,
4674− isDraining: isDraining}
4819+ isDraining: dp > 0}
4675 -- Draining servers report ok=False so load balancers stop routing traffic to them
4820 -- Draining servers report ok=False so load balancers stop routing traffic to them
4676 let allOk = not result.isDraining && result.pruneOk && result.heartbeatOk && result.capacityOk && result.federationOk && result.snapshotOk && result.watchdogOk
4821 let allOk = not result.isDraining && result.pruneOk && result.heartbeatOk && result.capacityOk && result.federationOk && result.snapshotOk && result.watchdogOk
4677 -- Cache health result for rate-limited responses
4822 -- Cache health result for rate-limited responses
@@ -4714,10 +4859,12 @@ handleStats = \authorization respond -> do
4715 -- Split into two atomic reads: core capacity state (consistent) and counters/timing
4860 -- Split into two atomic reads: core capacity state (consistent) and counters/timing
4716 -- (slight inconsistency acceptable for monitoring). This narrows the lock window.
4861 -- (slight inconsistency acceptable for monitoring). This narrows the lock window.
4862+ -- Forward stats computed per-bucket to avoid materializing all 16 buckets into one list.
4863+ -- Each bucket is read independently; counts and per-server stats are aggregated.
4864+ fwdStats <- computeForwardStats 0 t
4717 core <- atomic do
4865 core <- atomic do
4718 sc <- *totalSessionCount
4866 sc <- *totalSessionCount
4719 pc <- *totalPresenceCount
4867 pc <- *totalPresenceCount
4720− f <- readAllForwards
4721 c <- *challenges
4868 c <- *challenges
4722 a <- readAllAcked
4869 a <- readAllAcked
4723 ao <- *ackOverflow
4870 ao <- *ackOverflow
@@ -4734,7 +4881,7 @@ handleStats = \authorization respond -> do
4734 ger <- *gossipEventRates
4881 ger <- *gossipEventRates
4735 vir <- *verifyIpRates
4882 vir <- *verifyIpRates
4736 ukir <- *uniqueKeyIpRates
4883 ukir <- *uniqueKeyIpRates
4737− yield {sc: sc, pc: pc, f: f, c: c, a: a, ao: ao, d: d, ff: ff, ap: ap, tmc: tmc, rc: rc, fi: fi, rg: rg, fnCount: fnCount, pac: pac, bp: bp, ger: ger, vir: vir, ukir: ukir}
4884+ yield {sc: sc, pc: pc, c: c, a: a, ao: ao, d: d, ff: ff, ap: ap, tmc: tmc, rc: rc, fi: fi, rg: rg, fnCount: fnCount, pac: pac, bp: bp, ger: ger, vir: vir, ukir: ukir}
4738 counters <- atomic do
4885 counters <- atomic do
4739 sr <- *sendRates
4886 sr <- *sendRates
4740 cr <- *challengeRates
4887 cr <- *challengeRates
@@ -4751,7 +4898,7 @@ handleStats = \authorization respond -> do
4751 sa <- *startedAt
4898 sa <- *startedAt
4752 lp <- *lastPrunedAt
4899 lp <- *lastPrunedAt
4753 lh <- *lastHeartbeatAt
4900 lh <- *lastHeartbeatAt
4754− dr <- *draining
4901+ dr <- *drainPhase
4755 flt <- *federationLatencyTotal
4902 flt <- *federationLatencyTotal
4756 flc <- *federationLatencyCount
4903 flc <- *federationLatencyCount
4757 flm <- *federationLatencyMax
4904 flm <- *federationLatencyMax
@@ -4761,27 +4908,22 @@ handleStats = \authorization respond -> do
4761 let sessions = core.sc
4908 let sessions = core.sc
4762 let messages = core.tmc
4909 let messages = core.tmc
4763 let presence = core.pc
4910 let presence = core.pc
4764− let fwds = count core.f
4911+ let fwds = fwdStats.totalCount
4765 let msgPct = if maxMessages > 0 then messages * 100 / maxMessages else 0
4912 let msgPct = if maxMessages > 0 then messages * 100 / maxMessages else 0
4766 let presPct = if maxPresenceEntries > 0 then presence * 100 / maxPresenceEntries else 0
4913 let presPct = if maxPresenceEntries > 0 then presence * 100 / maxPresenceEntries else 0
4767 let fwdPct = if maxForwardsEntries > 0 then fwds * 100 / maxForwardsEntries else 0
4914 let fwdPct = if maxForwardsEntries > 0 then fwds * 100 / maxForwardsEntries else 0
4768 let sessPct = if maxSessions > 0 then sessions * 100 / maxSessions else 0
4915 let sessPct = if maxSessions > 0 then sessions * 100 / maxSessions else 0
4769 let fedAvg = if counters.flc > 0 then counters.flt / counters.flc else 0
4916 let fedAvg = if counters.flc > 0 then counters.flt / counters.flc else 0
4770 let hungFedSlots = count (do r <- core.fi; where r.count > 0 && r.firstStartedAt > 0 && t - r.firstStartedAt >= federationFetchTimeoutMs; yield r)
4917 let hungFedSlots = count (do r <- core.fi; where r.count > 0 && r.firstStartedAt > 0 && t - r.firstStartedAt >= federationFetchTimeoutMs; yield r)
4771− let oldestFwdAge = case minBy core.f (\x -> x.createdAt) of Nothing {} -> 0; Just {value: x} -> t - x.createdAt
4918+ let oldestFwdAge = fwdStats.oldestAge
4772 let oldestAckAge = case minBy core.a (\x -> x.ackedAt) of Nothing {} -> 0; Just {value: x} -> t - x.ackedAt
4919 let oldestAckAge = case minBy core.a (\x -> x.ackedAt) of Nothing {} -> 0; Just {value: x} -> t - x.ackedAt
4773− -- Per-server forward queue summary: count unique servers and max queue depth
4774− -- Two-pass aggregation: collect unique servers, then count per server.
4775− -- Avoids O(n*k) accumulator rebuilding in the fold (no list comprehension per step).
4776− let fwdServers = fold (\acc f -> if any acc (\s -> s == f.toServer) then acc else union acc [f.toServer]) [] core.f
4777− let fwdCounts = do s <- fwdServers; yield {server: s, count: count (do f <- core.f; where f.toServer == s; yield f)}
4778− let fwdServerCount = count fwdCounts
4920+ let fwdServerCount = fwdStats.serverCount
4779− let maxFwdDepth = fold (\acc c -> max acc c.count) 0 fwdCounts
4921+ let maxFwdDepth = fwdStats.maxDepth
4780 yield (respond {ok: True {},
4922 yield (respond {ok: True {},
4781 capacity: {sessions: sessions, sessionsMax: maxSessions, sessionsPct: sessPct, messages: messages, messagesMax: maxMessages, messagesPct: msgPct, totalMessageCount: messages, presence: presence, presenceMax: maxPresenceEntries, presencePct: presPct, forwards: fwds, forwardsMax: maxForwardsEntries, forwardsPct: fwdPct, challenges: count core.c, challengesMax: maxChallenges, ackedIds: count core.a, ackedIdsMax: maxAckedIds, ackOverflow: count core.ao, ackOverflowMax: maxAckOverflow, deliveredIds: count core.d, deliveredIdsMax: maxDeliveredIds, forwardNonces: core.fnCount, forwardNoncesMax: maxForwardNonces, activePolls: count core.ap, activePollsMax: maxActivePolls, recipientCounts: count core.rc},
4923 capacity: {sessions: sessions, sessionsMax: maxSessions, sessionsPct: sessPct, messages: messages, messagesMax: maxMessages, messagesPct: msgPct, totalMessageCount: messages, presence: presence, presenceMax: maxPresenceEntries, presencePct: presPct, forwards: fwds, forwardsMax: maxForwardsEntries, forwardsPct: fwdPct, challenges: count core.c, challengesMax: maxChallenges, ackedIds: count core.a, ackedIdsMax: maxAckedIds, ackOverflow: count core.ao, ackOverflowMax: maxAckOverflow, deliveredIds: count core.d, deliveredIdsMax: maxDeliveredIds, forwardNonces: core.fnCount, forwardNoncesMax: maxForwardNonces, activePolls: count core.ap, activePollsMax: maxActivePolls, recipientCounts: count core.rc},
4782 federation: {enabled: federationEnabled, inFlight: count core.fi, hungSlots: hungFedSlots, avgLatencyMs: fedAvg, maxLatencyMs: counters.flm, callCount: counters.flc, forwardFailures: count core.ff, oldestForwardAge: oldestFwdAge, oldestAckAge: oldestAckAge, recentGossip: count core.rg, blockedPeers: count core.bp, drainCompletedAt: counters.dca, forwardServers: fwdServerCount, maxForwardQueueDepth: maxFwdDepth},
4924 federation: {enabled: federationEnabled, inFlight: count core.fi, hungSlots: hungFedSlots, avgLatencyMs: fedAvg, maxLatencyMs: counters.flm, callCount: counters.flc, forwardFailures: count core.ff, oldestForwardAge: oldestFwdAge, oldestAckAge: oldestAckAge, recentGossip: count core.rg, blockedPeers: count core.bp, drainCompletedAt: counters.dca, forwardServers: fwdServerCount, maxForwardQueueDepth: maxFwdDepth},
4783 rates: {send: count counters.sr, challenge: count counters.cr, federation: count counters.fr, federationPair: count counters.fpr, global: counters.grc, ip: count counters.ir, lookup: count counters.lr, recv: count counters.rr, pair: count counters.prr, gossipEvent: count core.ger, verifyIp: count core.vir, uniqueKeyIp: count core.ukir, pendingAckCleanups: count core.pac},
4925 rates: {send: count counters.sr, challenge: count counters.cr, federation: count counters.fr, federationPair: count counters.fpr, global: counters.grc, ip: count counters.ir, lookup: count counters.lr, recv: count counters.rr, pair: count counters.prr, gossipEvent: count core.ger, verifyIp: count core.vir, uniqueKeyIp: count core.ukir, pendingAckCleanups: count core.pac},
4784− timing: {pruneCount: counters.pc, pruneRunning: counters.pr, heartbeatCount: counters.hc, uptime: t - counters.sa, lastPruneAt: counters.lp, lastHeartbeatAt: counters.lh, draining: counters.dr}})
4926+ timing: {pruneCount: counters.pc, pruneRunning: counters.pr, heartbeatCount: counters.hc, uptime: t - counters.sa, lastPruneAt: counters.lp, lastHeartbeatAt: counters.lh, draining: counters.dr > 0}})
4785 else yield (respond zeroStats)
4927 else yield (respond zeroStats)
4787handleStatsPeers = \authorization respond -> do
4929handleStatsPeers = \authorization respond -> do
@@ -4872,8 +5014,8 @@ handleRestore = \version presenceData ackedIdsData ackOverflowData deliveredIdsD
4872 yield (respond {ok: False {}, error: "incompatible_version", presenceCount: 0, ackedIdsCount: 0, deliveredIdsCount: 0})
5014 yield (respond {ok: False {}, error: "incompatible_version", presenceCount: 0, ackedIdsCount: 0, deliveredIdsCount: 0})
4874 t <- now
5016 t <- now
4875− isDraining <- *draining
5017+ dp <- *drainPhase
4876− if not isDraining then yield (respond {ok: False {}, error: "server_must_be_draining", presenceCount: 0, ackedIdsCount: 0, deliveredIdsCount: 0})
5018+ if dp < 2 then yield (respond {ok: False {}, error: "server_must_be_draining", presenceCount: 0, ackedIdsCount: 0, deliveredIdsCount: 0})
4878 -- Validate imported data: reject entries with invalid or self-referencing server names.
5020 -- Validate imported data: reject entries with invalid or self-referencing server names.
4879 -- originServer in ForwardEntry is untyped Text (not ServerName), so not validated by the type system.
5021 -- originServer in ForwardEntry is untyped Text (not ServerName), so not validated by the type system.
@@ -4951,10 +5093,11 @@ handleAckMessages = \messageIds authorization clientIp respond -> do
4951 withSessionAuth authorization t clientIp
5093 withSessionAuth authorization t clientIp
4952 (yield (respond {ok: False {}, authorized: False {}, ackedCount: 0}))
5094 (yield (respond {ok: False {}, authorized: False {}, ackedCount: 0}))
4953 (\pubkey -> do
5095 (\pubkey -> do
5096+ -- Uses clientWriteRates (separate from lookupRates) so ack bursts don't starve lookups
4954 ackRateLimited <- atomic do
5097 ackRateLimited <- atomic do
4955− lookupRates <- *lookupRates
5098+ cwr <- *clientWriteRates
4956− let rateResult = checkPubkeyRate lookupRates pubkey maxLookupRatePerKey lookupRateWindowMs maxLookupRateEntries t
5099+ let rateResult = checkPubkeyRate cwr pubkey maxClientWriteRatePerKey clientWriteRateWindowMs maxClientWriteRateEntries t
4957− set *lookupRates = rateResult.entries
5100+ set *clientWriteRates = rateResult.entries
4958 yield rateResult.limited
5101 yield rateResult.limited
4959 if ackRateLimited then yield (respond {ok: False {}, authorized: True {}, ackedCount: 0})
5102 if ackRateLimited then yield (respond {ok: False {}, authorized: True {}, ackedCount: 0})
@@ -4992,8 +5135,8 @@ handleAckMessages = \messageIds authorization clientIp respond -> do
4992-- Rate-limited via lookup rate limiter (read-only info lookup).
5135-- Rate-limited via lookup rate limiter (read-only info lookup).
4993handleListMessages = \authorization clientIp respond -> do
5136handleListMessages = \authorization clientIp respond -> do
4994 t <- now
5137 t <- now
4995− draining <- *draining
5138+ dp <- *drainPhase
4996− if draining then yield (respond {messages: [], authorized: False {}, count: 0})
5139+ if dp >= 2 then yield (respond {messages: [], authorized: False {}, count: 0})
4998 globalRateLimited <- checkGlobalRate t
5141 globalRateLimited <- checkGlobalRate t
4999 if globalRateLimited then yield (respond {messages: [], authorized: False {}, count: 0})
5142 if globalRateLimited then yield (respond {messages: [], authorized: False {}, count: 0})
@@ -5022,8 +5165,8 @@ handleListMessages = \authorization clientIp respond -> do
5022-- their forward entries instead of retrying until prune.
5165-- their forward entries instead of retrying until prune.
5023handleDeleteMessages = \messageIds authorization clientIp respond -> do
5166handleDeleteMessages = \messageIds authorization clientIp respond -> do
5024 t <- now
5167 t <- now
5025− draining <- *draining
5168+ dp <- *drainPhase
5026− if draining then yield (respond {ok: False {}, authorized: False {}, deletedCount: 0})
5169+ if dp >= 2 then yield (respond {ok: False {}, authorized: False {}, deletedCount: 0})
5028 globalRateLimited <- checkGlobalRate t
5171 globalRateLimited <- checkGlobalRate t
5029 if globalRateLimited then yield (respond {ok: False {}, authorized: False {}, deletedCount: 0})
5172 if globalRateLimited then yield (respond {ok: False {}, authorized: False {}, deletedCount: 0})
@@ -5031,10 +5174,11 @@ handleDeleteMessages = \messageIds authorization clientIp respond -> do
5031 withSessionAuth authorization t clientIp
5174 withSessionAuth authorization t clientIp
5032 (yield (respond {ok: False {}, authorized: False {}, deletedCount: 0}))
5175 (yield (respond {ok: False {}, authorized: False {}, deletedCount: 0}))
5033 (\pubkey -> do
5176 (\pubkey -> do
5177+ -- Uses clientWriteRates (separate from lookupRates) so delete bursts don't starve lookups
5034 deleteRateLimited <- atomic do
5178 deleteRateLimited <- atomic do
5035− lookupRates <- *lookupRates
5179+ cwr <- *clientWriteRates
5036− let rateResult = checkPubkeyRate lookupRates pubkey maxLookupRatePerKey lookupRateWindowMs maxLookupRateEntries t
5180+ let rateResult = checkPubkeyRate cwr pubkey maxClientWriteRatePerKey clientWriteRateWindowMs maxClientWriteRateEntries t
5037− set *lookupRates = rateResult.entries
5181+ set *clientWriteRates = rateResult.entries
5038 yield rateResult.limited
5182 yield rateResult.limited
5039 if deleteRateLimited then yield (respond {ok: False {}, authorized: True {}, deletedCount: 0})
5183 if deleteRateLimited then yield (respond {ok: False {}, authorized: True {}, deletedCount: 0})
@@ -5089,7 +5233,10 @@ handleSendBatch = \messages authorization clientIp respond -> do
5089 -- dedup between messages in the same batch is not caught by STM.
5233 -- dedup between messages in the same batch is not caught by STM.
5090 -- Messages with empty messageId (auto-generated) are always kept.
5234 -- Messages with empty messageId (auto-generated) are always kept.
5091 let capped = fold (\acc m -> if m.messageId != "" && any acc (\a -> a.messageId == m.messageId) then acc else union acc [m]) [] sized
5235 let capped = fold (\acc m -> if m.messageId != "" && any acc (\a -> a.messageId == m.messageId) then acc else union acc [m]) [] sized
5236+ -- Sort by pubkeyBucket(to) so sequential atomic blocks in sendBatchInner benefit
5237+ -- from STM cache locality when multiple messages target the same recipient bucket.
5238+ let sorted = sortBy (\m -> pubkeyBucket m.to) capped
5092− results <- sendBatchInner senderKey capped t []
5239+ results <- sendBatchInner senderKey sorted t []
5093 -- Fork federation forwarding: one forked task per target server (not per message).
5240 -- Fork federation forwarding: one forked task per target server (not per message).
5094 -- Within each task, forwards are sent sequentially to avoid per-message fork overhead.
5241 -- Within each task, forwards are sent sequentially to avoid per-message fork overhead.
5095 let forwardTargets = do r <- results; where count r.remoteServers > 0; yield r
5242 let forwardTargets = do r <- results; where count r.remoteServers > 0; yield r
@@ -5125,8 +5272,8 @@ sendBatchInner = \senderKey messages t acc ->
5126 let clampedTs = clampTimestamp msg.timestamp t
5273 let clampedTs = clampTimestamp msg.timestamp t
5127 sendResult <- atomic do
5274 sendResult <- atomic do
5128− draining <- *draining
5275+ dp <- *drainPhase
5129− if draining then yield {status: MsDraining {}, remoteServers: [], retryAfterMs: 0}
5276+ if dp >= 2 then yield {status: MsDraining {}, remoteServers: [], retryAfterMs: 0}
5130 else sendMessageCore senderKey msg.to msg.encryptedBlob clampedTs effectiveMsgId t
5277 else sendMessageCore senderKey msg.to msg.encryptedBlob clampedTs effectiveMsgId t
5131 case sendResult.status of
5278 case sendResult.status of
5132 MsRateLimited {} -> logWarn ("sendBatch: rate limited for " ++ effectiveMsgId)
5279 MsRateLimited {} -> logWarn ("sendBatch: rate limited for " ++ effectiveMsgId)
@@ -5342,6 +5489,7 @@ requestPath = \req -> case req of
5342 VerifyToken {} -> "POST /federation/verify"
5489 VerifyToken {} -> "POST /federation/verify"
5343 Prune {} -> "POST /maintenance/prune"
5490 Prune {} -> "POST /maintenance/prune"
5344 Drain {} -> "POST /maintenance/drain"
5491 Drain {} -> "POST /maintenance/drain"
5492+ Shutdown {} -> "POST /maintenance/shutdown"
5345 Undrain {} -> "POST /maintenance/undrain"
5493 Undrain {} -> "POST /maintenance/undrain"
5346 Snapshot {} -> "GET /maintenance/snapshot"
5494 Snapshot {} -> "GET /maintenance/snapshot"
5347 Restore {} -> "POST /maintenance/restore"
5495 Restore {} -> "POST /maintenance/restore"
@@ -5379,6 +5527,7 @@ requestClientIp = \req -> case req of
5379 FederationAck {} -> ""
5527 FederationAck {} -> ""
5380 Prune {} -> ""
5528 Prune {} -> ""
5381 Drain {} -> ""
5529 Drain {} -> ""
5382 Undrain {} -> ""
5531 Undrain {} -> ""
5383 Snapshot {} -> ""
5532 Snapshot {} -> ""
5384 Restore {} -> ""
5533 Restore {} -> ""
@@ -5418,7 +5567,8 @@ serveInner = \req -> case req of
5418 FederationAck {messageIds, originServer, fromServer, requestTimestamp, correlationId, authorization, respond} -> handleFederationAck messageIds originServer fromServer requestTimestamp correlationId authorization respond
5567 FederationAck {messageIds, originServer, fromServer, requestTimestamp, correlationId, authorization, respond} -> handleFederationAck messageIds originServer fromServer requestTimestamp correlationId authorization respond
5419 VerifyToken {token, scope, clientIp, respond} -> handleVerifyToken token scope clientIp respond
5568 VerifyToken {token, scope, clientIp, respond} -> handleVerifyToken token scope clientIp respond
5420 Prune {authorization, respond} -> handlePrune authorization respond
5569 Prune {authorization, respond} -> handlePrune authorization respond
5421− Drain {authorization, respond} -> handleDrain authorization respond
5570+ Drain {phase, authorization, respond} -> handleDrain phase authorization respond
5571+ Shutdown {authorization, respond} -> handleShutdown authorization respond
5422 Undrain {authorization, respond} -> handleUndrain authorization respond
5572 Undrain {authorization, respond} -> handleUndrain authorization respond
5423 Snapshot {authorization, respond} -> handleSnapshot authorization respond
5573 Snapshot {authorization, respond} -> handleSnapshot authorization respond
5424 Restore {version, presence, ackedIds, ackOverflow, deliveredIds, forwardFailures, forwards, forwardNonces, blockedPeers, authorization, respond} -> handleRestore version presence ackedIds ackOverflow deliveredIds forwardFailures forwards forwardNonces blockedPeers authorization respond
5574 Restore {version, presence, ackedIds, ackOverflow, deliveredIds, forwardFailures, forwards, forwardNonces, blockedPeers, authorization, respond} -> handleRestore version presence ackedIds ackOverflow deliveredIds forwardFailures forwards forwardNonces blockedPeers authorization respond
@@ -5470,8 +5620,12 @@ serve = \req -> do
5470backgroundPrune = do
5620backgroundPrune = do
5471 jitter <- randomInt 5000
5621 jitter <- randomInt 5000
5472 sleep (backgroundPruneInterval + jitter)
5622 sleep (backgroundPruneInterval + jitter)
5473− draining <- *draining
5623+ dp <- *drainPhase
5474− if draining then do
5624+ -- Check shutdown: if drain completed and shutdown requested, exit the loop
5625+ shutdownReq <- *shutdownRequested
5626+ dca <- *drainCompletedAt
5627+ if shutdownReq && dca > 0 then logInfo "backgroundPrune: shutdown complete, exiting loop"
5628+ else if dp >= 2 then do
5475 sleep backgroundPruneInterval
5629 sleep backgroundPruneInterval
5476 backgroundPrune
5630 backgroundPrune
@@ -5545,8 +5699,11 @@ backgroundPrune = do
5545pollHeartbeat = do
5699pollHeartbeat = do
5546 jitter <- randomInt 2000
5700 jitter <- randomInt 2000
5547 sleep (pollTimeoutMs + jitter)
5701 sleep (pollTimeoutMs + jitter)
5548− draining <- *draining
5702+ dp <- *drainPhase
5549− if draining then do
5703+ shutdownReq <- *shutdownRequested
5704+ dca <- *drainCompletedAt
5705+ if shutdownReq && dca > 0 then logInfo "pollHeartbeat: shutdown complete, exiting loop"
5706+ else if dp >= 2 then do
5550 t <- now
5707 t <- now
5551 atomic do set *pollWakeup = t
5708 atomic do set *pollWakeup = t
5552 -- Don't exit; sleep and retry so undrain can resume the loop
5709 -- Don't exit; sleep and retry so undrain can resume the loop
@@ -5569,8 +5726,11 @@ pollHeartbeat = do
5569-- (e.g., /var/lib/skrepka/snapshot.json) in addition to logging.
5726-- (e.g., /var/lib/skrepka/snapshot.json) in addition to logging.
5570backgroundSnapshot = do
5727backgroundSnapshot = do
5571 sleep autoSnapshotInterval
5728 sleep autoSnapshotInterval
5572− draining <- *draining
5729+ dp <- *drainPhase
5573− if not draining then do
5730+ shutdownReq <- *shutdownRequested
5731+ dca <- *drainCompletedAt
5732+ if shutdownReq && dca > 0 then logInfo "backgroundSnapshot: shutdown complete, exiting loop"
5733+ else if dp < 2 then do
5574 result <- atomic do
5734 result <- atomic do
5575 tmc <- *totalMessageCount
5735 tmc <- *totalMessageCount
5576 sc <- *totalSessionCount
5736 sc <- *totalSessionCount
@@ -5589,7 +5749,12 @@ backgroundSnapshot = do
5589 when (autoSnapshotExportInterval > 0) (do
5749 when (autoSnapshotExportInterval > 0) (do
5590 lastAutoExport <- *lastAutoSnapshotExportAt
5750 lastAutoExport <- *lastAutoSnapshotExportAt
5591 let hasState = result.presence > 0 || result.forwards > 0 || result.ackedIds > 0 || result.deliveredIds > 0
5751 let hasState = result.presence > 0 || result.forwards > 0 || result.ackedIds > 0 || result.deliveredIds > 0
5752+ -- Fingerprint: skip export when key counters haven't changed since last export.
5753+ -- Uses multiplicative hash to distinguish {messages:1,sessions:0} from {messages:0,sessions:1}.
5754+ let fingerprint = result.messages + result.sessions * 100003 + result.presence * 100019 + result.forwards * 100043 + result.ackedIds * 100049 + result.deliveredIds * 100057
5755+ lastFp <- *lastSnapshotFingerprint
5756+ let stateChanged = lastFp == 0 || fingerprint != lastFp
5592− when (hasState && (lastAutoExport == 0 || t - lastAutoExport >= autoSnapshotExportInterval)) (do
5757+ when (hasState && stateChanged && (lastAutoExport == 0 || t - lastAutoExport >= autoSnapshotExportInterval)) (do
5593 exportData <- atomic do
5758 exportData <- atomic do
5594 p <- readAllPresence
5759 p <- readAllPresence
5595 a <- readAllAcked
5760 a <- readAllAcked
@@ -5601,7 +5766,9 @@ backgroundSnapshot = do
5601 bp <- *blockedPeers
5766 bp <- *blockedPeers
5602 yield {version: snapshotVersion, presence: p, ackedIds: a, ackOverflow: ao, deliveredIds: d, forwardFailures: ff, forwards: f, forwardNonces: fn, blockedPeers: bp}
5767 yield {version: snapshotVersion, presence: p, ackedIds: a, ackOverflow: ao, deliveredIds: d, forwardFailures: ff, forwards: f, forwardNonces: fn, blockedPeers: bp}
5603 println ("[SNAPSHOT_DATA] " ++ toJson exportData)
5768 println ("[SNAPSHOT_DATA] " ++ toJson exportData)
5604− atomic do set *lastAutoSnapshotExportAt = t
5770+ set *lastAutoSnapshotExportAt = t
5771+ set *lastSnapshotFingerprint = fingerprint
5605 logInfo ("snapshot: auto-exported full state (" ++ show result.presence ++ " presence, " ++ show result.forwards ++ " forwards)")))
5772 logInfo ("snapshot: auto-exported full state (" ++ show result.presence ++ " presence, " ++ show result.forwards ++ " forwards)")))
5606 -- Warn if federation state hasn't been exported recently (will be lost on restart)
5773 -- Warn if federation state hasn't been exported recently (will be lost on restart)
5607 lastExport <- *lastSnapshotExportAt
5774 lastExport <- *lastSnapshotExportAt
@@ -5679,6 +5846,15 @@ startupProbeNext = \peers probed ->
5679 Ok r -> logInfo ("startup: peer " ++ server ++ " reachable (" ++ show latency ++ "ms)")
5846 Ok r -> logInfo ("startup: peer " ++ server ++ " reachable (" ++ show latency ++ "ms)")
5680 startupProbeNext rest (probed + 1)
5847 startupProbeNext rest (probed + 1)
5849+-- Reset federation latency metrics after startup probes so unauthenticated probe
5850+-- results (expected "unauthorized" responses) don't skew initial /stats readings.
5851+resetFederationLatencyMetrics = atomic do
5852+ set *federationLatencyTotal = 0
5853+ set *federationLatencyCount = 0
5854+ set *federationLatencyMax = 0
5856+ set *federationLatencyResetAt = t
5682main = do
5858main = do
5683 let configOk = maxMessages > 0 && maxMessagesPerRecipient > 0 && maxPresenceEntries > 0
5859 let configOk = maxMessages > 0 && maxMessagesPerRecipient > 0 && maxPresenceEntries > 0
5684 && maxSessions > 0 && maxChallenges > 0 && backgroundPruneInterval > 0
5860 && maxSessions > 0 && maxChallenges > 0 && backgroundPruneInterval > 0
@@ -5767,5 +5943,6 @@ main = do
5767 when federationEnabled (fork (do
5943 when federationEnabled (fork (do
5768 sleep 5000 -- wait for listener to be ready
5944 sleep 5000 -- wait for listener to be ready
5769 logInfo "startup: probing federation peers..."
5945 logInfo "startup: probing federation peers..."
5770− startupProbePeers))
5946+ startupProbePeers
5947+ resetFederationLatencyMetrics))
5771 listen 8080 serve
5948 listen 8080 serve