feat(cubestore): cap concurrent websocket connections per user - #11667
Conversation
Adds CUBESTORE_MAX_WS_CONNECTIONS_PER_USER, an optional per-node cap on how many concurrent WebSocket connections one authenticated user may hold. A client that opens connections faster than it closes them can otherwise consume the whole process file descriptor table, which affects every other user of that node. At the limit the user's oldest connection is closed to admit the new one, rather than the new one being refused: a client that still needs the closed connection reconnects, while one that had forgotten about it simply loses it. It is closed with a close frame, so the client sees a clean close and can reconnect on its own terms. A slot is taken at the websocket upgrade and released by a guard held for the lifetime of the connection task, so it comes back on any exit path: a close frame, a transport error, a panic, a cancelled task, or a connection that dies before the upgrade completes. Live connections are keyed by a monotonic id, so an evicting admission and the victim's own later drop cannot remove each other's entry, and the limit never overshoots. The limit defaults to 0, which disables the cap and leaves behaviour unchanged.
|
Preview deployment for your docs. Learn more about Mintlify Previews.
💡 Tip: Enable Workflows to automatically generate PRs for you. |
|
Claude finished @waralexrom's task in 2m 5s —— View job Review complete — 1 medium-high, 1 medium, 3 low. Solid design; the eviction path can hang and the cap is inert without auth.Full review
What's goodThe guard design is the right shape for this problem. Releasing from
Recovering the poisoned lock instead of unwrapping ( Findings
Smaller notes
|
| .and_then(move |tx: mpsc::Sender<(mpsc::Sender<Arc<HttpMessage>>, SqlQueryContext, HttpMessage)>, sql_query_context: SqlQueryContext, ws_connections: Arc<WsConnectionCounter>, ws: Ws| async move { | ||
| let tx_to_move = tx.clone(); | ||
| let sql_query_context = sql_query_context.clone(); | ||
| let connection_guard = ws_connections.acquire(sql_query_context.user.as_deref()); |
There was a problem hiding this comment.
sql_query_context.user is None whenever the request carries no auth header — authorize returns Ok(None) in that case even when SqlAuthService reports no password for the user (see authorize, ~L918-940). So on a Cube Store deployment without credentials configured (the common single-tenant setup), every WebSocket connection is untracked and the cap is a no-op regardless of what CUBESTORE_MAX_WS_CONNECTIONS_PER_USER is set to.
That's a defensible scope — you can't attribute an anonymous connection to a tenant — but the stated motivation is fd exhaustion, which is just as reachable without auth. Two things worth considering:
- Log once at startup (or on the first untracked connection) when the limit is set but connections arrive unauthenticated, so an operator who set the variable doesn't believe they're protected when they aren't.
- Either a global cap or an
anonymousbucket as a fallback would actually close the hole; a per-node total is cheap to add on top of the same counter.
At minimum the docs entry should say explicitly that unauthenticated connections are exempt — "a single authenticated user" implies it but a reader setting the variable for fd protection will likely miss it.
Sending the close frame flushes, so a peer whose receive window is closed would
park the connection task indefinitely with the descriptor still held. That
matters more than a stray task: the entry leaves the counter when the
connection is evicted rather than when its task ends, so the cap would read as
satisfied while the process held more connections than it allows. The send is
now bounded, and the bound is deliberately short — the steady-state overshoot
is the eviction rate times this timeout, and a close frame is a handful of
bytes, so a peer that cannot take them in two seconds is not going to.
The frame now carries close code 1013 ("Try Again Later") and a reason, so a
client can tell being recycled from a transport failure and reconnect rather
than report an error.
Adds an end-to-end test: with a cap of one, a second connection for the same
authenticated user evicts the first, which receives that close code and reason,
while a different user connects unaffected. It covers the parts the unit tests
cannot reach — the select branch, the close path and the guard living in the
connection task. The read is bounded, so a regression in the eviction path
fails the test rather than hanging it.
|
Claude finished @waralexrom's task in 2m 7s —— View job Re-review of
|
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | Medium | http/mod.rs:395 |
Carried over, unresolved. sql_query_context.user is None whenever no Authorization header is present — I re-read authorize (http/mod.rs:946-968) to confirm: credentials is None, so both the authenticated and unauthenticated arms return Ok(None). On a deployment without credentials, the cap is inert regardless of what the variable is set to. Existing thread — at minimum the docs entry should say unauthenticated connections are exempt. Fix this → |
| 2 | Low | http/mod.rs:403 |
tokio::select! picks branches at random, so an evicted connection with a busy response_rx can keep serving responses before it notices the cancellation. Adding biased; with the eviction branch first makes "evicted means evicted" deterministic at no cost — and it also tightens the overshoot bound the new EVICTED_CLOSE_TIMEOUT comment reasons about, which currently assumes the branch is taken promptly. |
| 3 | Low | http/mod.rs:406, 419 |
Style nit only: the eviction arm mixes bare error! (imported at the top) with fully-qualified log::warn!. Elsewhere the file uses the bare macros. |
Smaller notes
- The
send-then-breakpath never waits for the client's close reply before dropping the socket. Strictly, RFC 6455 has the initiator wait; in practice every client treats a received close frame plus FIN as a clean close, and waiting is precisely what the fix removed. Correct tradeoff, just recording it. - If the timeout fires mid-flush, the frame is dropped and the peer sees a bare FIN — the degraded case is exactly the old behaviour, which is the right floor.
WsConnectionCounter::count(http/mod.rs:235) still has no non-test consumer.- I did not run
cargo check/cargo testfor cubestore here (build time in this environment). The PR reports both clean; nothing in the new diff looks like it would fail to compile.
The counter's critical section is a map lookup and an insert with nothing awaited inside it — `acquire` is not an `async fn`, so nothing can be — which is the shape the `tokio::sync::Mutex` docs point at a standard-library lock for. That makes a stall unreachable by reasoning about the callers; taking the lock with a bound makes it unreachable without needing that reasoning at all. Switches to parking_lot for `try_lock_for`. It also removes poisoning, which had to be recovered by hand before: there is nothing meaningful to recover for a map of live connections, and a poisoned unwrap in the destructor would have aborted the process. parking_lot is already in the tree at this version, so the lockfile gains one dependency edge and no crate. Both callers fail open, because the cap is a safety net rather than a correctness invariant: an unavailable counter admits the connection untracked instead of refusing or stalling it, and a finished connection leaves its entry behind to be evicted — that entry is the oldest one of its user by construction, and cancelling a finished connection's token is a no-op. Both paths log, so a bound that is ever reached is visible instead of quietly weakening the cap. The bound is 100ms rather than something tighter because a thread holding the lock can be descheduled for milliseconds under CPU pressure, and giving up then would stop counting connections exactly when the cap matters most.
|
Claude finished @waralexrom's task in 2m 0s —— View job Re-review of
|
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | Low | http/mod.rs:302-304 |
The Drop timeout comment asserts the stranded entry "is the oldest one of that user by construction" — it isn't. Connections don't finish in admission order, so a mid-list entry can be stranded, and the next admission then evicts a live connection while the stale entry keeps a slot until churn reaches it. Self-healing and effectively unreachable, so no code change needed; the comment (and the commit message, which repeats the claim) shouldn't state an invariant the type doesn't have. inline |
Still open from earlier passes
| # | Severity | Where | Issue |
|---|---|---|---|
| 2 | Medium | http/mod.rs:437 |
Re-affirmed, not re-posted. sql_query_context.user is None when no Authorization header is present, so the cap is inert on deployments without credentials. I re-checked the docs entry in this pass — it still says "a single authenticated user" without stating that unauthenticated connections are exempt. Existing thread · Fix this → |
| 3 | Low | http/mod.rs:445 |
tokio::select! still picks branches at random; biased; with the eviction arm first would make "evicted means evicted" deterministic and tighten the overshoot bound EVICTED_CLOSE_TIMEOUT reasons about. |
Smaller notes
acquire_gives_up_on_a_stuck_lock_instead_of_waiting_on_it(http/mod.rs:1621) is a good test — it holds the lock from the main thread and probes from another, so it asserts the fail-open path rather than reentrancy. Two nits: the assert message has a stray run of spaces mid-sentence ("...admits the connection untracked...", L1635), and it costs a real 100ms of suite time. Nothing tests theDroptimeout path, which is where finding ReferenceError: self is not defined - Issue with whatwg-fetch when using the core lib with Node.js #1 lives.count(http/mod.rs:277) now returns0on lock timeout, conflating "unknown" with "none". Test-only consumer, so it doesn't matter today — just noting the collapse is silent.- If the 100ms bound were ever actually reached,
try_lock_forblocks a tokio worker thread for that long rather than yielding. That's the accepted cost of releasing from a destructor, and it's now bounded where before it wasn't — recording it rather than objecting. - No threads became stale in this pass, so nothing was resolved. The one resolved thread (the unbounded
close()) was already resolved in the prior pass. - I did not run
cargo check/cargo testfor cubestore here (build time in this environment). The PR reports both clean; nothing in the new diff looks like it would fail to compile —try_lock_foris available onparking_lot::Mutexwith default features.
| // The entry stays behind, but it is the oldest one of that user by | ||
| // construction, so the next admission evicts it: cancelling an | ||
| // already finished connection's token is a no-op. |
There was a problem hiding this comment.
The "oldest by construction" claim doesn't hold. Connections don't finish in the order they were admitted: with ids 3, 4, 5 live for one user, id 5 can be the one whose drop hits this branch, and its stale entry then sits behind 3 and 4 in the BTreeMap.
The consequence is slightly worse than "evicted on the next admission": the next admission picks 3 — a live, working connection — and closes it, while the stale entry keeps consuming a slot. It only clears once churn makes it the oldest, so until then the user's effective cap is limit - 1 and one legitimate connection is recycled early for nothing.
Self-healing and effectively unreachable, so I'd leave the code alone — but the comment shouldn't assert an invariant the type doesn't have. Something like:
| // The entry stays behind, but it is the oldest one of that user by | |
| // construction, so the next admission evicts it: cancelling an | |
| // already finished connection's token is a no-op. | |
| // The entry stays behind. It is not necessarily that user's oldest | |
| // — connections do not finish in admission order — so until churn | |
| // makes it the oldest it holds a slot, and the admission that would | |
| // have taken it evicts a live connection instead. Cancelling a | |
| // finished connection's token is a no-op, so it clears harmlessly. |
Same wording appears in the commit message ("that entry is the oldest one of its user by construction"), worth fixing there too if you amend.
…e taken The bound exists to make a stall impossible without reasoning about the callers, not to be tight, and the failure it guards against is one-sided: a thread holding the lock can be descheduled for a long time under CPU pressure, and giving up then would quietly stop counting connections exactly when the cap matters most. A second is far past anything a map lookup and an insert can take, and giving up is logged either way.
|
Claude finished @waralexrom's task in 1m 23s —— View job Re-review of
|
| # | Severity | Where | Issue |
|---|---|---|---|
| 1 | Medium | http/mod.rs:438 |
sql_query_context.user is None when no Authorization header is present, so the cap is inert on deployments without credentials. The docs entry still says "a single authenticated user" without stating that unauthenticated connections are exempt. Existing thread · Fix this → |
| 2 | Low | http/mod.rs:302-304 |
The Drop timeout comment still asserts the stranded entry "is the oldest one of that user by construction" — connections don't finish in admission order, so a mid-list entry can be stranded. Comment-only; the code is fine. Existing thread |
| 3 | Low | http/mod.rs:445 |
tokio::select! still picks branches at random; biased; with the eviction arm first would make "evicted means evicted" deterministic and tighten the overshoot bound EVICTED_CLOSE_TIMEOUT reasons about. |
Smaller notes
- The stray run of spaces in the stuck-lock test's assert message (
"...admits the connection untracked...",http/mod.rs:1635) is still there. count(http/mod.rs:280) still collapses "lock unavailable" into0, and still has no non-test consumer.- No threads became stale in this pass, so nothing was resolved.
- I did not run
cargo check/cargo testfor cubestore here (build time in this environment). A one-line constant change to an existingDurationcan't fail to compile.
Overall: the feature is in good shape. Nothing in this commit changes my view of it, and the two carried findings are both small — one docs sentence and one comment.
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #11667 +/- ##
==========================================
+ Coverage 58.20% 58.38% +0.18%
==========================================
Files 237 237
Lines 18852 18853 +1
Branches 3825 3830 +5
==========================================
+ Hits 10973 11008 +35
+ Misses 7323 7283 -40
- Partials 556 562 +6
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Summary
Cube Store accepts an unbounded number of concurrent WebSocket connections per authenticated user. A client that opens connections faster than it closes them consumes the whole process file descriptor table, and once that is full the node cannot accept connections for any other user on it either. This adds an optional per-user cap.
At the limit the user's oldest connection is closed to admit the new one, rather than the new one being refused. A client that still needs the closed connection reconnects; one that had forgotten about it simply loses it. Refusing instead would leave a user whose slots are all held by forgotten connections unable to open a working one at all, since nothing would ever release them.
Disabled by default (
0), so behaviour is unchanged unless a limit is set.Changes
CUBESTORE_MAX_WS_CONNECTIONS_PER_USER— new config option, default0(disabled), documented in the environment variable reference.WsConnectionCounter— live connections per authenticated user, keyed by a monotonic id so the first entry is the oldest. Unauthenticated connections and a disabled cap stay untracked; users whose last connection goes away are removed from the map, so a churn of one-off users cannot grow it.WsConnectionGuard— taken at the/wsupgrade and released fromDrop, so the slot comes back on every exit path: a close frame, a transport error, a panic, a cancelled task, or a connection that dies before the upgrade completes.CancellationTokenthe connection task selects on, next to the two branches it already had. The task sends a close frame before it leaves, so the client sees a clean close rather than a reset.drop, so the limit never overshoots; the monotonic id is what keeps those two removals from taking each other's entry.The counter is per process, so a user connecting to several nodes may hold up to the limit on each. Called out in the docs, along with the advice to set the limit comfortably above what a client legitimately keeps open at once — below that, working connections get recycled.
Changed behaviour at the limit
Testing
cargo test -p cubestore --lib http::— 16 passed. That includes the existing WebSocket tests (ws_test,ws_process_id_header_test,query_test,inline_tables_query_test, and the message-size tests), which run through the acquire/release path with the cap disabled, plus five new unit tests:dropdoes not remove another connection's entry;Droprecovers a poisoned lock instead of unwrapping it.cargo check -p cubestore --all-targetsis clean, no new warnings.Design note
The counter lock is a
std::sync::Mutexrather than atokioone behindutil::lock::acquire_lock, because entries are removed from a destructor and a destructor cannot.await. Releasing anywhere else — explicitly at the end of the connection loop, or from a spawned task — would leak the slot whenever the connection ends abnormally, and a cap whose slots only ever leak eventually admits nothing. Nothing awaits inside the critical section, which is what makes the synchronous lock correct here; that reasoning is recorded next to the field.