From 3681c1b608ab03d7ee47fbc9d029453934640ae1 Mon Sep 17 00:00:00 2001 From: Tom Dupuis <60640908+tomdps@users.noreply.github.com> Date: Sat, 25 Jul 2026 04:29:43 +0200 Subject: [PATCH] feat: implement #651 - feat(cluster-protocol): add the Rust WebSocket binding and cloud data-plane contract --- .../src/websocket.rs | 106 +++++++++++++----- .../src/websocket/tests.rs | 105 +++++++++++++++++ .../tests/websocket.rs | 101 +++++++++++++++++ 3 files changed, 286 insertions(+), 26 deletions(-) create mode 100644 crates/openengine-cluster-server/src/websocket/tests.rs diff --git a/crates/openengine-cluster-server/src/websocket.rs b/crates/openengine-cluster-server/src/websocket.rs index 6c7bea80..3298dad6 100644 --- a/crates/openengine-cluster-server/src/websocket.rs +++ b/crates/openengine-cluster-server/src/websocket.rs @@ -4,7 +4,7 @@ //! ([`crate::stdio::ConnectionState`], `classify_ndjson_line`, `dispatch_classified_line`) so //! results, events, and errors stay byte-equivalent between the stdio and WebSocket bindings. //! Framing rules unique to WebSocket -- binary rejection, the 1,048,576 UTF-8 byte bound, and -//! best-effort `$/cancelRequest` -- live only here; `stdio::serve_ndjson` itself is untouched. +//! race-free `$/cancelRequest` -- live only here; `stdio::serve_ndjson` itself is untouched. use std::collections::HashMap; use std::io; @@ -36,9 +36,12 @@ pub const MAX_FRAME_BYTES: usize = 1_048_576; /// Bounded per-connection outbound queue, matching `stdio::serve_ndjson`'s bound. const OUTBOUND_QUEUE_CAPACITY: usize = 256; -/// Per-connection registry of best-effort abort handles for in-flight passthrough (unary) request -/// tasks, keyed by their [`RequestId`]. `$/cancelRequest` looks an id up here and aborts it; an -/// absent id (unknown, or already completed and self-removed) is silently a no-op. Establishing +/// Per-connection registry of abort handles for in-flight passthrough (unary) request tasks, keyed +/// by their [`RequestId`]. Registration happens-before the spawned task can perform any work -- +/// including completing -- via `spawn_passthrough`'s readiness gate, and release is unconditional +/// on every exit path (normal completion, cancellation, or panic) via [`PassthroughGuard`]'s +/// `Drop` impl, so no stale or leaked entry is possible. `$/cancelRequest` looks an id up here and +/// aborts it; an absent id (unknown, or already completed) is silently a no-op. Establishing /// requests (`watch`/`logs`/`agent/attach`) are not registered here -- their in-flight lifetime is /// already covered by `subscription/cancel` once established, and is intentionally short. type CancelRegistry = Arc>>; @@ -53,7 +56,7 @@ pub fn websocket_config() -> WebSocketConfig { /// Serves one already-handshaken WebSocket connection: demultiplexes unary requests and `watch`/ /// `logs`/`agent/attach` subscriptions sharing this connection exactly like `stdio::serve_ndjson`, -/// plus per-connection best-effort `$/cancelRequest`. Binary frames close with code 1003; +/// plus per-connection race-free `$/cancelRequest`. Binary frames close with code 1003; /// oversized or capacity-rejected text frames close with code 1009. Never returns an `Err`: /// transport failures close the connection and this simply returns once torn down. pub async fn serve_websocket( @@ -144,7 +147,7 @@ enum FrameOutcome { } /// Per-frame handling context: the shared [`DispatchCtx`] plus the two handles unique to this -/// WebSocket binding (the best-effort `$/cancelRequest` registry and the deterministic-close +/// WebSocket binding (the race-free `$/cancelRequest` registry and the deterministic-close /// signal), bundled so [`handle_message`] and its helpers take one argument instead of an /// ever-growing list. struct WsCtx<'a, B> { @@ -202,8 +205,22 @@ where FrameOutcome::Continue } -/// Spawns the passthrough (non-subscription) request task for an admission-approved line, best- -/// effort registering its abort handle under `id` in `cancel_registry` for `$/cancelRequest`. +/// Spawns the passthrough (non-subscription) request task for an admission-approved line, +/// registering its abort handle under `id` in `cancel_registry` for `$/cancelRequest`. The spawned +/// task cannot perform any work -- including completing -- until the caller has finished inserting +/// that handle: a request that carries an `id` gates its own start on a one-shot readiness signal +/// the caller only fires after the insert, making registration happen-before any task work +/// deterministically regardless of which worker thread the scheduler picks it up on. A bare +/// notification (`id.is_none()`) has nothing to register and keeps starting immediately. +/// +/// [`PassthroughGuard`] is constructed here -- synchronously, on the connection task -- and moved +/// into the spawned future's captured environment rather than built inside the task's own body. +/// This matters because a task can be aborted before the runtime ever polls it even once (e.g. a +/// `$/cancelRequest` that lands while still parked behind the readiness gate above): such a task +/// never executes a single statement of its own body, so a guard created *inside* that body would +/// never come into existence and cleanup would silently never run. A value captured by an `async +/// move` block, by contrast, is part of the future's state from the moment the block is +/// constructed, so Rust drops it normally when the never-polled future itself is dropped. fn spawn_passthrough( ctx: &mut WsCtx<'_, B>, id: Option, @@ -215,23 +232,43 @@ fn spawn_passthrough( let task_dispatcher = ctx.dispatch.dispatcher.clone(); let task_state = ctx.dispatch.state.clone(); let task_cancel_registry = Arc::clone(ctx.cancel_registry); - let spawn_id = id.clone(); + let guard = id.clone().map(|id| PassthroughGuard { + id, + state: task_state.clone(), + cancel_registry: Arc::clone(&task_cancel_registry), + }); + let (ready_tx, ready_rx) = if id.is_some() { + let (tx, rx) = oneshot::channel::<()>(); + (Some(tx), Some(rx)) + } else { + (None, None) + }; let abort_handle = ctx.dispatch.tasks.spawn(async move { + if let Some(ready_rx) = ready_rx { + let _ = ready_rx.await; + } let _permit = permit; run_passthrough_request(PassthroughRequest { dispatcher: task_dispatcher, - id, line, state: task_state, - cancel_registry: task_cancel_registry, + guard, }) .await; }); - // Best-effort: a pathologically fast dispatch can self-remove (see `run_passthrough_request`) - // before this insert runs, leaving a harmless stale entry -- `$/cancelRequest` never claims a - // rollback guarantee, and the entry is bounded by this connection's lifetime and task-slot cap. - if let Some(spawn_id) = spawn_id { - ctx.cancel_registry.lock().insert(spawn_id, abort_handle); + // Test-only: widens the window between the spawned task being handed to the scheduler and + // this registration completing, so `tests::fast_completion_race_cannot_leave_stale_registry_entry` + // can deterministically manufacture -- rather than hope to get lucky on -- the exact adverse + // interleaving the readiness gate above exists to close. `cfg(test)` only activates while + // compiling this crate's own lib unit tests, so this never runs in a real binary or in + // `tests/websocket.rs`'s integration coverage (a separate crate that links this one normally). + #[cfg(test)] + std::thread::sleep(std::time::Duration::from_millis(2)); + if let Some(id) = id { + ctx.cancel_registry.lock().insert(id, abort_handle); + } + if let Some(ready_tx) = ready_tx { + let _ = ready_tx.send(()); } } @@ -239,32 +276,46 @@ fn spawn_passthrough( /// reasonable. struct PassthroughRequest { dispatcher: Dispatcher, - id: Option, line: String, state: ConnectionState, + guard: Option, +} + +/// Releases a passthrough request's `in_flight_ids` and `cancel_registry` entries exactly once, on +/// every exit path -- including `$/cancelRequest`'s task abortion, which drops this future in +/// place (whether at its current await point, or -- if aborted before ever being polled -- as part +/// of dropping the future's captured-but-never-executed initial state) instead of resuming past +/// it, so cleanup cannot depend on falling off the end of `run_passthrough_request`. +struct PassthroughGuard { + id: RequestId, + state: ConnectionState, cancel_registry: CancelRegistry, } +impl Drop for PassthroughGuard { + fn drop(&mut self) { + self.state.in_flight_ids.lock().remove(&self.id); + self.cancel_registry.lock().remove(&self.id); + } +} + /// Dispatches a non-subscription request or notification frame, releasing its in-flight id and -/// best-effort cancel registration (if any) once the backend call returns and before the response -/// is enqueued. Mirrors `stdio::run_passthrough_request`, plus the `cancel_registry` cleanup that -/// binding has no notion of. +/// cancel registration (if any, via the already-constructed [`PassthroughGuard`]) once the backend +/// call returns and before the response is enqueued -- unconditionally, even if the task is +/// aborted mid-flight. Mirrors `stdio::run_passthrough_request`, plus the `cancel_registry` cleanup +/// that binding has no notion of. async fn run_passthrough_request(request: PassthroughRequest) where B: ClusterBackend, { let PassthroughRequest { dispatcher, - id, line, state, - cancel_registry, + guard, } = request; let response = dispatcher.dispatch(&line).await; - if let Some(id) = &id { - state.in_flight_ids.lock().remove(id); - cancel_registry.lock().remove(id); - } + drop(guard); let _ = state.outbound_tx.send(response).await; } @@ -331,3 +382,6 @@ async fn run_writer( } let _ = sink.close().await; } + +#[cfg(test)] +mod tests; diff --git a/crates/openengine-cluster-server/src/websocket/tests.rs b/crates/openengine-cluster-server/src/websocket/tests.rs new file mode 100644 index 00000000..b91348bb --- /dev/null +++ b/crates/openengine-cluster-server/src/websocket/tests.rs @@ -0,0 +1,105 @@ +use std::time::Duration; + +use openengine_cluster_protocol::RunId; +use serde_json::{json, Value}; + +use super::*; +use crate::stdio::{ + classify_ndjson_line, dispatch_classified_line, new_connection_setup, ConnectionSetup, + DispatchCtx, LineDispatch, +}; +use crate::watch::fixtures::{FixtureBackend, FixtureStore}; +use crate::ConnectionContext; + +/// Regression test for the spawn-before-register race in `spawn_passthrough`: pre-fix, the newly +/// spawned passthrough task starts running immediately once handed to the scheduler and, once its +/// (near-instant, in-memory) dispatch resolves, removes its own `cancel_registry` entry as part of +/// normal cleanup -- a no-op if that happens before the caller's `insert` for the same id has run. +/// Nothing in the pre-fix code establishes a happens-before relationship between "the task is +/// schedulable" and "the caller's insert completes", so if the removal is observed first, the +/// caller's later insert leaves a stale handle for an already-finished request that nothing will +/// ever clean up. +/// +/// Real scheduling alone essentially never reproduces this: the caller attempts its insert within +/// nanoseconds of spawning (no other work in between), while the spawned task needs a genuine +/// cross-thread scheduling round-trip before it can even begin dispatch, so the safe +/// insert-then-remove ordering dominates overwhelmingly by default (independently confirmed by +/// running an all-real-timing version of this test hundreds of times against the pre-fix code +/// without once observing the race). To make the test actually exercise the interleaving the issue +/// requires proving fixed, `spawn_passthrough` widens that window under `#[cfg(test)]` only (a +/// no-op in every real binary and in `tests/websocket.rs`'s separate-crate integration coverage), +/// giving the spawned task's cleanup a deterministic chance to run first. +/// +/// Post-fix, `spawn_passthrough`'s readiness gate means the spawned task cannot attempt *any* work +/// -- including reaching its own cleanup -- until a readiness signal fires, which is sent only +/// after the caller's insert has already completed. Widening the same window therefore cannot move +/// the outcome post-fix: the spawned task is still parked behind the gate regardless of how long +/// the caller takes to reach its insert, so registration remains happens-before any request work by +/// construction. That asymmetry -- reliably broken pre-fix, structurally unbreakable post-fix -- is +/// exactly what this test proves. +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +async fn fast_completion_race_cannot_leave_stale_registry_entry() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let dispatcher = Dispatcher::new(FixtureBackend::new(store), ConnectionContext::default()); + let (outbound_tx, mut outbound_rx) = mpsc::channel::(256); + let ConnectionSetup { + task_slots, + mut tasks, + state, + .. + } = new_connection_setup(&outbound_tx); + let cancel_registry: CancelRegistry = Arc::new(Mutex::new(HashMap::new())); + + const TRIALS: i64 = 20; + for trial in 0..TRIALS { + let id = RequestId::Integer(trial); + let line = + json!({"jsonrpc": "2.0", "id": trial, "method": "get", "params": {}}).to_string(); + + let kind = classify_ndjson_line(&line); + let mut dispatch_ctx = DispatchCtx { + dispatcher: &dispatcher, + state: &state, + task_slots: &task_slots, + tasks: &mut tasks, + }; + let LineDispatch::Passthrough { id: got_id, permit } = + dispatch_classified_line(kind, &mut dispatch_ctx).await + else { + panic!("expected a passthrough dispatch for a `get` request"); + }; + assert_eq!(got_id.as_ref(), Some(&id)); + + let mut close_tx: Option> = None; + let mut ctx = WsCtx { + dispatch: DispatchCtx { + dispatcher: &dispatcher, + state: &state, + task_slots: &task_slots, + tasks: &mut tasks, + }, + cancel_registry: &cancel_registry, + close_tx: &mut close_tx, + }; + spawn_passthrough(&mut ctx, got_id, permit, line); + + // Drain this trial's task to completion (its post-dispatch cleanup, whichever way the + // race landed, has therefore already run) before inspecting registry state. + while tasks.join_next().await.is_some() {} + + let response = tokio::time::timeout(Duration::from_secs(1), outbound_rx.recv()) + .await + .expect("the request must receive its response promptly") + .expect("outbound channel must not close mid-test"); + let response: Value = serde_json::from_str(&response).unwrap(); + assert_eq!(response["id"], trial); + assert!(response.get("result").is_some(), "{response}"); + + assert!( + cancel_registry.lock().is_empty(), + "trial {trial}: a stale cancel_registry entry survived a fast completion racing the \ + caller's registration -- insert is not guaranteed to happen-before the spawned task \ + can observe/complete its own cleanup" + ); + } +} diff --git a/crates/openengine-cluster-server/tests/websocket.rs b/crates/openengine-cluster-server/tests/websocket.rs index 173c19e6..b1ebed9a 100644 --- a/crates/openengine-cluster-server/tests/websocket.rs +++ b/crates/openengine-cluster-server/tests/websocket.rs @@ -216,6 +216,107 @@ impl GatedHarnessSpawn for Harness { } } +#[tokio::test] +async fn cancel_request_releases_in_flight_id_and_permits_id_reuse() { + let (mut harness, gate) = spawn_gated_harness::().await; + + // Blocks on the gate: this request never completes, so its `in_flight_ids`/`cancel_registry` + // entries can only be released via the cancellation path, never via normal fall-through. + harness.send_get(1).await; + + let cancel = json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": {"id": 1}, + }) + .to_string(); + send_text(&mut harness.client, cancel).await; + + // `AbortHandle::abort()` only schedules the target task's cancellation; the guard that + // releases `in_flight_ids` runs whenever the runtime actually gets around to polling (and + // dropping) that task, not synchronously when `abort()` returns. Poll for the id becoming + // reusable via bounded cooperative retries rather than a fixed sleep: a duplicate rejection is + // a harmless, synchronous no-op (it never touches `in_flight_ids`, so resending is safe), and + // it always arrives promptly since it never waits on the still-closed gate -- whereas an + // *accepted* retry's task immediately blocks on that same still-closed gate exactly like the + // original did, so it produces no prompt response at all. That absence, not its content, + // is what distinguishes "admitted" from "still rejected" here, deterministically regardless + // of scheduler load. The attempt cap still fails the test if the id is never released -- the + // exact pre-fix bug (a cancelled request permanently reserving its id). + let mut accepted = false; + for attempt in 0..100_000 { + harness.send_get(1).await; + match tokio::time::timeout(Duration::from_millis(50), harness.recv_value()).await { + Ok(response) => { + assert_eq!(response["id"], 1); + assert_eq!( + response["error"]["data"]["code"], "DUPLICATE_REQUEST_ID", + "a prompt response before the gate is released can only be a duplicate \ + rejection (a real result would block on the still-closed gate); got \ + {response}" + ); + assert!( + attempt < 99_999, + "id 1 was never released after cancellation settled -- it is still \ + permanently reserved in in_flight_ids" + ); + tokio::task::yield_now().await; + } + Err(_) => { + // No prompt response: this retry was admitted and its task is now parked on the + // gate, exactly like the original was before cancellation. + accepted = true; + break; + } + } + } + assert!(accepted, "id 1 was never accepted for reuse"); + + // Releasing the gate now only unblocks the accepted retry's still-pending backend call; no + // response frame was ever emitted for the original cancelled request, so this is the only + // frame left to receive. + gate.notify_one(); + let second = tokio::time::timeout(Duration::from_secs(1), harness.recv_value()) + .await + .expect("the accepted retry must complete once the gate is released"); + assert_eq!(second["id"], 1); + assert!(second.get("result").is_some(), "{second}"); + + shut_down(harness).await; +} + +#[tokio::test] +async fn fast_completions_do_not_leak_or_corrupt_cancel_registry() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(store)).await; + + // Ungated backend: each `get` resolves as fast as the runtime allows. Repeatedly reusing the + // same id and firing a `$/cancelRequest` immediately after each completion -- when the id is + // already fully released and the notification can only ever hit an unknown-or-completed id -- + // is a high-volume regression/sanity net for that exact "silent no-op" contract and for the + // general fast-completion-plus-cancellation interaction: every iteration must still receive + // its own correct, uncorrupted response and the connection must never misbehave (unexpected + // close, mismatched response, panic, deadlock) across 200 rapid cycles. + for _ in 0..200 { + send_text(&mut harness.client, request_text(1, "get", json!({}))).await; + let response = tokio::time::timeout(Duration::from_secs(1), recv_json(&mut harness.client)) + .await + .expect("each fast completion must receive its own uncorrupted response"); + assert_eq!(response["id"], 1); + assert!(response.get("result").is_some(), "{response}"); + + let cancel = json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": {"id": 1}, + }) + .to_string(); + send_text(&mut harness.client, cancel).await; + } + + shut_down(harness).await; +} + #[tokio::test] async fn duplicate_in_flight_request_ids_are_rejected() { let (mut harness, gate) = spawn_gated_harness::().await;