fix(api): globalize WS connection/subscription caps via SharedStore [BUG-005] - #214
Open
Morenikeoa wants to merge 1 commit into
Open
fix(api): globalize WS connection/subscription caps via SharedStore [BUG-005]#214Morenikeoa wants to merge 1 commit into
Morenikeoa wants to merge 1 commit into
Conversation
…BUG-005] MAX_WS_CONNECTIONS and MAX_GLOBAL_SUBSCRIPTIONS were enforced against process-local state (clients.size, a plain globalSubscriptionCount variable) — under N horizontally-scaled replicas, the true fleet-wide cap became N×limit instead of limit, since each replica only sees its own slice. /ws/stats and /health's WS-utilization check had the same problem: totalConnections only reflected the answering replica's own clients.size. PR dcccrypto#189 already solved this exact class of bug for per-IP connection counts, rate-limit buckets, and auth-failure bans via a SharedStore abstraction (Upstash Redis when configured, in-memory fallback otherwise). This extends that same mechanism to the two caps it didn't cover: - Added addConnectionCount(key, delta) to SharedStore (both InMemoryStore and UpstashStore) — a signed-delta batch update, so the subscribe/unsubscribe loops (up to 50 channels per message) can flush one update after the loop instead of up to 50 sequential round-trips. incrementConnectionCount/decrementConnectionCount are unchanged and remain the right choice for the connection cap, which only changes by ±1 per connect/disconnect. - The subscription cap is checked once per message (read the global count at message start, track an in-loop local delta, compare against count+delta per channel) rather than once per channel — avoiding the round-trip cost while keeping the same non-atomic check-then-act semantics the per-IP caps in this file already use elsewhere. - getWebSocketMetrics() is now async and reports the true global totalConnections/totalSubscriptions from the shared store (plus the previously-unreported maxGlobalSubscriptions limit). connectionsPerSlab, messagesPerSec, and bytesPerSec stay local — they're inherently per-replica observability data, not safety caps, and aggregating them is a different, larger problem than this fix's scope. Also fixed clients.size >= MAX_WS_CONNECTIONS (line ~609), found while scoping this: identical per-replica-only bug, same file, same mechanism — fixing it alongside rather than leaving an identical bug next to the one being fixed. Added regression tests that seed the shared store directly (simulating "other replicas already at the cap") while this replica's own local state stays empty, proving the cap is now enforced fleet-wide rather than per-replica. Verified all three fail against the pre-fix code and pass against the fix. Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
|
@Princessdada is attempting to deploy a commit to the Khubair Nasir's projects Team on Vercel. A member of the Team first needs to authorize it. |
📝 WalkthroughWalkthroughAdds an atomic shared-store counter delta API, routes WebSocket connection and subscription limits through fleet-wide shared counters, makes WebSocket metrics async, and updates health endpoints and tests to use the resolved shared metrics. ChangesFleet-wide WebSocket counters
Sequence Diagram(s)sequenceDiagram
participant Client
participant WS as src/routes/ws.ts
participant Store as SharedStore
participant Health as src/routes/health.ts
participant Metrics as getWebSocketMetrics
Client->>WS: open WebSocket
WS->>Store: read global connection count
Store-->>WS: totalConnections
WS->>Store: addConnectionCount(globalConnections, 1)
WS-->>Client: accept or close 1008
Health->>Metrics: await metrics
Metrics->>Store: read totalConnections and totalSubscriptions
Store-->>Metrics: metrics
Metrics-->>Health: metrics
sequenceDiagram
participant Client
participant WS as src/routes/ws.ts
participant Store as SharedStore
Client->>WS: subscribe channels
WS->>Store: read global subscription count
Store-->>WS: totalSubscriptions
WS->>Store: addConnectionCount(globalSubscriptions, acceptedDelta)
WS-->>Client: subscribed or server subscription limit reached
Client->>WS: unsubscribe channels
WS->>Store: addConnectionCount(globalSubscriptions, -removedDelta)
Client->>WS: socket close
WS->>Store: addConnectionCount(globalConnections, -1)
WS->>Store: addConnectionCount(globalSubscriptions, -clientSubscriptions)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/routes/ws.ts`:
- Around line 895-907: The global subscription cap logic in the websocket
subscribe flow still uses a stale snapshot of GLOBAL_SUB_KEY, so concurrent
replicas can oversubscribe. Update the subscribe paths in ws.ts around the
current getSharedStore/globalSubsAtStart handling to use an atomic reservation
or capped increment in the shared store before mutating client.subscriptions,
and apply the same fix to the other subscribe blocks referenced in the comment
so MAX_GLOBAL_SUBSCRIPTIONS is enforced fleet-wide.
- Around line 628-631: The fleet-wide WS connection limit is still vulnerable to
check-then-act because `getConnectionCount(GLOBAL_CONN_KEY)` is read before the
later increment path, so concurrent upgrades can all pass the cap. Update
`src/routes/ws.ts` to use an atomic reserve/rollback flow through `SharedStore`
instead of reading `GLOBAL_CONN_KEY` and incrementing later in separate steps.
The fix should be centered around the global connection check near
`MAX_WS_CONNECTIONS` and the increment logic around the later connection setup
path, ensuring only one replica can reserve a slot when the cap is reached and
that any failed upgrade releases the reservation.
- Around line 707-709: The WebSocket connection setup in ws.ts should attach the
close cleanup handler before awaiting
getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY). Move the socket
close listener registration earlier in the connection flow, before the shared
counter write, so a disconnect during that await still triggers cleanup of
clients, the per-IP count, and the global count. Keep the fix localized around
the connection setup logic near logger.info and the shared store increment.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 76d9a320-c571-4567-b3b0-fec62e6a8208
📒 Files selected for processing (5)
src/middleware/shared-store.tssrc/routes/health.tssrc/routes/ws.tstests/middleware/shared-store.test.tstests/routes/ws-global-caps.test.ts
Comment on lines
+628
to
+631
| // H2: Reject if at max connections (global across all replicas, via SharedStore) | ||
| const globalConnCount = await getSharedStore().getConnectionCount(GLOBAL_CONN_KEY); | ||
| if (globalConnCount >= MAX_WS_CONNECTIONS) { | ||
| logger.warn("Max global WS connections reached", { ip: clientIp, count: globalConnCount }); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
The fleet-wide connection cap is still check-then-act.
Line 629 reads GLOBAL_CONN_KEY, and Line 707 increments it later. Concurrent upgrades on different replicas can all observe the same pre-cap value and all pass, so MAX_WS_CONNECTIONS can still be exceeded under burst load. This needs an atomic reserve/rollback path in SharedStore, not a read followed by a later increment.
Also applies to: 707-709
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/routes/ws.ts` around lines 628 - 631, The fleet-wide WS connection limit
is still vulnerable to check-then-act because
`getConnectionCount(GLOBAL_CONN_KEY)` is read before the later increment path,
so concurrent upgrades can all pass the cap. Update `src/routes/ws.ts` to use an
atomic reserve/rollback flow through `SharedStore` instead of reading
`GLOBAL_CONN_KEY` and incrementing later in separate steps. The fix should be
centered around the global connection check near `MAX_WS_CONNECTIONS` and the
increment logic around the later connection setup path, ensuring only one
replica can reserve a slot when the cap is reached and that any failed upgrade
releases the reservation.
Comment on lines
+707
to
+709
| await getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY); | ||
|
|
||
| logger.info("WebSocket connection established", { |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Register close before awaiting the shared counter write.
Line 707 can block on the shared store after the client has already been added locally and after the per-IP counter was incremented, but the close listener is only attached much later. If the peer disconnects during that await, cleanup never runs, so clients, the per-IP count, and potentially the global count leak until restart.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/routes/ws.ts` around lines 707 - 709, The WebSocket connection setup in
ws.ts should attach the close cleanup handler before awaiting
getSharedStore().incrementConnectionCount(GLOBAL_CONN_KEY). Move the socket
close listener registration earlier in the connection flow, before the shared
counter write, so a disconnect during that await still triggers cleanup of
clients, the per-IP count, and the global count. Keep the fix localized around
the connection setup logic near logger.info and the shared store increment.
Comment on lines
+895
to
+907
|
|
||
| // Check the global subscription cap once per message (not once | ||
| // per channel) and track local accepted-count as a delta, then | ||
| // flush a single batched update after the loop. This avoids up to | ||
| // MAX_CHANNELS_PER_MESSAGE (50) shared-store round-trips per | ||
| // message. The existing per-IP caps in this file are already | ||
| // non-atomic check-then-act, so the small intra-message overshoot | ||
| // window this introduces (bounded by MAX_CHANNELS_PER_MESSAGE) is | ||
| // not a new class of weaker guarantee. | ||
| const subStore = getSharedStore(); | ||
| const globalSubsAtStart = await subStore.getConnectionCount(GLOBAL_SUB_KEY); | ||
| let newSubsThisMessage = 0; | ||
|
|
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
The fleet-wide subscription cap can still overshoot under concurrency.
Both subscribe paths snapshot GLOBAL_SUB_KEY once and apply the delta afterward. Multiple clients/replicas can therefore all start below the cap and all commit, so MAX_GLOBAL_SUBSCRIPTIONS is not actually enforced fleet-wide. This needs an atomic capped-add/reservation in the shared store before mutating client.subscriptions.
Also applies to: 954-984, 1086-1106
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/routes/ws.ts` around lines 895 - 907, The global subscription cap logic
in the websocket subscribe flow still uses a stale snapshot of GLOBAL_SUB_KEY,
so concurrent replicas can oversubscribe. Update the subscribe paths in ws.ts
around the current getSharedStore/globalSubsAtStart handling to use an atomic
reservation or capped increment in the shared store before mutating
client.subscriptions, and apply the same fix to the other subscribe blocks
referenced in the comment so MAX_GLOBAL_SUBSCRIPTIONS is enforced fleet-wide.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
MAX_WS_CONNECTIONSandMAX_GLOBAL_SUBSCRIPTIONSwere enforced against process-local state —clients.size(aSet<WsClient>) andglobalSubscriptionCount(a plain module-levellet). Under N horizontally-scaled replicas, each replica only sees its own connections/subscriptions, so the true fleet-wide cap becomes N×limit instead of limit — e.g. with the 1000 default and 5 replicas, the fleet actually allows ~5000 before any individual replica starts rejecting new connections.Separately,
GET /ws/statsand the WS-utilization check insideGET /healthonly reportedtotalConnectionsfrom the answering replica's ownclients.size, giving operators an incomplete picture of true fleet-wide load.While scoping this I also found
clients.size >= MAX_WS_CONNECTIONS(the connection cap, distinct from the subscription cap that was the originally-confirmed finding) has the identical per-replica-only bug — same file, same root cause. Fixed both rather than leaving one sitting next to the other.Fix
PR #189 already solved this exact class of bug for per-IP connection counts, rate-limit buckets, and auth-failure bans via a
SharedStoreabstraction (Upstash Redis when configured, in-memory fallback for single-replica/dev). This extends that same mechanism to the two caps it didn't cover:addConnectionCount(key, delta)toSharedStore(bothInMemoryStoreandUpstashStore) — a signed-delta batch update. The subscribe/unsubscribe message handlers loop over up to 50 channels per message; without batching, enforcing the cap via the shared store would mean up to 50 sequential round-trips per message. Now: read the global count once at the start of the loop, track a local delta, flush one batched update at the end.incrementConnectionCount/decrementConnectionCountare untouched and remain correct for the connection cap, which only changes by ±1 per connect/disconnect (no batching needed there).getWebSocketMetrics()is nowasyncand reports the true globaltotalConnections/totalSubscriptions(plus the previously-unreportedmaxGlobalSubscriptionslimit).connectionsPerSlab,messagesPerSec,bytesPerSecstay local/per-replica — they're inherently per-process observability (this replica's own traffic distribution), not safety caps, and aggregating them across replicas is a different, larger problem than this fix's scope.health.tscall sites just neededawaitadded — no other changes, since both were already insideasynchandlers with existingtry/catch.Proof of Fix
3 new tests simulate "other replicas already at the cap" by seeding the shared store directly while this replica's own local state (
clients.size, etc.) stays empty — proving the cap now reads from the shared, fleet-wide count rather than this replica's empty local one:MAX_WS_CONNECTIONS=2, new connection from this (locally-empty) replica is rejected with close code 1008.getWebSocketMetrics().totalConnectionsreturns 7 even though this replica's ownclients.sizeis 0.Verified these are genuine regression tests: reverted the 3 source files and reran — all 3 failed exactly as expected (one timed out waiting for a rejection that never arrived, one hit a missing method, one asserted 0 instead of 7). Restored the fix and they pass.
tests/routes/health.test.ts's mocks ofgetWebSocketMetrics(one synchronous return, one synchronous throw) needed no changes despite the signature becomingasync—awaiton a non-Promise value resolves fine, and a synchronous throw inside an awaited call is still caught by the surroundingtry/catch.tsc --noEmitclean (no separate lint script in this repo).Test Output
Full suite: 304/305 passed (294 baseline + 10 new). The 1 failure (
tests/sdk-smoke.test.ts) is pre-existing and unrelated — it asserts on an exact@percolatorct/sdkerror-message string that has drifted from the locally-resolved SDK version in this environment.Related
Found during a broader API audit; no existing open issue/PR covers this.
Summary by CodeRabbit
New Features
Bug Fixes
Tests