From 8f3841d01057b39ba31976279d7e21f9dfae4b5f Mon Sep 17 00:00:00 2001 From: Tom Dupuis <60640908+tomdps@users.noreply.github.com> Date: Sat, 25 Jul 2026 02:29:13 +0200 Subject: [PATCH] feat: implement #651 - feat(cluster-protocol): add the Rust WebSocket binding and cloud data-plane contract --- Cargo.lock | 141 +++++++- Cargo.toml | 2 + crates/openengine-cluster-client/Cargo.toml | 2 + crates/openengine-cluster-client/src/lib.rs | 276 +++++---------- .../src/multiplex.rs | 264 ++++++++++++++ .../src/ndjson_agent_attach.rs | 18 +- .../src/ndjson_logs.rs | 16 +- .../src/ndjson_pump.rs | 61 +++- .../src/ndjson_subscription.rs | 46 ++- .../src/ndjson_watch.rs | 57 +-- .../src/websocket.rs | 116 ++++++ .../tests/backend_faults.rs | 3 + .../reconnect_support/cancel_scenario.rs | 74 ++++ .../tests/reconnect_support/mod.rs | 22 +- .../reconnect_support/ndjson_scenario.rs | 51 +-- .../reconnect_support/scenario_harness.rs | 51 +++ .../reconnect_support/websocket_scenario.rs | 43 +++ .../tests/subscription_ndjson.rs | 86 +---- .../tests/websocket.rs | 168 +++++++++ .../openengine-cluster-protocol/src/watch.rs | 14 +- crates/openengine-cluster-server/Cargo.toml | 2 + crates/openengine-cluster-server/src/lib.rs | 1 + crates/openengine-cluster-server/src/stdio.rs | 188 +++------- .../src/stdio/admission.rs | 11 +- .../src/stdio/agent_attach.rs | 3 +- .../src/stdio/dispatch.rs | 195 ++++++++++ .../src/stdio/logs.rs | 5 +- .../src/stdio/tests.rs | 2 + .../src/watch/fixtures.rs | 45 ++- .../src/websocket.rs | 333 ++++++++++++++++++ .../tests/admission_bound_support/mod.rs | 84 +++++ .../tests/gated_backend_support/mod.rs | 50 +++ .../tests/subscription_ndjson.rs | 116 ++---- .../tests/websocket.rs | 235 ++++++++++++ crates/openengine-cluster-testkit/Cargo.toml | 2 + .../src/artifacts.rs | 15 +- .../src/artifacts/openrpc.rs | 7 +- .../tests/protocol_ndjson.rs | 126 +------ .../tests/protocol_transcript_support/mod.rs | 136 +++++++ .../tests/protocol_websocket.rs | 197 +++++++++++ .../v1/data-plane.md | 75 ++++ docs/openengine-cluster-protocol/v1/watch.md | 6 +- protocol/openengine-cluster/v1/openrpc.json | 5 +- protocol/openengine-cluster/v1/schema.json | 50 ++- 44 files changed, 2692 insertions(+), 708 deletions(-) create mode 100644 crates/openengine-cluster-client/src/multiplex.rs create mode 100644 crates/openengine-cluster-client/src/websocket.rs create mode 100644 crates/openengine-cluster-client/tests/reconnect_support/cancel_scenario.rs create mode 100644 crates/openengine-cluster-client/tests/reconnect_support/scenario_harness.rs create mode 100644 crates/openengine-cluster-client/tests/reconnect_support/websocket_scenario.rs create mode 100644 crates/openengine-cluster-client/tests/websocket.rs create mode 100644 crates/openengine-cluster-server/src/stdio/dispatch.rs create mode 100644 crates/openengine-cluster-server/src/websocket.rs create mode 100644 crates/openengine-cluster-server/tests/admission_bound_support/mod.rs create mode 100644 crates/openengine-cluster-server/tests/gated_backend_support/mod.rs create mode 100644 crates/openengine-cluster-server/tests/websocket.rs create mode 100644 crates/openengine-cluster-testkit/tests/protocol_transcript_support/mod.rs create mode 100644 crates/openengine-cluster-testkit/tests/protocol_websocket.rs create mode 100644 docs/openengine-cluster-protocol/v1/data-plane.md diff --git a/Cargo.lock b/Cargo.lock index 585eb4f1..bc72c715 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -218,6 +218,12 @@ dependencies = [ "typenum", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "digest" version = "0.10.7" @@ -348,6 +354,17 @@ version = "0.3.32" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "futures-sink" version = "0.3.33" @@ -367,6 +384,8 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" dependencies = [ "futures-core", + "futures-macro", + "futures-sink", "futures-task", "pin-project-lite", "slab", @@ -438,6 +457,22 @@ dependencies = [ "hashbrown 0.15.5", ] +[[package]] +name = "http" +version = "1.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6970f50e31d6fc17d3fa27329444bfa74e196cf62e95052a3f6fee181dba6425" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + [[package]] name = "icu_collections" version = "2.1.1" @@ -629,6 +664,12 @@ dependencies = [ "scopeguard", ] +[[package]] +name = "log" +version = "0.4.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" + [[package]] name = "memchr" version = "2.8.3" @@ -735,6 +776,7 @@ name = "openengine-cluster-client" version = "0.1.0" dependencies = [ "async-trait", + "futures-util", "openengine-cluster-protocol", "openengine-cluster-server", "parking_lot", @@ -743,6 +785,7 @@ dependencies = [ "thiserror", "tokio", "tokio-stream", + "tokio-tungstenite", "tokio-util", ] @@ -763,6 +806,7 @@ name = "openengine-cluster-server" version = "0.1.0" dependencies = [ "async-trait", + "futures-util", "openengine-cluster-protocol", "parking_lot", "serde", @@ -770,6 +814,7 @@ dependencies = [ "thiserror", "tokio", "tokio-stream", + "tokio-tungstenite", "tokio-util", ] @@ -778,6 +823,7 @@ name = "openengine-cluster-testkit" version = "0.1.0" dependencies = [ "async-trait", + "futures-util", "jsonschema", "openengine-cluster-client", "openengine-cluster-protocol", @@ -788,6 +834,7 @@ dependencies = [ "sha2", "thiserror", "tokio", + "tokio-tungstenite", ] [[package]] @@ -921,8 +968,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "22f6172bdec972074665ed81ed53b71da00bfc44b65a753cfde883ec4c702a1a" dependencies = [ "libc", - "rand_chacha", - "rand_core", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", ] [[package]] @@ -932,7 +989,17 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", ] [[package]] @@ -944,6 +1011,15 @@ dependencies = [ "getrandom 0.2.17", ] +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -1066,7 +1142,7 @@ dependencies = [ "borsh", "bytes", "num-traits", - "rand", + "rand 0.8.7", "rkyv", "serde", "serde_json", @@ -1170,6 +1246,17 @@ dependencies = [ "zmij", ] +[[package]] +name = "sha1" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a978451301f4db1d02937a4ab3ccce137717b81826e79b7d49ffe3244a13c3b8" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "sha2" version = "0.10.9" @@ -1215,6 +1302,16 @@ version = "1.15.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" +[[package]] +name = "socket2" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d1e2c7f27f8d4cb10542a02c49005dbd6e93095799d6f3be745fae9f8fedd4" +dependencies = [ + "libc", + "windows-sys", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -1316,6 +1413,7 @@ dependencies = [ "mio", "pin-project-lite", "signal-hook-registry", + "socket2", "tokio-macros", "windows-sys", ] @@ -1342,6 +1440,18 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-tungstenite" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a9daff607c6d2bf6c16fd681ccb7eecc83e4e2cdc1ca067ffaadfca5de7f084" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.19" @@ -1386,6 +1496,23 @@ dependencies = [ "winnow", ] +[[package]] +name = "tungstenite" +version = "0.26.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4793cb5e56680ecbb1d843515b23b6de9a75eb04b66643e256a396d43be33c13" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.5", + "sha1", + "thiserror", + "utf-8", +] + [[package]] name = "typenum" version = "1.20.1" @@ -1398,6 +1525,12 @@ version = "1.0.24" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/Cargo.toml b/Cargo.toml index 3d8ef392..0052680b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -28,6 +28,8 @@ thiserror = "2.0.17" tokio = { version = "1.48.0", features = ["fs", "io-std", "io-util", "macros", "process", "rt-multi-thread", "sync", "time"] } tokio-util = { version = "0.7", features = ["codec"] } tokio-stream = "0.1" +tokio-tungstenite = "0.26" +futures-util = "0.3" jsonschema = { version = "0.29.1", default-features = false } rusqlite = { version = "0.37.0", features = ["bundled"] } diff --git a/crates/openengine-cluster-client/Cargo.toml b/crates/openengine-cluster-client/Cargo.toml index 5b3c4f49..3d01c815 100644 --- a/crates/openengine-cluster-client/Cargo.toml +++ b/crates/openengine-cluster-client/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] async-trait.workspace = true +futures-util.workspace = true openengine-cluster-protocol.workspace = true openengine-cluster-server.workspace = true parking_lot.workspace = true @@ -16,4 +17,5 @@ serde_json.workspace = true thiserror.workspace = true tokio.workspace = true tokio-stream.workspace = true +tokio-tungstenite.workspace = true tokio-util.workspace = true diff --git a/crates/openengine-cluster-client/src/lib.rs b/crates/openengine-cluster-client/src/lib.rs index c4f9e565..c7482de2 100644 --- a/crates/openengine-cluster-client/src/lib.rs +++ b/crates/openengine-cluster-client/src/lib.rs @@ -1,35 +1,36 @@ //! Typed transport-neutral Cluster Protocol client. +mod multiplex; mod ndjson_pump; mod ndjson_subscription; -use ndjson_pump::forward_notification; - pub mod agent_attach; pub mod logs; pub mod ndjson_agent_attach; pub mod ndjson_logs; pub mod ndjson_watch; pub mod watch; +pub mod websocket; pub use agent_attach::*; pub use logs::*; pub use ndjson_agent_attach::*; pub use ndjson_logs::*; pub use ndjson_watch::*; pub use watch::*; +pub use websocket::*; -use std::collections::{hash_map::Entry, HashMap}; -use std::sync::atomic::{AtomicBool, AtomicI64, AtomicU64, Ordering}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; use std::sync::Arc; use parking_lot::Mutex as ParkingMutex; use async_trait::async_trait; use openengine_cluster_protocol::{ ApplyParams, ApplyResult, DeleteParams, DeleteResult, GetParams, GetResult, InitializeParams, - InitializeResult, JsonRpcError, JsonRpcErrorResponse, JsonRpcNotification, JsonRpcRequest, - JsonRpcSuccess, PlanParams, PlanResult, RequestId, ResubmitParams, ResubmitResult, RetryParams, - RetryResult, StopParams, StopResult, SubscriptionCancelParams, SubscriptionId, UpdateParams, - UpdateResult, JSON_RPC_VERSION, PROTOCOL_VERSION, + InitializeResult, JsonRpcError, JsonRpcErrorResponse, JsonRpcRequest, JsonRpcSuccess, + PlanParams, PlanResult, RequestId, ResubmitParams, ResubmitResult, RetryParams, RetryResult, + StopParams, StopResult, SubscriptionId, UpdateParams, UpdateResult, JSON_RPC_VERSION, + PROTOCOL_VERSION, }; use openengine_cluster_server::{ClusterBackend, Dispatcher}; use serde::de::DeserializeOwned; @@ -38,7 +39,6 @@ use serde_json::Value; use thiserror::Error; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; use tokio::sync::{mpsc, oneshot, Mutex}; -use tokio::task::JoinHandle; use tokio_stream::StreamExt; use tokio_util::codec::{FramedRead, LinesCodec}; @@ -68,6 +68,39 @@ where } } +/// Transport-neutral generic subscription framing: establishing a `watch`/`logs`/`agent/attach` +/// request, cancelling an established subscription, and best-effort cancelling any in-flight +/// unary request by id. Implemented once per wire transport ([`NdjsonTransport`], +/// [`crate::websocket::WebSocketTransport`]) so [`crate::watch::ReconnectingEventStream`]-shaped +/// subscription clients (`WatchSubscriptionClient` and the `logs`/`agent_attach` counterparts +/// generated by [`ndjson_subscription::impl_ndjson_event_subscription`]) drive either transport +/// through the exact same generic code. +#[async_trait] +pub trait SubscriptionTransport: JsonRpcTransport { + /// Sends a subscription-establishing request and returns its response line plus, only on + /// success, the receiver registered for its subscription's notifications. + async fn open_subscription( + &self, + request: String, + id: RequestId, + ) -> Result<(String, Option), TransportError>; + + /// Sends a `subscription/cancel` notification. Fire-and-forget: cancellation has no response + /// on the wire, so this only reports a write failure. + async fn cancel_subscription(&self, id: SubscriptionId) -> Result<(), TransportError>; + + /// Sends a `$/cancelRequest` notification best-effort cancelling an in-flight unary request by + /// id. Fire-and-forget, exactly like [`Self::cancel_subscription`]; the server silently + /// no-ops an unknown or already-completed id and never claims a rollback after commit. + async fn cancel_request(&self, id: RequestId) -> Result<(), TransportError>; + + /// Mints the next request id for a subscription-establishing request, from a counter shared by + /// this transport's connection rather than a client-local one, so independently constructed + /// subscription clients on one connection cannot replace each other's pending response + /// waiters. + fn next_watch_request_id(&self) -> RequestId; +} + pub struct InProcessTransport { dispatcher: Dispatcher, } @@ -108,7 +141,11 @@ struct PumpedResponse { subscription: Option, } -pub(crate) struct PumpedSubscription { +/// Opaque per-subscription handle handed from [`SubscriptionTransport::open_subscription`] to a +/// subscription client: the receiver end of that subscription's forwarded notification lines, +/// plus the shared overflow flag its pump sets on a full/abandoned local queue. Its fields are +/// crate-private; only [`SubscriptionTransport`] implementors and subscription clients wire them. +pub struct PumpedSubscription { pub(crate) receiver: mpsc::Receiver, pub(crate) overflowed: Arc, } @@ -122,15 +159,34 @@ struct SubscriptionRegistration { type PendingMap = Arc>>>; type SubscriptionMap = Arc>>; +/// Writes one line (a trailing `\n`, flushed) to a shared writer -- the [`multiplex::FrameSink`] +/// implementation backing [`NdjsonTransport`]. +struct NdjsonFrameSink { + writer: Arc>, +} + +#[async_trait] +impl multiplex::FrameSink for NdjsonFrameSink +where + W: AsyncWrite + Send + Unpin, +{ + async fn send_frame(&self, frame: String) -> Result<(), TransportError> { + let mut writer = self.writer.lock().await; + writer.write_all(frame.as_bytes()).await?; + writer.write_all(b"\n").await?; + writer.flush().await?; + Ok(()) + } +} + /// NDJSON stdio transport that demultiplexes unary request/response traffic and generic `watch` /// subscription notifications sharing one connection, correlating by request id and subscription -/// id respectively. A background pump task owns the read half; [`Self::request`] and the -/// subscription-establishing/cancelling methods below only ever take the write half's lock. +/// id respectively. Holds one [`multiplex::MultiplexedTransport`], which owns the demux state +/// (write sink, pending-request map, pump task, watch-id counter) and implements every +/// [`JsonRpcTransport`]/[`SubscriptionTransport`] method against it -- see +/// [`crate::websocket::WebSocketTransport`] for the identical WebSocket-frame binding. pub struct NdjsonTransport { - writer: Arc>, - pending: PendingMap, - pump: JoinHandle<()>, - next_watch_id: AtomicU64, + inner: multiplex::MultiplexedTransport>, _reader: std::marker::PhantomData R>, } @@ -143,112 +199,29 @@ where pub fn new(reader: R, writer: W) -> Self { let pending: PendingMap = Arc::new(ParkingMutex::new(HashMap::new())); let subscriptions: SubscriptionMap = Arc::new(ParkingMutex::new(HashMap::new())); - let writer = Arc::new(Mutex::new(writer)); + let sink = NdjsonFrameSink { + writer: Arc::new(Mutex::new(writer)), + }; let pump = tokio::spawn(run_pump( reader, Arc::clone(&pending), subscriptions, - Arc::clone(&writer), + NdjsonFrameSink { + writer: Arc::clone(&sink.writer), + }, )); Self { - writer, - pending, - pump, - next_watch_id: AtomicU64::new(1), + inner: multiplex::MultiplexedTransport::new(sink, pending, pump), _reader: std::marker::PhantomData, } } - - async fn write_line(&self, line: &str) -> Result<(), TransportError> { - write_shared_line(&self.writer, line).await - } - - /// Registers `id` as pending, writes `request`, and awaits its demultiplexed response. - async fn send_request( - &self, - request: String, - id: RequestId, - ) -> Result { - let (sender, receiver) = oneshot::channel(); - { - let mut pending = self.pending.lock(); - match pending.entry(id.clone()) { - Entry::Vacant(entry) => { - entry.insert(sender); - } - Entry::Occupied(_) => { - return Err(TransportError::Protocol(format!( - "request id is already pending: {id:?}" - ))); - } - } - } - if let Err(error) = self.write_line(&request).await { - self.pending.lock().remove(&id); - return Err(error); - } - receiver.await.map_err(|_| { - TransportError::Protocol("server closed the connection before responding".to_owned()) - }) - } - - /// Sends a subscription-establishing request and returns its response line plus, only on - /// success, the receiver registered for its subscription's notifications. Establishment can - /// legitimately fail with a JSON-RPC error (for example `agent/attach` rejecting an unknown or - /// inactive `ExecutionRef`) -- that case carries no `subscriptionId` but is not a transport - /// fault, so it is left for the caller's response parser to surface as a typed - /// [`crate::ClientError::Rpc`] rather than being collapsed into a generic [`TransportError`] - /// here. - pub(crate) async fn open_subscription( - &self, - request: String, - id: RequestId, - ) -> Result<(String, Option), TransportError> { - let response = self.send_request(request, id).await?; - Ok((response.line, response.subscription)) - } - - pub(crate) fn next_watch_request_id(&self) -> RequestId { - RequestId::String(format!( - "watch-{}", - self.next_watch_id.fetch_add(1, Ordering::Relaxed) - )) - } - - /// Sends a `subscription/cancel` notification. Fire-and-forget: cancellation has no response - /// on the wire, so this only reports a write failure. - pub async fn cancel_subscription( - &self, - subscription_id: SubscriptionId, - ) -> Result<(), TransportError> { - let notification = serde_json::to_string(&JsonRpcNotification { - jsonrpc: JSON_RPC_VERSION.to_owned(), - method: "subscription/cancel".to_owned(), - params: SubscriptionCancelParams { subscription_id }, - }) - .expect("subscription cancel notification serialization must succeed"); - self.write_line(¬ification).await - } } -impl Drop for NdjsonTransport { - fn drop(&mut self) { - self.pump.abort(); - } -} - -#[async_trait] -impl JsonRpcTransport for NdjsonTransport -where +multiplex::impl_multiplexed_transport!( + NdjsonTransport where R: AsyncRead + Send + Unpin + 'static, - W: AsyncWrite + Send + Unpin + 'static, -{ - async fn request(&self, request: String) -> Result { - let id = extract_request_id(&request); - let response = self.send_request(request, id).await?; - Ok(response.line) - } -} + W: AsyncWrite + Send + Unpin + 'static +); /// Extracts the `id` from an outgoing request this crate serialized itself. Panics on a malformed /// id, since that indicates an internal bug in request construction rather than bad external @@ -265,90 +238,25 @@ fn extract_request_id(request: &str) -> RequestId { } } -/// Drives the read half: decodes bounded NDJSON lines and demultiplexes each one by request id -/// (unary responses, resolving the matching pending oneshot) or by `params.subscriptionId` -/// (subscription notifications, forwarded to the matching registered channel). A `watch` -/// response carrying a `result.subscriptionId` registers that subscription's channel before -/// resolving the pending oneshot, so no `event` racing the response can be missed. On stream end -/// every pending request fails and every open subscription ends (dropping its sender). +/// Drives the read half: decodes bounded NDJSON lines and routes each one via +/// [`multiplex::route_and_maybe_cancel`] -- shared verbatim with +/// [`crate::websocket::WebSocketTransport`]'s pump, which routes the exact same decoded JSON +/// bodies sourced from `Message::Text` frames instead of NDJSON lines. On stream end every pending +/// request fails and every open subscription ends (dropping its sender). async fn run_pump( reader: R, pending: PendingMap, subscriptions: SubscriptionMap, - writer: Arc>, + sink: NdjsonFrameSink, ) where R: AsyncRead + Unpin, - W: AsyncWrite + Unpin, + W: AsyncWrite + Unpin + Send, { let mut lines = FramedRead::new(reader, LinesCodec::new_with_max_length(MAX_FRAME_BYTES)); while let Some(Ok(line)) = lines.next().await { - let Ok(value) = serde_json::from_str::(&line) else { - continue; - }; - if value.get("method").is_some() { - if let Some(subscription_id) = forward_notification(&value, line, &subscriptions) { - let _ = write_subscription_cancel(&writer, subscription_id).await; - } - continue; - } - let Some(id) = value.get("id").and_then(RequestId::from_json_value) else { - continue; - }; - let Some(sender) = pending.lock().remove(&id) else { - continue; - }; - let subscription = value - .get("result") - .and_then(|result| result.get("subscriptionId")) - .and_then(Value::as_str) - .map(|subscription_id| { - let (sender, receiver) = mpsc::channel(SUBSCRIPTION_QUEUE_CAPACITY); - let overflowed = Arc::new(AtomicBool::new(false)); - subscriptions.lock().insert( - SubscriptionId::new(subscription_id), - SubscriptionRegistration { - sender, - overflowed: Arc::clone(&overflowed), - }, - ); - PumpedSubscription { - receiver, - overflowed, - } - }); - let _ = sender.send(PumpedResponse { line, subscription }); + multiplex::route_and_maybe_cancel(line, &pending, &subscriptions, &sink).await; } - for (_, sender) in pending.lock().drain() { - drop(sender); - } - subscriptions.lock().clear(); -} - -async fn write_subscription_cancel( - writer: &Arc>, - subscription_id: SubscriptionId, -) -> Result<(), TransportError> -where - W: AsyncWrite + Unpin, -{ - let notification = serde_json::to_string(&JsonRpcNotification { - jsonrpc: JSON_RPC_VERSION.to_owned(), - method: "subscription/cancel".to_owned(), - params: SubscriptionCancelParams { subscription_id }, - }) - .expect("subscription cancel notification serialization must succeed"); - write_shared_line(writer, ¬ification).await -} - -async fn write_shared_line(writer: &Arc>, line: &str) -> Result<(), TransportError> -where - W: AsyncWrite + Unpin, -{ - let mut writer = writer.lock().await; - writer.write_all(line.as_bytes()).await?; - writer.write_all(b"\n").await?; - writer.flush().await?; - Ok(()) + multiplex::finish_pump(&pending, &subscriptions); } pub struct ClusterClient { diff --git a/crates/openengine-cluster-client/src/multiplex.rs b/crates/openengine-cluster-client/src/multiplex.rs new file mode 100644 index 00000000..e80604aa --- /dev/null +++ b/crates/openengine-cluster-client/src/multiplex.rs @@ -0,0 +1,264 @@ +//! Shared request/subscription demultiplexing machinery, implemented once and reused verbatim by +//! [`crate::NdjsonTransport`] (NDJSON lines) and [`crate::websocket::WebSocketTransport`] +//! (`Message::Text` frames) so both wire bindings drive this crate's [`crate::JsonRpcTransport`]/ +//! [`crate::SubscriptionTransport`] methods through identical code regardless of frame shape. + +use std::collections::hash_map::Entry; +use std::sync::atomic::{AtomicU64, Ordering}; + +use async_trait::async_trait; +use openengine_cluster_protocol::{ + CancelRequestParams, JsonRpcNotification, RequestId, SubscriptionCancelParams, SubscriptionId, + JSON_RPC_VERSION, +}; +use tokio::sync::oneshot; +use tokio::task::JoinHandle; + +use crate::ndjson_pump::route_pumped_message; +use crate::{ + extract_request_id, JsonRpcTransport, PendingMap, PumpedResponse, PumpedSubscription, + SubscriptionMap, SubscriptionTransport, TransportError, +}; + +/// Abstraction over "write one already-serialized JSON-RPC frame to the peer", implemented once +/// per wire transport ([`crate::NdjsonFrameSink`], [`crate::websocket::WebSocketFrameSink`]) so +/// the demultiplexing logic below is implemented exactly once regardless of the underlying frame +/// shape (NDJSON line vs. `Message::Text`). +#[async_trait] +pub(crate) trait FrameSink: Send + Sync { + async fn send_frame(&self, frame: String) -> Result<(), TransportError>; +} + +/// Registers `id` as pending, writes `request`, and awaits its demultiplexed response. Shared body +/// of `NdjsonTransport::send_request` and `WebSocketTransport::send_request`. +pub(crate) async fn send_request( + sink: &F, + pending: &PendingMap, + request: String, + id: RequestId, +) -> Result { + let (sender, receiver) = oneshot::channel(); + { + let mut pending = pending.lock(); + match pending.entry(id.clone()) { + Entry::Vacant(entry) => { + entry.insert(sender); + } + Entry::Occupied(_) => { + return Err(TransportError::Protocol(format!( + "request id is already pending: {id:?}" + ))); + } + } + } + if let Err(error) = sink.send_frame(request).await { + pending.lock().remove(&id); + return Err(error); + } + receiver.await.map_err(|_| { + TransportError::Protocol("server closed the connection before responding".to_owned()) + }) +} + +/// Shared body of both transports' [`crate::JsonRpcTransport::request`] impl. +pub(crate) async fn request( + sink: &F, + pending: &PendingMap, + request: String, +) -> Result { + let id = extract_request_id(&request); + Ok(send_request(sink, pending, request, id).await?.line) +} + +/// Shared body of both transports' [`crate::SubscriptionTransport::open_subscription`] impl. +pub(crate) async fn open_subscription( + sink: &F, + pending: &PendingMap, + request: String, + id: RequestId, +) -> Result<(String, Option), TransportError> { + let response = send_request(sink, pending, request, id).await?; + Ok((response.line, response.subscription)) +} + +/// Shared body of both transports' [`crate::SubscriptionTransport::cancel_subscription`] impl. +pub(crate) async fn cancel_subscription( + sink: &F, + subscription_id: SubscriptionId, +) -> Result<(), TransportError> { + let notification = serde_json::to_string(&JsonRpcNotification { + jsonrpc: JSON_RPC_VERSION.to_owned(), + method: "subscription/cancel".to_owned(), + params: SubscriptionCancelParams { subscription_id }, + }) + .expect("subscription cancel notification serialization must succeed"); + sink.send_frame(notification).await +} + +/// Shared body of both transports' [`crate::SubscriptionTransport::cancel_request`] impl. +pub(crate) async fn cancel_request( + sink: &F, + id: RequestId, +) -> Result<(), TransportError> { + let notification = serde_json::to_string(&JsonRpcNotification { + jsonrpc: JSON_RPC_VERSION.to_owned(), + method: "$/cancelRequest".to_owned(), + params: CancelRequestParams { id }, + }) + .expect("cancel request notification serialization must succeed"); + sink.send_frame(notification).await +} + +/// Shared body of both transports' [`crate::SubscriptionTransport::next_watch_request_id`] impl. +pub(crate) fn next_watch_id(counter: &AtomicU64) -> RequestId { + RequestId::String(format!("watch-{}", counter.fetch_add(1, Ordering::Relaxed))) +} + +/// Routes one decoded message body via [`route_pumped_message`] and, if it named a subscription +/// whose local queue has overflowed or been abandoned, best-effort sends its cancellation -- +/// shared by both transports' pump loops. +pub(crate) async fn route_and_maybe_cancel( + line: String, + pending: &PendingMap, + subscriptions: &SubscriptionMap, + sink: &F, +) { + if let Some(subscription_id) = route_pumped_message(line, pending, subscriptions) { + let _ = cancel_subscription(sink, subscription_id).await; + } +} + +/// Fails every still-pending request and ends every open subscription (dropping its sender) once +/// a pump's read half ends -- shared tail of both transports' pump loops. +pub(crate) fn finish_pump(pending: &PendingMap, subscriptions: &SubscriptionMap) { + for (_, sender) in pending.lock().drain() { + drop(sender); + } + subscriptions.lock().clear(); +} + +/// Owns one connection's demultiplexing state -- write sink, pending-request map, read-half pump +/// task, and per-connection watch-id counter -- and implements [`crate::JsonRpcTransport`]/ +/// [`crate::SubscriptionTransport`] exactly once against it. [`crate::NdjsonTransport`] and +/// [`crate::websocket::WebSocketTransport`] each hold one of these behind their public, +/// frame-shape-specific type and forward every trait method to it single-line, so the demux +/// wiring -- not just the routing logic in the free functions above -- is implemented once +/// regardless of frame shape. (A blanket `impl JsonRpcTransport for +/// MultiplexedTransport` cannot itself satisfy `NdjsonTransport`/`WebSocketTransport`'s trait +/// bounds without a wrapper: Rust's coherence rules reject a second blanket impl over an +/// unconstrained type parameter alongside the existing `impl +/// JsonRpcTransport for &T` forwarding impl, since `T` could unify with `&_`.) +pub(crate) struct MultiplexedTransport { + sink: F, + pending: PendingMap, + pump: JoinHandle<()>, + next_watch_id: AtomicU64, +} + +impl MultiplexedTransport { + pub(crate) fn new(sink: F, pending: PendingMap, pump: JoinHandle<()>) -> Self { + Self { + sink, + pending, + pump, + next_watch_id: AtomicU64::new(1), + } + } +} + +impl Drop for MultiplexedTransport { + fn drop(&mut self) { + self.pump.abort(); + } +} + +#[async_trait] +impl JsonRpcTransport for MultiplexedTransport { + async fn request(&self, req: String) -> Result { + request(&self.sink, &self.pending, req).await + } +} + +#[async_trait] +impl SubscriptionTransport for MultiplexedTransport { + async fn open_subscription( + &self, + req: String, + id: RequestId, + ) -> Result<(String, Option), TransportError> { + open_subscription(&self.sink, &self.pending, req, id).await + } + + async fn cancel_subscription( + &self, + subscription_id: SubscriptionId, + ) -> Result<(), TransportError> { + cancel_subscription(&self.sink, subscription_id).await + } + + async fn cancel_request(&self, id: RequestId) -> Result<(), TransportError> { + cancel_request(&self.sink, id).await + } + + fn next_watch_request_id(&self) -> RequestId { + next_watch_id(&self.next_watch_id) + } +} + +/// Generates `JsonRpcTransport`/`SubscriptionTransport` for a wire-transport wrapper type holding +/// an `inner: MultiplexedTransport<_>` field, forwarding every method to it. Written once here (as +/// a macro, not a blanket impl -- see [`MultiplexedTransport`]'s doc comment for why a blanket impl +/// does not typecheck) so [`crate::NdjsonTransport`] and [`crate::websocket::WebSocketTransport`] +/// each get one macro invocation instead of hand-writing the same forwarding source twice. +macro_rules! impl_multiplexed_transport { + ($ty:ident < $($generic:ident),+ > where $($bound:tt)+) => { + #[async_trait::async_trait] + impl<$($generic),+> crate::JsonRpcTransport for $ty<$($generic),+> + where + $($bound)+ + { + async fn request(&self, request: String) -> Result { + self.inner.request(request).await + } + } + + #[async_trait::async_trait] + impl<$($generic),+> crate::SubscriptionTransport for $ty<$($generic),+> + where + $($bound)+ + { + /// Establishment can legitimately fail with a JSON-RPC error (for example + /// `agent/attach` rejecting an unknown or inactive `ExecutionRef`) -- that case + /// carries no `subscriptionId` but is not a transport fault, so it is left for the + /// caller's response parser to surface as a typed `ClientError::Rpc` rather than + /// being collapsed into a generic `TransportError` here. + async fn open_subscription( + &self, + request: String, + id: openengine_cluster_protocol::RequestId, + ) -> Result<(String, Option), crate::TransportError> { + self.inner.open_subscription(request, id).await + } + + async fn cancel_subscription( + &self, + subscription_id: openengine_cluster_protocol::SubscriptionId, + ) -> Result<(), crate::TransportError> { + self.inner.cancel_subscription(subscription_id).await + } + + async fn cancel_request( + &self, + id: openengine_cluster_protocol::RequestId, + ) -> Result<(), crate::TransportError> { + self.inner.cancel_request(id).await + } + + fn next_watch_request_id(&self) -> openengine_cluster_protocol::RequestId { + self.inner.next_watch_request_id() + } + } + }; +} + +pub(crate) use impl_multiplexed_transport; diff --git a/crates/openengine-cluster-client/src/ndjson_agent_attach.rs b/crates/openengine-cluster-client/src/ndjson_agent_attach.rs index d893fa08..a94953f9 100644 --- a/crates/openengine-cluster-client/src/ndjson_agent_attach.rs +++ b/crates/openengine-cluster-client/src/ndjson_agent_attach.rs @@ -1,14 +1,18 @@ -//! NDJSON-bound `agent/attach` subscription client. Drives `agent/attach`/`event`/ -//! `subscription/cancel`/`subscription/closed` notifications over [`crate::NdjsonTransport`], -//! reusing the exact same generic subscription framing [`crate::ndjson_watch`]/ -//! [`crate::ndjson_logs`] use. There is no dedup or reconnect logic here, unlike -//! [`crate::NdjsonReconnectingEventStream`] -- `agent/attach` has no cursor to resume from. +//! [`crate::SubscriptionTransport`]-generic `agent/attach` subscription client. Drives +//! `agent/attach`/`event`/`subscription/cancel`/`subscription/closed` notifications over any +//! [`crate::SubscriptionTransport`], reusing the exact same generic subscription framing +//! [`crate::ndjson_watch`]/[`crate::ndjson_logs`] use. [`NdjsonAgentAttachClient`]/ +//! [`NdjsonAgentAttachEventStream`] alias this machinery to [`crate::NdjsonTransport`]. There is +//! no dedup or reconnect logic here, unlike [`crate::NdjsonReconnectingEventStream`] -- +//! `agent/attach` has no cursor to resume from. use crate::ndjson_subscription::impl_ndjson_event_subscription; impl_ndjson_event_subscription! { - client: NdjsonAgentAttachClient, - stream: NdjsonAgentAttachEventStream, + generic_client: AgentAttachSubscriptionClient, + generic_stream: AgentAttachSubscriptionEventStream, + ndjson_client: NdjsonAgentAttachClient, + ndjson_stream: NdjsonAgentAttachEventStream, event_or_closed: AgentAttachEventOrClosed, method_fn: agent_attach, method_name: "agent/attach", diff --git a/crates/openengine-cluster-client/src/ndjson_logs.rs b/crates/openengine-cluster-client/src/ndjson_logs.rs index dae280d2..98bf26f9 100644 --- a/crates/openengine-cluster-client/src/ndjson_logs.rs +++ b/crates/openengine-cluster-client/src/ndjson_logs.rs @@ -1,13 +1,17 @@ -//! NDJSON-bound `logs` subscription client. Drives `logs`/`event`/`subscription/cancel`/ -//! `subscription/closed` notifications over [`crate::NdjsonTransport`], reusing the exact same -//! generic subscription framing [`crate::ndjson_watch`] uses. There is no dedup or reconnect logic -//! here, unlike [`crate::NdjsonReconnectingEventStream`] -- `logs` has no cursor to resume from. +//! [`crate::SubscriptionTransport`]-generic `logs` subscription client. Drives `logs`/`event`/ +//! `subscription/cancel`/`subscription/closed` notifications over any [`crate::SubscriptionTransport`], +//! reusing the exact same generic subscription framing [`crate::ndjson_watch`] uses. +//! [`NdjsonLogsClient`]/[`NdjsonLogsEventStream`] alias this machinery to [`crate::NdjsonTransport`]. +//! There is no dedup or reconnect logic here, unlike [`crate::NdjsonReconnectingEventStream`] -- +//! `logs` has no cursor to resume from. use crate::ndjson_subscription::impl_ndjson_event_subscription; impl_ndjson_event_subscription! { - client: NdjsonLogsClient, - stream: NdjsonLogsEventStream, + generic_client: LogsSubscriptionClient, + generic_stream: LogsSubscriptionEventStream, + ndjson_client: NdjsonLogsClient, + ndjson_stream: NdjsonLogsEventStream, event_or_closed: LogEventOrClosed, method_fn: logs, method_name: "logs", diff --git a/crates/openengine-cluster-client/src/ndjson_pump.rs b/crates/openengine-cluster-client/src/ndjson_pump.rs index 2134d1b7..0c34e7f5 100644 --- a/crates/openengine-cluster-client/src/ndjson_pump.rs +++ b/crates/openengine-cluster-client/src/ndjson_pump.rs @@ -1,12 +1,65 @@ -//! Non-blocking notification routing for the shared NDJSON response pump. +//! Non-blocking pumped-message routing shared by every transport's response pump +//! ([`NdjsonTransport`](crate::NdjsonTransport)'s NDJSON-line pump and +//! [`WebSocketTransport`](crate::websocket::WebSocketTransport)'s `Message::Text`-frame pump): +//! resolving a unary response's pending oneshot (registering a freshly minted subscription's +//! channel first, so no `event` racing the response can be missed), or forwarding a `watch`/ +//! `logs`/`agent_attach` notification to its already-registered subscription channel. -use std::sync::atomic::Ordering; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; -use openengine_cluster_protocol::SubscriptionId; +use openengine_cluster_protocol::{RequestId, SubscriptionId}; use serde_json::Value; use tokio::sync::mpsc; -use super::SubscriptionMap; +use super::{ + PendingMap, PumpedResponse, PumpedSubscription, SubscriptionMap, SubscriptionRegistration, + SUBSCRIPTION_QUEUE_CAPACITY, +}; + +/// Decodes and routes one pumped line: a notification is forwarded live (see +/// [`forward_notification`]); a unary response resolves its pending oneshot, registering a +/// freshly minted subscription's channel first when the response is a successful `watch`-shaped +/// result carrying `result.subscriptionId`. Malformed JSON, a notification/response with no +/// resolvable identity, or an unknown/already-resolved request id are silently dropped -- the +/// same permissive handling `run_pump` always applied inline before this was extracted. Returns +/// the subscription id the caller must write a `subscription/cancel` for, exactly like +/// [`forward_notification`], when a live notification could not be delivered. +pub(super) fn route_pumped_message( + line: String, + pending: &PendingMap, + subscriptions: &SubscriptionMap, +) -> Option { + let Ok(value) = serde_json::from_str::(&line) else { + return None; + }; + if value.get("method").is_some() { + return forward_notification(&value, line, subscriptions); + } + let id = value.get("id").and_then(RequestId::from_json_value)?; + let sender = pending.lock().remove(&id)?; + let subscription = value + .get("result") + .and_then(|result| result.get("subscriptionId")) + .and_then(Value::as_str) + .map(|subscription_id| { + let (sender, receiver) = mpsc::channel(SUBSCRIPTION_QUEUE_CAPACITY); + let overflowed = Arc::new(AtomicBool::new(false)); + subscriptions.lock().insert( + SubscriptionId::new(subscription_id), + SubscriptionRegistration { + sender, + overflowed: Arc::clone(&overflowed), + }, + ); + PumpedSubscription { + receiver, + overflowed, + } + }); + let _ = sender.send(PumpedResponse { line, subscription }); + None +} /// Forwards one `event`/`subscription/closed` notification without waiting on a consumer. /// Returns the subscription id when the local receiver is full or gone and the server must be diff --git a/crates/openengine-cluster-client/src/ndjson_subscription.rs b/crates/openengine-cluster-client/src/ndjson_subscription.rs index 93f205eb..e65a878d 100644 --- a/crates/openengine-cluster-client/src/ndjson_subscription.rs +++ b/crates/openengine-cluster-client/src/ndjson_subscription.rs @@ -1,12 +1,16 @@ -//! Shared NDJSON "one unary response, then live `event`/`subscription/closed` notifications with -//! no dedup or reconnect" client machinery for future-only subscription capabilities (`logs`, -//! `agent_attach`). Generated once per capability via [`impl_ndjson_event_subscription`] rather -//! than hand-copied, so the request/parse/`next`/`cancel` logic exists exactly once. `watch` has -//! different (dedup + reconnect) semantics and is not implemented via this macro. +//! Shared [`crate::SubscriptionTransport`]-generic "one unary response, then live `event`/ +//! `subscription/closed` notifications with no dedup or reconnect" client machinery for +//! future-only subscription capabilities (`logs`, `agent_attach`). Generated once per capability +//! via [`impl_ndjson_event_subscription`] rather than hand-copied, so the request/parse/`next`/ +//! `cancel` logic exists exactly once and is driven identically by [`crate::NdjsonTransport`] and +//! [`crate::websocket::WebSocketTransport`] alike. `watch` has different (dedup + reconnect) +//! semantics and is not implemented via this macro. macro_rules! impl_ndjson_event_subscription { ( - client: $client:ident, - stream: $stream:ident, + generic_client: $client:ident, + generic_stream: $stream:ident, + ndjson_client: $ndjson_client:ident, + ndjson_stream: $ndjson_stream:ident, event_or_closed: $event_or_closed:ident, method_fn: $method_fn:ident, method_name: $method_name:literal, @@ -27,24 +31,26 @@ macro_rules! impl_ndjson_event_subscription { }, } - pub struct $client<'a, R, W> { - transport: &'a crate::NdjsonTransport, + pub struct $client<'a, T> { + transport: &'a T, } - impl<'a, R, W> $client<'a, R, W> + #[doc = concat!("[`", stringify!($client), "`] bound to [`crate::NdjsonTransport`].")] + pub type $ndjson_client<'a, R, W> = $client<'a, crate::NdjsonTransport>; + + impl<'a, T> $client<'a, T> where - R: tokio::io::AsyncRead + Send + Unpin + 'static, - W: tokio::io::AsyncWrite + Send + Unpin + 'static, + T: crate::SubscriptionTransport, { #[must_use] - pub const fn new(transport: &'a crate::NdjsonTransport) -> Self { + pub const fn new(transport: &'a T) -> Self { Self { transport } } pub async fn $method_fn( &self, params: $params_ty, - ) -> Result<($result_ty, $stream<'a, R, W>), crate::ClientError> { + ) -> Result<($result_ty, $stream<'a, T>), crate::ClientError> { let id = self.transport.next_watch_request_id(); let request = serde_json::to_string(&openengine_cluster_protocol::JsonRpcRequest { @@ -100,17 +106,19 @@ macro_rules! impl_ndjson_event_subscription { Ok(response.result) } - pub struct $stream<'a, R, W> { - transport: &'a crate::NdjsonTransport, + pub struct $stream<'a, T> { + transport: &'a T, receiver: tokio::sync::mpsc::Receiver, overflowed: std::sync::Arc, subscription_id: openengine_cluster_protocol::SubscriptionId, } - impl<'a, R, W> $stream<'a, R, W> + #[doc = concat!("[`", stringify!($stream), "`] bound to [`crate::NdjsonTransport`].")] + pub type $ndjson_stream<'a, R, W> = $stream<'a, crate::NdjsonTransport>; + + impl<'a, T> $stream<'a, T> where - R: tokio::io::AsyncRead + Send + Unpin + 'static, - W: tokio::io::AsyncWrite + Send + Unpin + 'static, + T: crate::SubscriptionTransport, { /// Returns the next live event, or a terminal close. Returns `None` once the /// subscription's channel ends (cancelled locally, or the transport's connection diff --git a/crates/openengine-cluster-client/src/ndjson_watch.rs b/crates/openengine-cluster-client/src/ndjson_watch.rs index effc1f61..92e35454 100644 --- a/crates/openengine-cluster-client/src/ndjson_watch.rs +++ b/crates/openengine-cluster-client/src/ndjson_watch.rs @@ -1,7 +1,10 @@ -//! NDJSON-bound watch subscription client. Mirrors [`crate::watch::ReconnectingEventStream`]'s -//! `(runId, cursor)` dedup and reconnect-from-last-delivered-cursor semantics, but drives them -//! over [`crate::NdjsonTransport`]'s wire-framed `watch`/`event`/`subscription/cancel`/ +//! [`SubscriptionTransport`]-generic watch subscription client. Mirrors +//! [`crate::watch::ReconnectingEventStream`]'s `(runId, cursor)` dedup and +//! reconnect-from-last-delivered-cursor semantics, but drives them over any +//! [`SubscriptionTransport`]'s wire-framed `watch`/`event`/`subscription/cancel`/ //! `subscription/closed` notifications instead of the in-process [`Dispatcher`] passthrough. +//! [`NdjsonWatchClient`]/[`NdjsonReconnectingEventStream`] alias this machinery to +//! [`crate::NdjsonTransport`]; [`crate::websocket::WebSocketTransport`] reuses it unchanged. use std::collections::HashSet; use std::sync::atomic::Ordering; @@ -13,34 +16,37 @@ use openengine_cluster_protocol::{ }; use openengine_cluster_server::watch::PublicEventRecord; use serde_json::Value; -use tokio::io::{AsyncRead, AsyncWrite}; use tokio::sync::mpsc; use crate::watch::admit_event; use crate::PumpedSubscription; -use crate::{validate_response_identity, ClientError, EventOrClosed, NdjsonTransport}; +use crate::{ + validate_response_identity, ClientError, EventOrClosed, NdjsonTransport, SubscriptionTransport, +}; -/// Typed NDJSON watch client. Request ids come from the shared [`NdjsonTransport`] rather than a -/// client-local counter, so independently constructed watch clients on one connection cannot -/// replace each other's pending response waiters. -pub struct NdjsonWatchClient<'a, R, W> { - transport: &'a NdjsonTransport, +/// Typed watch client generic over any [`SubscriptionTransport`]. Request ids come from the +/// shared transport rather than a client-local counter, so independently constructed watch +/// clients on one connection cannot replace each other's pending response waiters. +pub struct WatchSubscriptionClient<'a, T> { + transport: &'a T, } -impl<'a, R, W> NdjsonWatchClient<'a, R, W> +/// [`WatchSubscriptionClient`] bound to [`NdjsonTransport`]. +pub type NdjsonWatchClient<'a, R, W> = WatchSubscriptionClient<'a, NdjsonTransport>; + +impl<'a, T> WatchSubscriptionClient<'a, T> where - R: AsyncRead + Send + Unpin + 'static, - W: AsyncWrite + Send + Unpin + 'static, + T: SubscriptionTransport, { #[must_use] - pub const fn new(transport: &'a NdjsonTransport) -> Self { + pub const fn new(transport: &'a T) -> Self { Self { transport } } pub async fn watch( &self, params: WatchParams, - ) -> Result<(WatchResult, NdjsonReconnectingEventStream<'a, R, W>), ClientError> { + ) -> Result<(WatchResult, WatchSubscriptionEventStream<'a, T>), ClientError> { let id = self.next_request_id(); let request = serde_json::to_string(&JsonRpcRequest { jsonrpc: JSON_RPC_VERSION.to_owned(), @@ -57,7 +63,7 @@ where receiver, overflowed, } = subscription.expect("a successful watch response must carry a subscriptionId"); - let stream = NdjsonReconnectingEventStream { + let stream = WatchSubscriptionEventStream { transport: self.transport, receiver, overflowed, @@ -91,9 +97,9 @@ fn parse_watch_response(line: &str, expected_id: &RequestId) -> Result { - transport: &'a NdjsonTransport, +/// sourced from wire notifications forwarded by a [`SubscriptionTransport`]'s pump. +pub struct WatchSubscriptionEventStream<'a, T> { + transport: &'a T, receiver: mpsc::Receiver, overflowed: std::sync::Arc, subscription_id: SubscriptionId, @@ -102,10 +108,13 @@ pub struct NdjsonReconnectingEventStream<'a, R, W> { run_id: Option, } -impl<'a, R, W> NdjsonReconnectingEventStream<'a, R, W> +/// [`WatchSubscriptionEventStream`] bound to [`NdjsonTransport`]. +pub type NdjsonReconnectingEventStream<'a, R, W> = + WatchSubscriptionEventStream<'a, NdjsonTransport>; + +impl<'a, T> WatchSubscriptionEventStream<'a, T> where - R: AsyncRead + Send + Unpin + 'static, - W: AsyncWrite + Send + Unpin + 'static, + T: SubscriptionTransport, { /// Returns the next logically new event, transparently dropping legal duplicate physical /// deliveries, or a terminal close. Returns `None` once the subscription's channel ends @@ -176,8 +185,8 @@ where /// reconnect so a duplicate delivered before and after reconnect is still suppressed once. pub async fn reconnect( self, - ) -> Result<(WatchResult, NdjsonReconnectingEventStream<'a, R, W>), ClientError> { - let watch_client = NdjsonWatchClient::new(self.transport); + ) -> Result<(WatchResult, WatchSubscriptionEventStream<'a, T>), ClientError> { + let watch_client = WatchSubscriptionClient::new(self.transport); let params = WatchParams { run_id: self.run_id, from_cursor: self.last_delivered, diff --git a/crates/openengine-cluster-client/src/websocket.rs b/crates/openengine-cluster-client/src/websocket.rs new file mode 100644 index 00000000..3debe397 --- /dev/null +++ b/crates/openengine-cluster-client/src/websocket.rs @@ -0,0 +1,116 @@ +//! Production WebSocket transport for the typed Cluster Protocol client: demultiplexes unary +//! request/response traffic and generic `watch`/`logs`/`agent_attach` subscription notifications +//! sharing one WebSocket connection, correlating by request id and subscription id respectively. +//! [`WebSocketFrameSink`] backs [`crate::multiplex::FrameSink`], and [`WebSocketTransport`] holds +//! one [`crate::multiplex::MultiplexedTransport`] built from it -- the exact same demux state and +//! [`crate::JsonRpcTransport`]/[`crate::SubscriptionTransport`] wiring [`crate::NdjsonTransport`] +//! holds via [`crate::NdjsonFrameSink`] -- so only the underlying frame shape (NDJSON line vs. +//! `Message::Text`) differs between the two transports. + +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use futures_util::stream::{SplitSink, SplitStream}; +use futures_util::{SinkExt, StreamExt}; +use parking_lot::Mutex as ParkingMutex; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::sync::Mutex; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::WebSocketStream; + +use crate::multiplex; +use crate::{PendingMap, SubscriptionMap, TransportError}; + +/// Sends one already-serialized JSON-RPC frame as a `Message::Text` -- the +/// [`multiplex::FrameSink`] implementation backing [`WebSocketTransport`]. +struct WebSocketFrameSink { + sink: Arc, Message>>>, +} + +#[async_trait] +impl multiplex::FrameSink for WebSocketFrameSink +where + S: AsyncRead + AsyncWrite + Send + Unpin + 'static, +{ + async fn send_frame(&self, frame: String) -> Result<(), TransportError> { + let mut sink = self.sink.lock().await; + sink.send(Message::text(frame)) + .await + .map_err(|error| TransportError::Protocol(error.to_string())) + } +} + +/// WebSocket transport that demultiplexes unary request/response traffic and generic `watch` +/// subscription notifications sharing one connection. Holds one +/// [`multiplex::MultiplexedTransport`], which owns the demux state (write sink, pending-request +/// map, pump task, watch-id counter) and implements every [`JsonRpcTransport`]/ +/// [`SubscriptionTransport`] method against it. +pub struct WebSocketTransport { + inner: multiplex::MultiplexedTransport>, +} + +impl WebSocketTransport +where + S: AsyncRead + AsyncWrite + Send + Unpin + 'static, +{ + #[must_use] + pub fn new(ws: WebSocketStream) -> Self { + let (sink, stream) = ws.split(); + let pending: PendingMap = Arc::new(ParkingMutex::new(HashMap::new())); + let subscriptions: SubscriptionMap = Arc::new(ParkingMutex::new(HashMap::new())); + let sink = WebSocketFrameSink { + sink: Arc::new(Mutex::new(sink)), + }; + let pump = tokio::spawn(run_pump( + stream, + Arc::clone(&pending), + subscriptions, + WebSocketFrameSink { + sink: Arc::clone(&sink.sink), + }, + )); + Self { + inner: multiplex::MultiplexedTransport::new(sink, pending, pump), + } + } +} + +multiplex::impl_multiplexed_transport!( + WebSocketTransport where + S: AsyncRead + AsyncWrite + Send + Unpin + 'static +); + +/// Drives the read half: decodes `Message::Text` frames (one JSON-RPC object per frame -- no +/// reassembly needed, unlike NDJSON's newline-delimited lines) and routes each one via +/// [`multiplex::route_and_maybe_cancel`] -- shared verbatim with [`crate::NdjsonTransport`]'s pump, +/// which routes the exact same decoded JSON bodies sourced from NDJSON lines instead of +/// `Message::Text` frames. Non-text frames (`Binary`/`Ping`/`Pong`/`Frame`) are ignored; a `Close` +/// frame or a read error ends the pump, exactly like NDJSON's stream-end handling. On stream end +/// every pending request fails and every open subscription ends (dropping its sender). +async fn run_pump( + mut stream: SplitStream>, + pending: PendingMap, + subscriptions: SubscriptionMap, + sink: WebSocketFrameSink, +) where + S: AsyncRead + AsyncWrite + Send + Unpin + 'static, +{ + while let Some(Ok(message)) = stream.next().await { + let text = match message { + Message::Text(text) => text, + Message::Close(_) => break, + Message::Binary(_) | Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => { + continue; + } + }; + multiplex::route_and_maybe_cancel( + text.as_str().to_owned(), + &pending, + &subscriptions, + &sink, + ) + .await; + } + multiplex::finish_pump(&pending, &subscriptions); +} diff --git a/crates/openengine-cluster-client/tests/backend_faults.rs b/crates/openengine-cluster-client/tests/backend_faults.rs index 4efa8985..3e6b6be5 100644 --- a/crates/openengine-cluster-client/tests/backend_faults.rs +++ b/crates/openengine-cluster-client/tests/backend_faults.rs @@ -22,6 +22,9 @@ use tokio::task::JoinHandle; mod reconnect_support; use reconnect_support::FIXTURE_QUEUE_CAPACITY; +#[path = "reconnect_support/scenario_harness.rs"] +mod scenario_harness; + #[path = "reconnect_support/ndjson_scenario.rs"] mod ndjson_scenario; use ndjson_scenario::ndjson_overflow_and_reconnect_scenario; diff --git a/crates/openengine-cluster-client/tests/reconnect_support/cancel_scenario.rs b/crates/openengine-cluster-client/tests/reconnect_support/cancel_scenario.rs new file mode 100644 index 00000000..b64cf0ef --- /dev/null +++ b/crates/openengine-cluster-client/tests/reconnect_support/cancel_scenario.rs @@ -0,0 +1,74 @@ +//! Transport-generic driver for the shared `subscription/cancel`-stops-delivery scenario, used +//! directly by `tests/subscription_ndjson.rs` and `tests/websocket.rs`. Kept out of +//! `scenario_harness.rs` (rather than folded in) so `tests/backend_faults.rs`, which never +//! exercises `subscription/cancel`, doesn't need to compile or reference this scenario. + +use std::sync::Arc; +use std::time::Duration; + +use openengine_cluster_client::{ClusterClient, EventOrClosed, WatchSubscriptionClient}; +use openengine_cluster_protocol::{Cursor, GetParams, RunId, WatchEvent, WatchParams}; +use openengine_cluster_server::watch::fixtures::{FixtureBackend, FixtureStore}; + +use crate::scenario_harness::ScenarioTransport; + +/// Connects a fresh `FixtureStore`/`T` pair, cancels a live watch subscription after its first +/// event, and proves cancellation stops delivery: a synchronous `get` round trip on the same +/// connection forces the server to have already applied the preceding fire-and-forget +/// `subscription/cancel` (requests on one connection are read and handled strictly in order), then +/// at most one further event -- the one the server-side subscription task may already have been +/// parked awaiting when it observed cancellation -- may still leak before delivery stops for good. +/// Generic over [`ScenarioTransport`] so every wire binding shares this exact scenario logic +/// rather than duplicating it per transport. +pub async fn run_cancel_stops_delivery_scenario() { + let run_id = RunId::new("run-1"); + let store = Arc::new(FixtureStore::new(run_id, Vec::new(), 8)); + let (transport, server) = T::spawn(FixtureBackend::new(Arc::clone(&store))).await; + + let watch_client = WatchSubscriptionClient::new(&transport); + let (_result, mut stream) = watch_client.watch(WatchParams::default()).await.unwrap(); + + store.publish(WatchEvent::Bookmark).await; + match stream.next().await.unwrap() { + EventOrClosed::Event(record) => assert_eq!(record.cursor, Cursor::new("cursor-1")), + other => panic!("expected an event, got {other:?}"), + } + + stream.cancel().await.unwrap(); + + ClusterClient::new(&transport) + .get(GetParams::default()) + .await + .unwrap(); + + // The server-side subscription task may already be parked awaiting the next live event at + // the moment cancellation is processed, so at most one further event (the one immediately + // following cancellation) may still be delivered before it observes cancellation and stops + // for good. + store.publish(WatchEvent::Bookmark).await; // may leak as cursor-2 + store.publish(WatchEvent::Bookmark).await; // cursor-3, must never arrive + + let mut leaked = Vec::new(); + loop { + match tokio::time::timeout(Duration::from_millis(300), stream.next()).await { + Ok(Some(EventOrClosed::Event(record))) => leaked.push(record.cursor), + Ok(Some(other)) => panic!("unexpected notification after cancel: {other:?}"), + Ok(None) | Err(_) => break, + } + } + assert!( + leaked.len() <= 1, + "cancelled subscription received more than one post-cancel event: {leaked:?}" + ); + if let Some(cursor) = leaked.first() { + assert_eq!( + *cursor, + Cursor::new("cursor-2"), + "cancellation failed to stop delivery before the next published event" + ); + } + + drop(stream); + drop(transport); + T::shutdown(server).await; +} diff --git a/crates/openengine-cluster-client/tests/reconnect_support/mod.rs b/crates/openengine-cluster-client/tests/reconnect_support/mod.rs index 8f1df40c..d62f9cae 100644 --- a/crates/openengine-cluster-client/tests/reconnect_support/mod.rs +++ b/crates/openengine-cluster-client/tests/reconnect_support/mod.rs @@ -1,15 +1,17 @@ //! Shared slow-consumer-overflow / gap-free-reconnect-replay / dedup-and-`Finished` scenario //! driven identically by `tests/reconnect.rs` (in-process `WatchClient`) and, via -//! `ndjson_scenario`, `tests/subscription_ndjson.rs` / `tests/backend_faults.rs` -//! (`NdjsonWatchClient` over the wire), proving both transports produce byte-equivalent -//! event/cursor sequences and that the `fault` `WatchEvent` variant needs no special-casing. +//! `ndjson_scenario`/`websocket_scenario` on top of `scenario_harness`, `tests/subscription_ndjson.rs` +//! / `tests/websocket.rs` / `tests/backend_faults.rs` (`NdjsonWatchClient`/`WebSocketTransport` over +//! the wire), proving all transports produce byte-equivalent event/cursor sequences and that the +//! `fault` `WatchEvent` variant needs no special-casing. -use openengine_cluster_client::{EventOrClosed, NdjsonReconnectingEventStream, ReconnectingEventStream}; +use openengine_cluster_client::{ + EventOrClosed, ReconnectingEventStream, SubscriptionTransport, WatchSubscriptionEventStream, +}; use openengine_cluster_protocol::{ ClusterStatus, Cursor, StopMode, SubscriptionCloseReason, WatchEvent, }; use openengine_cluster_server::watch::fixtures::FixtureStore; -use tokio::io::{AsyncRead, AsyncWrite}; /// The fixture's queue is bounded to 2 entries; the third publish overflows it, forcing a /// deterministic `SLOW_CONSUMER` close for these tests to reconnect from. @@ -27,10 +29,14 @@ impl EventStream for ReconnectingEventStream { } } -impl<'a, R, W> EventStream for NdjsonReconnectingEventStream<'a, R, W> +/// Covers every [`SubscriptionTransport`]-generic wire binding at once (NDJSON's +/// `NdjsonReconnectingEventStream` alias and `WebSocketTransport`'s +/// `WatchSubscriptionEventStream<'a, WebSocketTransport>`), since both are the same generic +/// type parameterized by transport -- proving the two wire bindings share this exact scenario +/// logic rather than each needing their own impl. +impl<'a, T> EventStream for WatchSubscriptionEventStream<'a, T> where - R: AsyncRead + Send + Unpin + 'static, - W: AsyncWrite + Send + Unpin + 'static, + T: SubscriptionTransport, { async fn next_event(&mut self) -> Option { self.next().await diff --git a/crates/openengine-cluster-client/tests/reconnect_support/ndjson_scenario.rs b/crates/openengine-cluster-client/tests/reconnect_support/ndjson_scenario.rs index 401d3def..025c6f78 100644 --- a/crates/openengine-cluster-client/tests/reconnect_support/ndjson_scenario.rs +++ b/crates/openengine-cluster-client/tests/reconnect_support/ndjson_scenario.rs @@ -2,17 +2,30 @@ //! (`Bookmark`) and `tests/backend_faults.rs` (`fault`) so the two scenarios stay byte-for-byte //! identical apart from the published event. Kept separate from `reconnect_support/mod.rs` so //! `tests/reconnect.rs`, which drives the in-process `WatchClient` instead of NDJSON, doesn't -//! need to compile or reference NDJSON-only wiring it never calls. +//! need to compile or reference NDJSON-only wiring it never calls. The scenario body itself lives +//! in `scenario_harness`'s transport-generic `run_overflow_and_reconnect_scenario`; this module +//! only wires `NdjsonTransport`'s spawn/shutdown into that shared driver. -use std::sync::Arc; +use std::io; -use openengine_cluster_client::{NdjsonTransport, NdjsonWatchClient}; -use openengine_cluster_protocol::{RunId, WatchEvent, WatchParams}; -use openengine_cluster_server::watch::fixtures::{ - await_ndjson_shutdown, spawn_ndjson, FixtureBackend, FixtureStore, -}; +use openengine_cluster_client::NdjsonTransport; +use openengine_cluster_protocol::{RunId, WatchEvent}; +use openengine_cluster_server::watch::fixtures::{await_ndjson_shutdown, spawn_ndjson, FixtureBackend}; +use tokio::io::DuplexStream; +use tokio::task::JoinHandle; -use crate::reconnect_support::{assert_reconnect_replays_and_dedups, overflow_and_close_with}; +use crate::scenario_harness::{run_overflow_and_reconnect_scenario, ScenarioTransport}; + +impl ScenarioTransport for NdjsonTransport { + async fn spawn(backend: FixtureBackend) -> (Self, JoinHandle>) { + let (client_write, client_read, server) = spawn_ndjson(backend); + (NdjsonTransport::new(client_read, client_write), server) + } + + async fn shutdown(server: JoinHandle>) { + await_ndjson_shutdown(server).await; + } +} /// Connects a fresh `FixtureStore`/`serve_ndjson` pair for `run_id` with the given subscription /// queue `capacity`, then drives the full overflow / reconnect / dedup sequence with `event()` as @@ -20,22 +33,10 @@ use crate::reconnect_support::{assert_reconnect_replays_and_dedups, overflow_and pub async fn ndjson_overflow_and_reconnect_scenario( run_id: RunId, capacity: usize, - mut event: impl FnMut() -> WatchEvent, + event: impl FnMut() -> WatchEvent, ) { - let store = Arc::new(FixtureStore::new(run_id.clone(), Vec::new(), capacity)); - let (client_write, client_read, server) = spawn_ndjson(FixtureBackend::new(Arc::clone(&store))); - - let transport = NdjsonTransport::new(client_read, client_write); - let watch_client = NdjsonWatchClient::new(&transport); - let (result, mut stream) = watch_client.watch(WatchParams::default()).await.unwrap(); - assert_eq!(result.run_id, Some(run_id)); - - let received = overflow_and_close_with(&store, &mut stream, &mut event).await; - - let (_result, mut stream) = stream.reconnect().await.unwrap(); - assert_reconnect_replays_and_dedups(&store, &mut stream, received).await; - - drop(stream); - drop(transport); - await_ndjson_shutdown(server).await; + run_overflow_and_reconnect_scenario::>( + run_id, capacity, event, + ) + .await; } diff --git a/crates/openengine-cluster-client/tests/reconnect_support/scenario_harness.rs b/crates/openengine-cluster-client/tests/reconnect_support/scenario_harness.rs new file mode 100644 index 00000000..5b947881 --- /dev/null +++ b/crates/openengine-cluster-client/tests/reconnect_support/scenario_harness.rs @@ -0,0 +1,51 @@ +//! Transport-generic driver for the shared overflow/reconnect/dedup scenario in +//! `reconnect_support::mod`, used by `ndjson_scenario` and `websocket_scenario` so every wire +//! binding runs the exact same scenario body instead of each copying it. Kept out of +//! `reconnect_support/mod.rs` (rather than folded in) so `tests/reconnect.rs`, which drives the +//! in-process `WatchClient` and never spawns a wire transport, doesn't need to compile or +//! reference this wire-only glue. + +use std::io; +use std::sync::Arc; + +use openengine_cluster_client::{SubscriptionTransport, WatchSubscriptionClient}; +use openengine_cluster_protocol::{RunId, WatchEvent, WatchParams}; +use openengine_cluster_server::watch::fixtures::{FixtureBackend, FixtureStore}; +use tokio::task::JoinHandle; + +use crate::reconnect_support::{assert_reconnect_replays_and_dedups, overflow_and_close_with}; + +/// Spawn/shutdown glue letting [`run_overflow_and_reconnect_scenario`] (and, via +/// `cancel_scenario`, `run_cancel_stops_delivery_scenario`) drive identical scenarios over any +/// [`SubscriptionTransport`]-generic wire binding (NDJSON, WebSocket, ...) without each binding +/// needing its own copy of the scenario body. +pub trait ScenarioTransport: SubscriptionTransport + Sized { + async fn spawn(backend: FixtureBackend) -> (Self, JoinHandle>); + async fn shutdown(server: JoinHandle>); +} + +/// Connects a fresh `FixtureStore`/`T` pair for `run_id` with the given subscription queue +/// `capacity`, then drives the full overflow / reconnect / dedup sequence with `event()` as the +/// published `WatchEvent`. Generic over [`ScenarioTransport`] so every wire binding shares this +/// exact scenario logic rather than duplicating it per transport. +pub async fn run_overflow_and_reconnect_scenario( + run_id: RunId, + capacity: usize, + mut event: impl FnMut() -> WatchEvent, +) { + let store = Arc::new(FixtureStore::new(run_id.clone(), Vec::new(), capacity)); + let (transport, server) = T::spawn(FixtureBackend::new(Arc::clone(&store))).await; + + let watch_client = WatchSubscriptionClient::new(&transport); + let (result, mut stream) = watch_client.watch(WatchParams::default()).await.unwrap(); + assert_eq!(result.run_id, Some(run_id)); + + let received = overflow_and_close_with(&store, &mut stream, &mut event).await; + + let (_result, mut stream) = stream.reconnect().await.unwrap(); + assert_reconnect_replays_and_dedups(&store, &mut stream, received).await; + + drop(stream); + drop(transport); + T::shutdown(server).await; +} diff --git a/crates/openengine-cluster-client/tests/reconnect_support/websocket_scenario.rs b/crates/openengine-cluster-client/tests/reconnect_support/websocket_scenario.rs new file mode 100644 index 00000000..5591f855 --- /dev/null +++ b/crates/openengine-cluster-client/tests/reconnect_support/websocket_scenario.rs @@ -0,0 +1,43 @@ +//! End-to-end WebSocket overflow/reconnect scenario, mirroring `reconnect_support::ndjson_scenario` +//! but driven over `WebSocketTransport` and `serve_websocket` instead of NDJSON, proving the two +//! wire bindings reproduce byte-equivalent event/cursor sequences through the exact same generic +//! `SubscriptionTransport`/`WatchSubscriptionClient` machinery. The scenario body itself lives in +//! `scenario_harness`'s transport-generic `run_overflow_and_reconnect_scenario`; this module only +//! wires `WebSocketTransport`'s spawn/shutdown into that shared driver. + +use std::io; + +use openengine_cluster_client::WebSocketTransport; +use openengine_cluster_protocol::{RunId, WatchEvent}; +use openengine_cluster_server::watch::fixtures::{ + await_websocket_shutdown, spawn_websocket, FixtureBackend, +}; +use tokio::io::DuplexStream; +use tokio::task::JoinHandle; + +use crate::scenario_harness::{run_overflow_and_reconnect_scenario, ScenarioTransport}; + +impl ScenarioTransport for WebSocketTransport { + async fn spawn(backend: FixtureBackend) -> (Self, JoinHandle>) { + let (ws, server) = spawn_websocket(backend).await; + (WebSocketTransport::new(ws), server) + } + + async fn shutdown(server: JoinHandle>) { + await_websocket_shutdown(server).await; + } +} + +/// Connects a fresh `FixtureStore`/`serve_websocket` pair for `run_id` with the given subscription +/// queue `capacity`, then drives the full overflow / reconnect / dedup sequence with `event()` as +/// the published `WatchEvent`. +pub async fn websocket_overflow_and_reconnect_scenario( + run_id: RunId, + capacity: usize, + event: impl FnMut() -> WatchEvent, +) { + run_overflow_and_reconnect_scenario::>( + run_id, capacity, event, + ) + .await; +} diff --git a/crates/openengine-cluster-client/tests/subscription_ndjson.rs b/crates/openengine-cluster-client/tests/subscription_ndjson.rs index f027007c..88ec3448 100644 --- a/crates/openengine-cluster-client/tests/subscription_ndjson.rs +++ b/crates/openengine-cluster-client/tests/subscription_ndjson.rs @@ -4,18 +4,10 @@ //! duplicates, driven over the wire against `serve_ndjson` instead of the in-process //! `Dispatcher::watch` passthrough exercised by `tests/reconnect.rs`. -use std::sync::Arc; use std::time::Duration; -use openengine_cluster_client::{ - ClusterClient, EventOrClosed, JsonRpcTransport, NdjsonTransport, NdjsonWatchClient, -}; -use openengine_cluster_protocol::{ - Cursor, GetParams, RunId, SubscriptionCloseReason, WatchEvent, WatchParams, -}; -use openengine_cluster_server::watch::fixtures::{ - await_ndjson_shutdown, spawn_ndjson, FixtureBackend, FixtureStore, -}; +use openengine_cluster_client::{EventOrClosed, JsonRpcTransport, NdjsonTransport, NdjsonWatchClient}; +use openengine_cluster_protocol::{Cursor, RunId, SubscriptionCloseReason, WatchEvent, WatchParams}; use serde_json::{json, Value}; use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream}; @@ -23,6 +15,13 @@ use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, DuplexStream}; mod reconnect_support; use reconnect_support::FIXTURE_QUEUE_CAPACITY; +#[path = "reconnect_support/scenario_harness.rs"] +mod scenario_harness; + +#[path = "reconnect_support/cancel_scenario.rs"] +mod cancel_scenario; +use cancel_scenario::run_cancel_stops_delivery_scenario; + #[path = "reconnect_support/ndjson_scenario.rs"] mod ndjson_scenario; use ndjson_scenario::ndjson_overflow_and_reconnect_scenario; @@ -35,66 +34,17 @@ async fn reconnect_after_slow_consumer_recovers_with_no_gap_and_dedups_duplicate .await; } +// `cancel()` only writes a fire-and-forget notification line -- it does not wait for the server +// to have applied it. `run_cancel_stops_delivery_scenario` forces a synchronous `get` round trip +// on the same connection so the subsequent publishes are guaranteed to happen only after the +// server's read loop has already processed (and synchronously applied) the preceding cancel line; +// NDJSON lines are read and handled strictly in order, so a response to a request sent after +// `cancel` can only arrive once the cancel itself has already been read and applied. No +// `subscription/closed` follows a plain cancel, so absence of further delivery is observed as the +// stream's `next()` simply never resolving again. #[tokio::test] async fn cancel_stops_further_delivery() { - let run_id = RunId::new("run-1"); - let store = Arc::new(FixtureStore::new(run_id, Vec::new(), 8)); - let (client_write, client_read, server) = spawn_ndjson(FixtureBackend::new(Arc::clone(&store))); - - let transport = NdjsonTransport::new(client_read, client_write); - let watch_client = NdjsonWatchClient::new(&transport); - let (_result, mut stream) = watch_client.watch(WatchParams::default()).await.unwrap(); - - store.publish(WatchEvent::Bookmark).await; - match stream.next().await.unwrap() { - EventOrClosed::Event(record) => assert_eq!(record.cursor, Cursor::new("cursor-1")), - other => panic!("expected an event, got {other:?}"), - } - - stream.cancel().await.unwrap(); - - // `cancel()` only writes a fire-and-forget notification line -- it does not wait for the - // server to have applied it. Force a synchronous round trip on the same connection so the - // subsequent publishes are guaranteed to happen only after the server's read loop has already - // processed (and synchronously applied) the preceding cancel line; NDJSON lines are read and - // handled strictly in order, so a response to a request sent after `cancel` can only arrive - // once the cancel itself has already been read and applied. - ClusterClient::new(&transport) - .get(GetParams::default()) - .await - .unwrap(); - - // The server-side subscription task may already be parked awaiting the next live event at - // the moment cancellation is processed, so at most one further event (the one immediately - // following cancellation) may still be delivered before it observes cancellation and stops - // for good; no `subscription/closed` follows a plain cancel, so absence of further delivery - // is observed as this stream's `next()` simply never resolving again. - store.publish(WatchEvent::Bookmark).await; // may leak as cursor-2 - store.publish(WatchEvent::Bookmark).await; // cursor-3, must never arrive - - let mut leaked = Vec::new(); - loop { - match tokio::time::timeout(Duration::from_millis(300), stream.next()).await { - Ok(Some(EventOrClosed::Event(record))) => leaked.push(record.cursor), - Ok(Some(other)) => panic!("unexpected notification after cancel: {other:?}"), - Ok(None) | Err(_) => break, - } - } - assert!( - leaked.len() <= 1, - "cancelled subscription received more than one post-cancel event: {leaked:?}" - ); - if let Some(cursor) = leaked.first() { - assert_eq!( - *cursor, - Cursor::new("cursor-2"), - "cancellation failed to stop delivery before the next published event" - ); - } - - drop(stream); - drop(transport); - await_ndjson_shutdown(server).await; + run_cancel_stops_delivery_scenario::>().await; } async fn write_json_line(writer: &mut DuplexStream, value: Value) { diff --git a/crates/openengine-cluster-client/tests/websocket.rs b/crates/openengine-cluster-client/tests/websocket.rs new file mode 100644 index 00000000..f095fea0 --- /dev/null +++ b/crates/openengine-cluster-client/tests/websocket.rs @@ -0,0 +1,168 @@ +//! Client-side WebSocket transport: exercises `WebSocketTransport` -- the production WebSocket +//! counterpart to `NdjsonTransport` -- against a real `serve_websocket` server for unary +//! round-trips, `WatchSubscriptionClient` event/closed transcripts (slow-consumer overflow and +//! gap-free reconnect), `subscription/cancel`, and best-effort `$/cancelRequest` in-flight +//! cancellation. The counterpart of `openengine-cluster-server`'s own `tests/websocket.rs`, which +//! drives the same `serve_websocket` binding from raw tungstenite frames instead of this crate's +//! typed client. + +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; +use std::time::Duration; + +use async_trait::async_trait; +use openengine_cluster_client::{ + ClusterClient, JsonRpcTransport, SubscriptionTransport, WebSocketTransport, +}; +use openengine_cluster_protocol::{ + GetParams, GetResult, InitializeParams, InitializeResult, RequestId, RunId, WatchEvent, + WatchParams, PROTOCOL_VERSION, +}; +use openengine_cluster_server::watch::fixtures::{ + await_websocket_shutdown, spawn_websocket, FixtureBackend, FixtureStore, +}; +use openengine_cluster_server::watch::{WatchEventStream, WatchHandle}; +use openengine_cluster_server::{BackendError, ClusterBackend, ConnectionContext}; +use serde_json::json; +use tokio::io::DuplexStream; +use tokio::sync::Notify; + +#[path = "reconnect_support/mod.rs"] +mod reconnect_support; +use reconnect_support::FIXTURE_QUEUE_CAPACITY; + +#[path = "reconnect_support/scenario_harness.rs"] +mod scenario_harness; + +#[path = "reconnect_support/cancel_scenario.rs"] +mod cancel_scenario; +use cancel_scenario::run_cancel_stops_delivery_scenario; + +#[path = "reconnect_support/websocket_scenario.rs"] +mod websocket_scenario; +use websocket_scenario::websocket_overflow_and_reconnect_scenario; + +#[tokio::test] +async fn unary_initialize_and_get_round_trip_over_websocket_transport() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let (ws, server) = spawn_websocket(FixtureBackend::new(store)).await; + let transport = WebSocketTransport::new(ws); + let client = ClusterClient::new(&transport); + + let init = client.initialize().await.unwrap(); + assert_eq!(init.protocol_version, PROTOCOL_VERSION); + let get_result = client.get(GetParams::default()).await.unwrap(); + assert!(get_result.spec.is_none()); + + drop(transport); + await_websocket_shutdown(server).await; +} + +#[tokio::test] +async fn reconnect_after_slow_consumer_recovers_with_no_gap_and_dedups_duplicates_over_websocket() { + websocket_overflow_and_reconnect_scenario(RunId::new("run-1"), FIXTURE_QUEUE_CAPACITY, || { + WatchEvent::Bookmark + }) + .await; +} + +// `cancel()` only writes a fire-and-forget notification -- it does not wait for the server to +// have applied it. `run_cancel_stops_delivery_scenario` forces a synchronous `get` round trip on +// the same connection so the subsequent publishes are guaranteed to happen only after the +// server's read loop has already processed (and synchronously applied) the preceding cancel, +// mirroring `subscription_ndjson.rs`'s identical NDJSON cancel test. +#[tokio::test] +async fn cancel_subscription_ends_the_stream_over_websocket() { + run_cancel_stops_delivery_scenario::>().await; +} + +/// Wraps [`FixtureBackend`], gating only `get` on an explicit [`Notify`] and counting only calls +/// that actually complete past the gate -- so a `$/cancelRequest`-aborted call can never bump the +/// counter, letting a test prove no backend effect occurred. `initialize` is left ungated so a +/// test can prove the connection remains usable after aborting a gated `get`. +struct GatedBackend { + inner: FixtureBackend, + gate: Arc, + completed_gets: Arc, +} + +#[async_trait] +impl ClusterBackend for GatedBackend { + async fn initialize( + &self, + context: &ConnectionContext, + params: InitializeParams, + ) -> Result { + self.inner.initialize(context, params).await + } + + async fn get( + &self, + context: &ConnectionContext, + params: GetParams, + ) -> Result { + self.gate.notified().await; + let result = self.inner.get(context, params).await; + self.completed_gets.fetch_add(1, Ordering::SeqCst); + result + } + + async fn watch( + &self, + context: &ConnectionContext, + params: WatchParams, + queue_capacity: usize, + ) -> Result { + self.inner.watch(context, params, queue_capacity).await + } +} + +type WatchResultTriple = ( + openengine_cluster_protocol::WatchResult, + WatchEventStream, + WatchHandle, +); + +#[tokio::test] +async fn cancel_request_aborts_in_flight_response_with_no_effect_and_keeps_connection_usable() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let completed_gets = Arc::new(AtomicUsize::new(0)); + let gate = Arc::new(Notify::new()); + let (ws, server) = spawn_websocket(GatedBackend { + inner: FixtureBackend::new(store), + gate: Arc::clone(&gate), + completed_gets: Arc::clone(&completed_gets), + }) + .await; + let transport = WebSocketTransport::new(ws); + + let request = json!({ + "jsonrpc": "2.0", "id": "gated-get", "method": "get", "params": {} + }) + .to_string(); + let mut pending = transport.request(request); + tokio::select! { + _ = &mut pending => panic!("gated get must not resolve before its gate is released"), + () = tokio::time::sleep(Duration::from_millis(50)) => {} + } + + // AC5: best-effort `$/cancelRequest` aborts response delivery for the in-flight request. + transport + .cancel_request(RequestId::String("gated-get".to_owned())) + .await + .unwrap(); + + let outcome = tokio::time::timeout(Duration::from_millis(300), pending).await; + assert!(outcome.is_err(), "a cancelled request must never resolve"); + assert_eq!( + completed_gets.load(Ordering::SeqCst), + 0, + "committed backend state must be unchanged -- the aborted call must never have completed" + ); + + // AC5: the connection remains usable after the cancellation. + ClusterClient::new(&transport).initialize().await.unwrap(); + + drop(transport); + await_websocket_shutdown(server).await; +} diff --git a/crates/openengine-cluster-protocol/src/watch.rs b/crates/openengine-cluster-protocol/src/watch.rs index a5731180..53c7fadf 100644 --- a/crates/openengine-cluster-protocol/src/watch.rs +++ b/crates/openengine-cluster-protocol/src/watch.rs @@ -5,8 +5,8 @@ use serde::{Deserialize, Serialize}; use serde_json::Value; use crate::{ - BackendFault, ClusterStatus, Cursor, GraphSpec, NodeName, PositiveInteger, RunId, StopMode, - WorkerOutcome, + BackendFault, ClusterStatus, Cursor, GraphSpec, NodeName, PositiveInteger, RequestId, RunId, + StopMode, WorkerOutcome, }; pub const NOT_FOUND: &str = "NOT_FOUND"; @@ -117,6 +117,16 @@ pub struct SubscriptionCancelParams { pub subscription_id: SubscriptionId, } +/// Wire body of the `$/cancelRequest` client notification: best-effort cancellation of an +/// in-flight unary request by its `RequestId`. Unknown or already-completed ids are a silent +/// no-op; cancelling after the backend has committed leaves that committed state unchanged and +/// emits no response or compensation. +#[derive(Clone, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)] +#[serde(deny_unknown_fields, rename_all = "camelCase")] +pub struct CancelRequestParams { + pub id: RequestId, +} + #[derive(Clone, Copy, Debug, Deserialize, Eq, JsonSchema, PartialEq, Serialize)] pub enum SubscriptionCloseReason { #[serde(rename = "done")] diff --git a/crates/openengine-cluster-server/Cargo.toml b/crates/openengine-cluster-server/Cargo.toml index db2d4069..e9edd5c8 100644 --- a/crates/openengine-cluster-server/Cargo.toml +++ b/crates/openengine-cluster-server/Cargo.toml @@ -8,6 +8,7 @@ repository.workspace = true [dependencies] async-trait.workspace = true +futures-util.workspace = true openengine-cluster-protocol.workspace = true parking_lot.workspace = true serde.workspace = true @@ -15,4 +16,5 @@ serde_json.workspace = true thiserror.workspace = true tokio.workspace = true tokio-stream.workspace = true +tokio-tungstenite.workspace = true tokio-util.workspace = true diff --git a/crates/openengine-cluster-server/src/lib.rs b/crates/openengine-cluster-server/src/lib.rs index 5564cd71..9ad83e85 100644 --- a/crates/openengine-cluster-server/src/lib.rs +++ b/crates/openengine-cluster-server/src/lib.rs @@ -7,6 +7,7 @@ pub mod lifecycle; pub mod logs; pub mod stdio; pub mod watch; +pub mod websocket; pub mod worker_registry; mod dispatch; diff --git a/crates/openengine-cluster-server/src/stdio.rs b/crates/openengine-cluster-server/src/stdio.rs index f6a7dc31..ce7e81e2 100644 --- a/crates/openengine-cluster-server/src/stdio.rs +++ b/crates/openengine-cluster-server/src/stdio.rs @@ -1,17 +1,22 @@ //! NDJSON stdio transport: multiplexes unary JSON-RPC request/response traffic and generic //! `watch`/`logs`/`agent/attach` subscription notifications over one bounded-frame connection. -mod admission; -mod agent_attach; -mod logs; -mod subscription; +pub(crate) mod admission; +pub(crate) mod agent_attach; +pub(crate) mod dispatch; +pub(crate) mod logs; +pub(crate) mod subscription; -use admission::{acquire_task_slot, reject_duplicate, run_writer, InFlightIds, MAX_CONNECTION_TASKS}; +pub(crate) use dispatch::{ + dispatch_classified_line, new_connection_setup, shutdown_connection, ConnectionSetup, + DispatchCtx, LineDispatch, ShutdownArgs, +}; + +use admission::{run_writer, InFlightIds}; -use std::collections::{HashMap, HashSet}; +use std::collections::HashMap; use std::io; use std::sync::Arc; -use std::time::Duration; use parking_lot::Mutex; use openengine_cluster_protocol::{ @@ -21,8 +26,7 @@ use openengine_cluster_protocol::{ }; use serde_json::Value; use tokio::io::{AsyncRead, AsyncWrite, AsyncWriteExt}; -use tokio::sync::{mpsc, Notify, Semaphore}; -use tokio::task::JoinSet; +use tokio::sync::{mpsc, Notify}; use tokio_stream::StreamExt; use tokio_util::codec::{Framed, LinesCodec, LinesCodecError}; @@ -38,25 +42,22 @@ const MAX_FRAME_BYTES: usize = 1_048_576; /// than growing memory without bound. const OUTBOUND_QUEUE_CAPACITY: usize = 256; -/// Grace period given to already-spawned bounded backend dispatches to finish once the connection -/// closes. Subscription tasks are notified through their cancellation handles before shutdown; -/// any backend operation that does not finish inside this bound is force-aborted. -const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_millis(200); - /// Per-subscription cancellation signal: notifying it wakes `run_watch_subscription`'s streaming /// loop immediately, even while parked awaiting the next live event, instead of relying solely on /// `WatchEventStream`'s own cancelled flag, which is only re-checked at the top of `next()` and so /// never observed on an idle run once the task is parked inside `next_live`'s /// `receiver.recv().await`. -type SubscriptionMap = Arc>>>; +pub(crate) type SubscriptionMap = Arc>>>; /// Per-connection state shared by every spawned request/subscription task: the outbound write -/// queue and the tracking maps used for cancellation and duplicate-id rejection. +/// queue and the tracking maps used for cancellation and duplicate-id rejection. Shared verbatim +/// with the sibling `websocket` transport module so both bindings drive the exact same +/// subscription-establishment and cancellation machinery. #[derive(Clone)] -struct ConnectionState { - outbound_tx: mpsc::Sender, - subscriptions: SubscriptionMap, - in_flight_ids: InFlightIds, +pub(crate) struct ConnectionState { + pub(crate) outbound_tx: mpsc::Sender, + pub(crate) subscriptions: SubscriptionMap, + pub(crate) in_flight_ids: InFlightIds, } /// Races `next` against `cancel`, `biased` toward the cancellation so a `subscription/cancel` that @@ -64,8 +65,9 @@ struct ConnectionState { /// being observed the next time the stream is polled -- which never happens again on an idle run. /// `biased` also ensures a pending cancellation is never starved by an unbounded run of /// already-buffered stream items. Shared by `run_watch_subscription` and -/// `subscription::run_bounded_event_subscription`. -pub(super) async fn race_cancel_or_next( +/// `subscription::run_bounded_event_subscription`, and by the sibling `websocket` transport +/// module's identical subscription runner reuse. +pub(crate) async fn race_cancel_or_next( cancel: &Notify, next: impl std::future::Future>, ) -> Option { @@ -91,15 +93,12 @@ where let (outbound_tx, outbound_rx) = mpsc::channel::(OUTBOUND_QUEUE_CAPACITY); let writer_task = tokio::spawn(run_writer(writer, outbound_rx)); - let subscriptions: SubscriptionMap = Arc::new(Mutex::new(HashMap::new())); - let in_flight_ids: InFlightIds = Arc::new(Mutex::new(HashSet::new())); - let task_slots = Arc::new(Semaphore::new(MAX_CONNECTION_TASKS)); - let mut tasks: JoinSet<()> = JoinSet::new(); - let state = ConnectionState { - outbound_tx: outbound_tx.clone(), - subscriptions: Arc::clone(&subscriptions), - in_flight_ids: Arc::clone(&in_flight_ids), - }; + let ConnectionSetup { + subscriptions, + task_slots, + mut tasks, + state, + } = new_connection_setup(&outbound_tx); let mut lines = Framed::new(reader, LinesCodec::new_with_max_length(MAX_FRAME_BYTES)); loop { @@ -141,107 +140,33 @@ where None => break, }; - match classify_ndjson_line(&line) { - NdjsonLineKind::Cancel(subscription_id) => { - if let Some(cancel) = subscriptions.lock().remove(&subscription_id) { - cancel.notify_one(); - } - } - NdjsonLineKind::Watch { id, params } => { - if reject_duplicate(&in_flight_ids, &outbound_tx, id.clone()).await { - continue; - } - let Some(permit) = - acquire_task_slot(&task_slots, &outbound_tx, Some(id.clone())).await - else { - in_flight_ids.lock().remove(&id); - continue; - }; - let task_dispatcher = dispatcher.clone(); - let task_state = state.clone(); - tasks.spawn(async move { - let _permit = permit; - run_watch_subscription(task_dispatcher, id, params, task_state).await; - }); - } - NdjsonLineKind::Logs { id, params } => { - if reject_duplicate(&in_flight_ids, &outbound_tx, id.clone()).await { - continue; - } - let Some(permit) = - acquire_task_slot(&task_slots, &outbound_tx, Some(id.clone())).await - else { - in_flight_ids.lock().remove(&id); - continue; - }; - let task_dispatcher = dispatcher.clone(); - let task_state = state.clone(); - tasks.spawn(async move { - let _permit = permit; - logs::run_logs_subscription(task_dispatcher, id, params, task_state).await; - }); - } - NdjsonLineKind::AgentAttach { id, params } => { - if reject_duplicate(&in_flight_ids, &outbound_tx, id.clone()).await { - continue; - } - let Some(permit) = - acquire_task_slot(&task_slots, &outbound_tx, Some(id.clone())).await - else { - in_flight_ids.lock().remove(&id); - continue; - }; - let task_dispatcher = dispatcher.clone(); - let task_state = state.clone(); - tasks.spawn(async move { - let _permit = permit; - agent_attach::run_agent_attach_subscription( - task_dispatcher, - id, - params, - task_state, - ) - .await; - }); - } - NdjsonLineKind::Passthrough { id } => { - if let Some(id) = id.clone() { - if reject_duplicate(&in_flight_ids, &outbound_tx, id).await { - continue; - } - } - let Some(permit) = acquire_task_slot(&task_slots, &outbound_tx, id.clone()).await - else { - if let Some(id) = id { - in_flight_ids.lock().remove(&id); - } - continue; - }; - let task_dispatcher = dispatcher.clone(); - let task_state = state.clone(); - tasks.spawn(async move { - let _permit = permit; - run_passthrough_request(task_dispatcher, id, line, task_state).await; - }); - } + let kind = classify_ndjson_line(&line); + let mut ctx = DispatchCtx { + dispatcher: &dispatcher, + state: &state, + task_slots: &task_slots, + tasks: &mut tasks, + }; + if let LineDispatch::Passthrough { id, permit } = + dispatch_classified_line(kind, &mut ctx).await + { + let task_dispatcher = dispatcher.clone(); + let task_state = state.clone(); + tasks.spawn(async move { + let _permit = permit; + run_passthrough_request(task_dispatcher, id, line, task_state).await; + }); } } - for cancel in subscriptions.lock().drain().map(|(_, cancel)| cancel) { - cancel.notify_one(); - } - - let drain_naturally = async { while tasks.join_next().await.is_some() {} }; - if tokio::time::timeout(SHUTDOWN_GRACE_PERIOD, drain_naturally) - .await - .is_err() - { - tasks.shutdown().await; - } - drop(subscriptions); - drop(outbound_tx); - drop(state); - let _ = writer_task.await; + shutdown_connection(ShutdownArgs { + subscriptions, + tasks, + outbound_tx, + state, + writer_task, + }) + .await; Ok(()) } @@ -267,8 +192,9 @@ async fn run_passthrough_request( /// `subscription::run_established_subscription` -- the same shared establish/loop/cleanup /// `subscription::run_bounded_event_subscription` uses. The established [`WatchHandle`] is kept /// alive for the duration purely to hold its backing flag false: dropping it early would trip -/// `WatchEventStream`'s own cancellation check before anything ever streams. -async fn run_watch_subscription( +/// `WatchEventStream`'s own cancellation check before anything ever streams. Reused verbatim by +/// the sibling `websocket` transport module. +pub(crate) async fn run_watch_subscription( dispatcher: Dispatcher, id: RequestId, params: Value, diff --git a/crates/openengine-cluster-server/src/stdio/admission.rs b/crates/openengine-cluster-server/src/stdio/admission.rs index 41834901..248b3f54 100644 --- a/crates/openengine-cluster-server/src/stdio/admission.rs +++ b/crates/openengine-cluster-server/src/stdio/admission.rs @@ -10,15 +10,16 @@ use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore}; use super::serialize_error; -/// Maximum concurrent backend dispatch/subscription tasks per connection. -pub(super) const MAX_CONNECTION_TASKS: usize = 256; +/// Maximum concurrent backend dispatch/subscription tasks per connection. Shared verbatim with +/// the sibling `websocket` transport module. +pub(crate) const MAX_CONNECTION_TASKS: usize = 256; const DUPLICATE_REQUEST_ID: &str = "DUPLICATE_REQUEST_ID"; const SERVER_BUSY: &str = "SERVER_BUSY"; -pub(super) type InFlightIds = Arc>>; +pub(crate) type InFlightIds = Arc>>; /// Registers `id` as in flight, or reports that it was already in flight. -pub(super) async fn reject_duplicate( +pub(crate) async fn reject_duplicate( in_flight_ids: &InFlightIds, outbound_tx: &mpsc::Sender, id: RequestId, @@ -40,7 +41,7 @@ pub(super) async fn reject_duplicate( /// Acquires one task slot. Excess requests receive a deterministic application error; /// notifications have no JSON-RPC response and are dropped. -pub(super) async fn acquire_task_slot( +pub(crate) async fn acquire_task_slot( task_slots: &Arc, outbound_tx: &mpsc::Sender, id: Option, diff --git a/crates/openengine-cluster-server/src/stdio/agent_attach.rs b/crates/openengine-cluster-server/src/stdio/agent_attach.rs index 70661fe3..0f4fd009 100644 --- a/crates/openengine-cluster-server/src/stdio/agent_attach.rs +++ b/crates/openengine-cluster-server/src/stdio/agent_attach.rs @@ -19,7 +19,8 @@ use crate::{serialize_backend_error, serialize_error, serialize_success, Cluster /// `subscription/closed` notifications until the stream ends (overflow, backend close, or /// cancellation). See `run_logs_subscription`/`run_watch_subscription` for the identical /// registration-ordering and cancellation-race notes [`run_bounded_event_subscription`] mirrors. -pub(super) async fn run_agent_attach_subscription( +/// Reused verbatim by the sibling `websocket` transport module. +pub(crate) async fn run_agent_attach_subscription( dispatcher: Dispatcher, id: RequestId, params: Value, diff --git a/crates/openengine-cluster-server/src/stdio/dispatch.rs b/crates/openengine-cluster-server/src/stdio/dispatch.rs new file mode 100644 index 00000000..f7d2789c --- /dev/null +++ b/crates/openengine-cluster-server/src/stdio/dispatch.rs @@ -0,0 +1,195 @@ +//! Shared per-connection dispatch/setup/shutdown machinery driving both `serve_ndjson`'s and the +//! sibling `websocket` transport module's `serve_websocket`'s classify-then-spawn loop, so +//! results, events, and errors stay byte-equivalent between the stdio and WebSocket bindings. + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use std::time::Duration; + +use openengine_cluster_protocol::RequestId; +use parking_lot::Mutex; +use serde_json::Value; +use tokio::sync::{mpsc, OwnedSemaphorePermit, Semaphore}; +use tokio::task::{JoinHandle, JoinSet}; + +use super::admission::{acquire_task_slot, reject_duplicate, InFlightIds, MAX_CONNECTION_TASKS}; +use super::{ + agent_attach, logs, run_watch_subscription, ConnectionState, NdjsonLineKind, SubscriptionMap, +}; +use crate::{ClusterBackend, Dispatcher}; + +/// Grace period given to already-spawned bounded backend dispatches to finish once the connection +/// closes. Subscription tasks are notified through their cancellation handles before shutdown; +/// any backend operation that does not finish inside this bound is force-aborted. +const SHUTDOWN_GRACE_PERIOD: Duration = Duration::from_millis(200); + +/// Result of [`dispatch_classified_line`]'s shared subscription/cancel/admission handling for one +/// classified line, identical between `serve_ndjson` and the sibling `websocket` transport +/// module's `serve_websocket`. `Passthrough` is handed back to the caller -- after its own +/// duplicate-id rejection and task-slot acquisition have already run here -- only because each +/// binding's actual passthrough dispatch differs: the WebSocket binding additionally tracks a +/// best-effort `$/cancelRequest` registry stdio has no notion of. +pub(crate) enum LineDispatch { + Handled, + Passthrough { + id: Option, + permit: OwnedSemaphorePermit, + }, +} + +/// Bundles the four per-connection handles [`dispatch_classified_line`] and its helpers thread +/// through, so callers -- `serve_ndjson` here and the sibling `websocket` transport module's +/// `serve_websocket` -- pass one context value instead of an ever-growing parameter list. +pub(crate) struct DispatchCtx<'a, B> { + pub(crate) dispatcher: &'a Dispatcher, + pub(crate) state: &'a ConnectionState, + pub(crate) task_slots: &'a Arc, + pub(crate) tasks: &'a mut JoinSet<()>, +} + +/// Spawns `run` as a bounded, duplicate-id-rejecting, admission-controlled subscription task for +/// `id`/`params` -- shared by `dispatch_classified_line`'s `Watch`/`Logs`/`AgentAttach` arms, which +/// otherwise differ only in which subscription runner they spawn. +async fn spawn_subscription_task( + ctx: &mut DispatchCtx<'_, B>, + id: RequestId, + params: Value, + run: F, +) where + B: ClusterBackend, + F: FnOnce(Dispatcher, RequestId, Value, ConnectionState) -> Fut + Send + 'static, + Fut: std::future::Future + Send + 'static, +{ + if reject_duplicate(&ctx.state.in_flight_ids, &ctx.state.outbound_tx, id.clone()).await { + return; + } + let Some(permit) = + acquire_task_slot(ctx.task_slots, &ctx.state.outbound_tx, Some(id.clone())).await + else { + ctx.state.in_flight_ids.lock().remove(&id); + return; + }; + let task_dispatcher = ctx.dispatcher.clone(); + let task_state = ctx.state.clone(); + ctx.tasks.spawn(async move { + let _permit = permit; + run(task_dispatcher, id, params, task_state).await; + }); +} + +/// Handles everything shared between `serve_ndjson` and `serve_websocket` for one classified line: +/// `subscription/cancel` notifications and `watch`/`logs`/`agent/attach` establishment, including +/// their admission control (duplicate-id rejection, task-slot acquisition). Only +/// [`NdjsonLineKind::Passthrough`] is handed back to the caller as [`LineDispatch::Passthrough`], +/// with its own admission already applied, so each binding can run its own passthrough dispatch. +pub(crate) async fn dispatch_classified_line( + kind: NdjsonLineKind, + ctx: &mut DispatchCtx<'_, B>, +) -> LineDispatch +where + B: ClusterBackend, +{ + match kind { + NdjsonLineKind::Cancel(subscription_id) => { + if let Some(cancel) = ctx.state.subscriptions.lock().remove(&subscription_id) { + cancel.notify_one(); + } + LineDispatch::Handled + } + NdjsonLineKind::Watch { id, params } => { + spawn_subscription_task(ctx, id, params, run_watch_subscription).await; + LineDispatch::Handled + } + NdjsonLineKind::Logs { id, params } => { + spawn_subscription_task(ctx, id, params, logs::run_logs_subscription).await; + LineDispatch::Handled + } + NdjsonLineKind::AgentAttach { id, params } => { + spawn_subscription_task(ctx, id, params, agent_attach::run_agent_attach_subscription) + .await; + LineDispatch::Handled + } + NdjsonLineKind::Passthrough { id } => { + if let Some(id) = id.clone() { + if reject_duplicate(&ctx.state.in_flight_ids, &ctx.state.outbound_tx, id).await { + return LineDispatch::Handled; + } + } + let Some(permit) = + acquire_task_slot(ctx.task_slots, &ctx.state.outbound_tx, id.clone()).await + else { + if let Some(id) = id { + ctx.state.in_flight_ids.lock().remove(&id); + } + return LineDispatch::Handled; + }; + LineDispatch::Passthrough { id, permit } + } + } +} + +/// The per-connection tracking state shared by every spawned task, freshly constructed identically +/// by `serve_ndjson` and the sibling `websocket` transport module's `serve_websocket` -- everything +/// except the outbound queue itself (each binding constructs that differently: a plain +/// `mpsc::channel` here vs. one paired with a close-signal channel there). +pub(crate) struct ConnectionSetup { + pub(crate) subscriptions: SubscriptionMap, + pub(crate) task_slots: Arc, + pub(crate) tasks: JoinSet<()>, + pub(crate) state: ConnectionState, +} + +pub(crate) fn new_connection_setup(outbound_tx: &mpsc::Sender) -> ConnectionSetup { + let subscriptions: SubscriptionMap = Arc::new(Mutex::new(HashMap::new())); + let in_flight_ids: InFlightIds = Arc::new(Mutex::new(HashSet::new())); + let task_slots = Arc::new(Semaphore::new(MAX_CONNECTION_TASKS)); + let state = ConnectionState { + outbound_tx: outbound_tx.clone(), + subscriptions: Arc::clone(&subscriptions), + in_flight_ids, + }; + ConnectionSetup { + subscriptions, + task_slots, + tasks: JoinSet::new(), + state, + } +} + +/// Grouped arguments for [`shutdown_connection`], keeping that function's argument count +/// reasonable. +pub(crate) struct ShutdownArgs { + pub(crate) subscriptions: SubscriptionMap, + pub(crate) tasks: JoinSet<()>, + pub(crate) outbound_tx: mpsc::Sender, + pub(crate) state: ConnectionState, + pub(crate) writer_task: JoinHandle<()>, +} + +/// Tears down one connection identically for `serve_ndjson` and `serve_websocket`: wakes every +/// live subscription's cancellation, gives already-spawned tasks [`SHUTDOWN_GRACE_PERIOD`] to +/// finish naturally before force-aborting the rest, then drops every sender so the writer task +/// drains and exits. +pub(crate) async fn shutdown_connection(args: ShutdownArgs) { + let ShutdownArgs { + subscriptions, + mut tasks, + outbound_tx, + state, + writer_task, + } = args; + for cancel in subscriptions.lock().drain().map(|(_, cancel)| cancel) { + cancel.notify_one(); + } + let drain_naturally = async { while tasks.join_next().await.is_some() {} }; + if tokio::time::timeout(SHUTDOWN_GRACE_PERIOD, drain_naturally) + .await + .is_err() + { + tasks.shutdown().await; + } + drop(subscriptions); + drop(outbound_tx); + drop(state); + let _ = writer_task.await; +} diff --git a/crates/openengine-cluster-server/src/stdio/logs.rs b/crates/openengine-cluster-server/src/stdio/logs.rs index 9cb3f943..aafdb354 100644 --- a/crates/openengine-cluster-server/src/stdio/logs.rs +++ b/crates/openengine-cluster-server/src/stdio/logs.rs @@ -17,8 +17,9 @@ use crate::{serialize_backend_error, serialize_error, serialize_success, Cluster /// Establishes a `logs` subscription and, on success, streams its `event`/`subscription/closed` /// notifications until the stream ends (overflow, backend close, or cancellation). See /// `run_watch_subscription` for the identical registration-ordering and cancellation-race notes -/// [`run_bounded_event_subscription`] mirrors. -pub(super) async fn run_logs_subscription( +/// [`run_bounded_event_subscription`] mirrors. Reused verbatim by the sibling `websocket` +/// transport module. +pub(crate) async fn run_logs_subscription( dispatcher: Dispatcher, id: RequestId, params: Value, diff --git a/crates/openengine-cluster-server/src/stdio/tests.rs b/crates/openengine-cluster-server/src/stdio/tests.rs index 6e5c1ed1..6be99032 100644 --- a/crates/openengine-cluster-server/src/stdio/tests.rs +++ b/crates/openengine-cluster-server/src/stdio/tests.rs @@ -1,3 +1,5 @@ +use std::collections::HashSet; + use openengine_cluster_protocol::RunId; use super::*; diff --git a/crates/openengine-cluster-server/src/watch/fixtures.rs b/crates/openengine-cluster-server/src/watch/fixtures.rs index 165524f4..61e0803e 100644 --- a/crates/openengine-cluster-server/src/watch/fixtures.rs +++ b/crates/openengine-cluster-server/src/watch/fixtures.rs @@ -2,7 +2,9 @@ //! contract independent of `openengine-cluster-testkit`'s production-shaped //! `InMemoryAdmissionStore`. Shared by this crate's own integration tests and by //! `openengine-cluster-client`'s reconnect tests so both exercise identical -//! subscribe/replay/overflow semantics against one implementation. +//! subscribe/replay/overflow semantics against one implementation. Also bundles [`spawn_ndjson`]/ +//! [`spawn_websocket`] connection harnesses reused across this crate, the client crate, and the +//! testkit crate's own wire-transport tests. use std::io; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; @@ -16,6 +18,7 @@ use openengine_cluster_protocol::{ use tokio::io::DuplexStream; use tokio::sync::{mpsc, Mutex}; use tokio::task::JoinHandle; +use tokio_tungstenite::WebSocketStream; use super::{ subscribe_and_stream, ObservationStore, PublicEventRecord, ReplayPageRequest, @@ -23,6 +26,7 @@ use super::{ }; use crate::admission::StoreError; use crate::stdio::serve_ndjson; +use crate::websocket::{serve_websocket, websocket_config}; use crate::{BackendError, ClusterBackend, ConnectionContext, Dispatcher}; /// Wires `backend` to a fresh [`serve_ndjson`] task over an in-memory duplex pipe pair, returning @@ -56,6 +60,45 @@ pub async fn await_ndjson_shutdown(server: JoinHandle>) { .expect("serve_ndjson task returned an error"); } +/// Wires `backend` to a fresh [`serve_websocket`] task over an in-memory duplex pipe pair, +/// returning the pipe's already-handshaken client-facing [`WebSocketStream`] and the server task's +/// join handle. Shared by this crate's own WebSocket framing tests and by +/// `openengine-cluster-client`'s and `openengine-cluster-testkit`'s WebSocket transport tests so +/// all three drive the exact same wiring, mirroring [`spawn_ndjson`]'s role for the NDJSON +/// binding. +pub async fn spawn_websocket( + backend: B, +) -> (WebSocketStream, JoinHandle>) +where + B: ClusterBackend, +{ + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let dispatcher = Dispatcher::new(backend, ConnectionContext::default()); + let server = tokio::spawn(async move { + let ws = tokio_tungstenite::accept_async_with_config(server_io, Some(websocket_config())) + .await + .expect("server handshake must succeed"); + serve_websocket(dispatcher, ws).await + }); + let (client, _response) = + tokio_tungstenite::client_async("ws://localhost/websocket", client_io) + .await + .expect("client handshake must succeed"); + (client, server) +} + +/// Awaits a [`spawn_websocket`] server task's join handle within a short bound, panicking with a +/// descriptive message if it hangs, panicked, or returned an error. Shared by this crate's own and +/// `openengine-cluster-client`'s/`openengine-cluster-testkit`'s WebSocket tests for asserting +/// deterministic shutdown, mirroring [`await_ndjson_shutdown`]. +pub async fn await_websocket_shutdown(server: JoinHandle>) { + tokio::time::timeout(std::time::Duration::from_secs(2), server) + .await + .expect("serve_websocket task did not terminate within the timeout") + .expect("serve_websocket task panicked") + .expect("serve_websocket task returned an error"); +} + #[derive(Default)] struct FixtureState { history: Vec, diff --git a/crates/openengine-cluster-server/src/websocket.rs b/crates/openengine-cluster-server/src/websocket.rs new file mode 100644 index 00000000..6c7bea80 --- /dev/null +++ b/crates/openengine-cluster-server/src/websocket.rs @@ -0,0 +1,333 @@ +//! Production WebSocket transport binding the backend-neutral [`Dispatcher`] and generic +//! subscription framing to the wire. One JSON-RPC object per text message; this module reuses +//! `stdio::serve_ndjson`'s exact classification, admission, and subscription-streaming machinery +//! ([`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. + +use std::collections::HashMap; +use std::io; +use std::sync::Arc; + +use futures_util::stream::SplitStream; +use futures_util::{Sink, SinkExt, StreamExt}; +use openengine_cluster_protocol::{CancelRequestParams, JsonRpcNotification, RequestId}; +use parking_lot::Mutex; +use tokio::io::{AsyncRead, AsyncWrite}; +use tokio::sync::{mpsc, oneshot, OwnedSemaphorePermit}; +use tokio::task::{AbortHandle, JoinSet}; +use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; +use tokio_tungstenite::tungstenite::protocol::{CloseFrame, WebSocketConfig}; +use tokio_tungstenite::tungstenite::{Error as WsError, Message, Utf8Bytes}; +use tokio_tungstenite::WebSocketStream; + +use crate::stdio::{ + classify_ndjson_line, dispatch_classified_line, new_connection_setup, shutdown_connection, + ConnectionSetup, ConnectionState, DispatchCtx, LineDispatch, ShutdownArgs, +}; +use crate::{ClusterBackend, Dispatcher}; + +/// Bounded WebSocket text-frame length: a text message whose UTF-8 byte length exceeds this +/// closes the connection with code 1009 (message too big), matching `stdio::serve_ndjson`'s +/// NDJSON line bound. +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 +/// 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>>; + +/// WebSocket configuration enforcing [`MAX_FRAME_BYTES`] as the incoming message size bound, so +/// tungstenite itself rejects an oversized message during frame reassembly rather than buffering +/// it without limit ahead of this binding's own explicit length check. +#[must_use] +pub fn websocket_config() -> WebSocketConfig { + WebSocketConfig::default().max_message_size(Some(MAX_FRAME_BYTES)) +} + +/// 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; +/// 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( + dispatcher: Dispatcher, + ws: WebSocketStream, +) -> io::Result<()> +where + B: ClusterBackend, + S: AsyncRead + AsyncWrite + Unpin + Send + 'static, +{ + let (sink, mut stream) = ws.split(); + let (outbound_tx, outbound_rx) = mpsc::channel::(OUTBOUND_QUEUE_CAPACITY); + let (close_tx, close_rx) = oneshot::channel::(); + let writer_task = tokio::spawn(run_writer(sink, outbound_rx, close_rx)); + + let ConnectionSetup { + subscriptions, + task_slots, + mut tasks, + state, + } = new_connection_setup(&outbound_tx); + let cancel_registry: CancelRegistry = Arc::new(Mutex::new(HashMap::new())); + let mut close_tx = Some(close_tx); + + while let Some(message) = next_incoming_message(&mut tasks, &mut stream, &mut close_tx).await { + 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, + }; + if matches!(handle_message(message, &mut ctx).await, FrameOutcome::Break) { + break; + } + } + + shutdown_connection(ShutdownArgs { + subscriptions, + tasks, + outbound_tx, + state, + writer_task, + }) + .await; + Ok(()) +} + +/// Awaits the next inbound message, reaping completed request tasks in the meantime exactly like +/// `stdio::serve_ndjson` -- a concurrency cap alone is insufficient since `JoinSet` retains every +/// completed output until it is joined. Returns `None` once the stream ends or a transport-level +/// capacity error forces a close (signalled via `close_tx` beforehand, per [`signal_close`]). +async fn next_incoming_message( + tasks: &mut JoinSet<()>, + stream: &mut SplitStream>, + close_tx: &mut Option>, +) -> Option +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let next_message = loop { + tokio::select! { + completed = tasks.join_next(), if !tasks.is_empty() => { + let _ = completed; + } + message = stream.next() => break message, + } + }; + match next_message { + Some(Ok(message)) => Some(message), + Some(Err(error)) => { + if matches!(error, WsError::Capacity(_)) { + signal_close(close_tx, CloseCode::Size, "message too big"); + } + None + } + None => None, + } +} + +/// Whether the connection loop should keep reading after handling one frame. +enum FrameOutcome { + Continue, + Break, +} + +/// 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 +/// signal), bundled so [`handle_message`] and its helpers take one argument instead of an +/// ever-growing list. +struct WsCtx<'a, B> { + dispatch: DispatchCtx<'a, B>, + cancel_registry: &'a CancelRegistry, + close_tx: &'a mut Option>, +} + +/// Dispatches one inbound frame by kind. Binary frames and a peer-initiated close both end the +/// connection; ping/pong/raw frames are ignored (tungstenite auto-answers pings); text frames are +/// handled by [`handle_text_frame`]. +async fn handle_message(message: Message, ctx: &mut WsCtx<'_, B>) -> FrameOutcome +where + B: ClusterBackend, +{ + match message { + Message::Text(text) => handle_text_frame(text, ctx).await, + Message::Binary(_) => { + signal_close( + ctx.close_tx, + CloseCode::Unsupported, + "binary frames are not supported", + ); + FrameOutcome::Break + } + Message::Ping(_) | Message::Pong(_) | Message::Frame(_) => FrameOutcome::Continue, + Message::Close(_) => FrameOutcome::Break, + } +} + +/// Handles one `Message::Text` frame: enforces [`MAX_FRAME_BYTES`], routes a `$/cancelRequest` +/// notification inline, and otherwise classifies and dispatches the line exactly like +/// `stdio::serve_ndjson` via [`dispatch_classified_line`] -- spawning this binding's own +/// passthrough task (with `cancel_registry` tracking) for a [`LineDispatch::Passthrough`] result. +async fn handle_text_frame(text: Utf8Bytes, ctx: &mut WsCtx<'_, B>) -> FrameOutcome +where + B: ClusterBackend, +{ + if text.len() > MAX_FRAME_BYTES { + signal_close(ctx.close_tx, CloseCode::Size, "message too big"); + return FrameOutcome::Break; + } + if let Some(cancel_id) = parse_cancel_request(&text) { + if let Some(handle) = ctx.cancel_registry.lock().remove(&cancel_id) { + handle.abort(); + } + return FrameOutcome::Continue; + } + let kind = classify_ndjson_line(&text); + if let LineDispatch::Passthrough { id, permit } = + dispatch_classified_line(kind, &mut ctx.dispatch).await + { + spawn_passthrough(ctx, id, permit, text.as_str().to_owned()); + } + 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`. +fn spawn_passthrough( + ctx: &mut WsCtx<'_, B>, + id: Option, + permit: OwnedSemaphorePermit, + line: String, +) where + B: ClusterBackend, +{ + 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 abort_handle = ctx.dispatch.tasks.spawn(async move { + let _permit = permit; + run_passthrough_request(PassthroughRequest { + dispatcher: task_dispatcher, + id, + line, + state: task_state, + cancel_registry: task_cancel_registry, + }) + .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); + } +} + +/// Grouped arguments for [`run_passthrough_request`], keeping that function's argument count +/// reasonable. +struct PassthroughRequest { + dispatcher: Dispatcher, + id: Option, + line: String, + state: ConnectionState, + cancel_registry: CancelRegistry, +} + +/// 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. +async fn run_passthrough_request(request: PassthroughRequest) +where + B: ClusterBackend, +{ + let PassthroughRequest { + dispatcher, + id, + line, + state, + cancel_registry, + } = request; + let response = dispatcher.dispatch(&line).await; + if let Some(id) = &id { + state.in_flight_ids.lock().remove(id); + cancel_registry.lock().remove(id); + } + let _ = state.outbound_tx.send(response).await; +} + +/// Parses `text` as a `$/cancelRequest` notification, returning the target `RequestId` only when +/// both the JSON-RPC method matches and the body deserializes -- anything else (including every +/// other recognized or unrecognized method) is left for `classify_ndjson_line` to route. +fn parse_cancel_request(text: &str) -> Option { + let notification: JsonRpcNotification = serde_json::from_str(text).ok()?; + (notification.method == "$/cancelRequest").then_some(notification.params.id) +} + +/// Sends `code`/`reason` on `close_tx` exactly once -- a no-op if a close was already signalled -- +/// so the writer task emits one deterministic `Message::Close` and stops. +fn signal_close( + close_tx: &mut Option>, + code: CloseCode, + reason: &'static str, +) { + if let Some(tx) = close_tx.take() { + let _ = tx.send(CloseFrame { + code, + reason: Utf8Bytes::from_static(reason), + }); + } +} + +/// Drains the bounded outbound queue, writing each line as a `Message::Text` frame, until the +/// peer closes, every sender is dropped, or `close_tx` (see [`signal_close`]) fires -- in which +/// case one `Message::Close` is sent and the writer stops immediately without draining further +/// queued lines, mirroring `stdio::run_writer`'s "drain until torn down" shape plus the close +/// frame WebSocket framing requires that NDJSON has no equivalent of. +async fn run_writer( + mut sink: Si, + mut outbound_rx: mpsc::Receiver, + mut close_rx: oneshot::Receiver, +) where + Si: Sink + Unpin, +{ + loop { + // Biased: once a close is signalled it must win even if `outbound_tx` is dropped in the + // same tick (the caller drops it immediately after signalling close during shutdown) -- + // otherwise an unbiased `select!` picks between two simultaneously ready branches at + // random, sometimes discarding this deterministic close code for `sink.close()`'s generic + // code-less fallback below. + tokio::select! { + biased; + close = &mut close_rx => { + if let Ok(frame) = close { + let _ = sink.send(Message::Close(Some(frame))).await; + } + break; + } + line = outbound_rx.recv() => { + match line { + Some(line) => { + if sink.send(Message::text(line)).await.is_err() { + break; + } + } + None => break, + } + } + } + } + let _ = sink.close().await; +} diff --git a/crates/openengine-cluster-server/tests/admission_bound_support/mod.rs b/crates/openengine-cluster-server/tests/admission_bound_support/mod.rs new file mode 100644 index 00000000..d1ff4a48 --- /dev/null +++ b/crates/openengine-cluster-server/tests/admission_bound_support/mod.rs @@ -0,0 +1,84 @@ +//! Transport-generic duplicate-in-flight-id and bounded-task-admission scenarios, used identically +//! by `tests/subscription_ndjson.rs` (over NDJSON) and `tests/websocket.rs` (over WebSocket) since +//! both drive the exact same admission/dedup behavior against the shared `Dispatcher`/ +//! connection-task machinery, independent of wire framing. + +use std::sync::Arc; +use std::time::Duration; + +use openengine_cluster_protocol::RunId; +use openengine_cluster_server::watch::fixtures::{FixtureBackend, FixtureStore}; +use serde_json::Value; +use tokio::sync::Notify; + +use crate::gated_backend_support::GatedBackend; + +/// Minimal capability both wire harnesses share: writing a `get` request with a given id and +/// reading the next decoded JSON-RPC frame, regardless of wire framing (NDJSON lines vs WebSocket +/// text frames). +pub trait RequestChannel { + async fn send_get(&mut self, id: i64); + async fn recv_value(&mut self) -> Value; +} + +/// Spawn glue letting [`spawn_gated_harness`] construct any wire harness `H` against a fresh +/// gated backend without each binding needing its own copy of that setup. +pub trait GatedHarnessSpawn: Sized { + async fn spawn_gated(backend: GatedBackend) -> Self; +} + +/// Constructs a fresh `FixtureStore`/[`GatedBackend`] pair (gating only `get`) and spawns `H` +/// against it, returning the harness alongside the gate so a test can release it once its +/// in-flight-request assertions are set up. +pub async fn spawn_gated_harness() -> (H, Arc) { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let gate = Arc::new(Notify::new()); + let harness = H::spawn_gated(GatedBackend { + inner: FixtureBackend::new(store), + gate: Arc::clone(&gate), + }) + .await; + (harness, gate) +} + +/// Sends two `get` requests sharing request id `1` while `gate` blocks the first from completing, +/// asserts the second is rejected as a synchronous `DUPLICATE_REQUEST_ID` error (the first request +/// is still blocked on the gate, so the only frame that can possibly exist yet is the duplicate +/// rejection for the second), then releases the gate and asserts the first request completes +/// normally. +pub async fn assert_duplicate_in_flight_ids_are_rejected( + harness: &mut H, + gate: &Notify, +) { + harness.send_get(1).await; + harness.send_get(1).await; + + let duplicate = harness.recv_value().await; + assert_eq!(duplicate["id"], 1); + assert_eq!(duplicate["error"]["code"], -32600); + assert_eq!(duplicate["error"]["data"]["code"], "DUPLICATE_REQUEST_ID"); + + gate.notify_one(); + let first = harness.recv_value().await; + assert_eq!(first["id"], 1); + assert!(first.get("result").is_some(), "{first}"); +} + +/// Sends `max_connection_tasks + 1` distinct-id `get` requests, all blocked on `harness`'s gated +/// backend, and asserts the request past the bound is rejected with a synchronous `SERVER_BUSY` +/// error that does not wait for any of the blocked backend calls to complete. +pub async fn assert_excess_requests_are_rejected_with_server_busy( + harness: &mut H, + max_connection_tasks: i64, +) { + for id in 1..=max_connection_tasks + 1 { + harness.send_get(id).await; + } + + let rejected = tokio::time::timeout(Duration::from_secs(1), harness.recv_value()) + .await + .expect("the bounded admission rejection must not wait for blocked backend calls"); + assert_eq!(rejected["id"], max_connection_tasks + 1); + assert_eq!(rejected["error"]["code"], -32000); + assert_eq!(rejected["error"]["data"]["code"], "SERVER_BUSY"); +} diff --git a/crates/openengine-cluster-server/tests/gated_backend_support/mod.rs b/crates/openengine-cluster-server/tests/gated_backend_support/mod.rs new file mode 100644 index 00000000..ff130b62 --- /dev/null +++ b/crates/openengine-cluster-server/tests/gated_backend_support/mod.rs @@ -0,0 +1,50 @@ +//! Shared `ClusterBackend` wrapper that gates only `get` behind an explicit [`Notify`], making +//! duplicate-in-flight-request-id and bounded-task-admission cases deterministic instead of +//! racing the backend call. Used identically by `tests/subscription_ndjson.rs` (over NDJSON) and +//! `tests/websocket.rs` (over WebSocket) since both drive the exact same admission/dedup behavior +//! against the shared `Dispatcher`/connection-task machinery, independent of wire framing. + +use std::sync::Arc; + +use async_trait::async_trait; +use openengine_cluster_protocol::{ + GetParams, GetResult, InitializeParams, InitializeResult, WatchParams, WatchResult, +}; +use openengine_cluster_server::watch::fixtures::FixtureBackend; +use openengine_cluster_server::watch::{WatchEventStream, WatchHandle}; +use openengine_cluster_server::{BackendError, ClusterBackend, ConnectionContext}; +use tokio::sync::Notify; + +pub struct GatedBackend { + pub inner: FixtureBackend, + pub gate: Arc, +} + +#[async_trait] +impl ClusterBackend for GatedBackend { + async fn initialize( + &self, + context: &ConnectionContext, + params: InitializeParams, + ) -> Result { + self.inner.initialize(context, params).await + } + + async fn get( + &self, + context: &ConnectionContext, + params: GetParams, + ) -> Result { + self.gate.notified().await; + self.inner.get(context, params).await + } + + async fn watch( + &self, + context: &ConnectionContext, + params: WatchParams, + queue_capacity: usize, + ) -> Result<(WatchResult, WatchEventStream, WatchHandle), BackendError> { + self.inner.watch(context, params, queue_capacity).await + } +} diff --git a/crates/openengine-cluster-server/tests/subscription_ndjson.rs b/crates/openengine-cluster-server/tests/subscription_ndjson.rs index 710e89c6..5b8509bb 100644 --- a/crates/openengine-cluster-server/tests/subscription_ndjson.rs +++ b/crates/openengine-cluster-server/tests/subscription_ndjson.rs @@ -8,25 +8,31 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::Duration; -use async_trait::async_trait; -use openengine_cluster_protocol::{ - GetParams, GetResult, InitializeParams, InitializeResult, RunId, WatchEvent, WatchParams, - WatchResult, -}; +use openengine_cluster_protocol::{RunId, WatchEvent}; use openengine_cluster_server::watch::fixtures::{ await_ndjson_shutdown, spawn_ndjson, FixtureBackend, FixtureStore, }; -use openengine_cluster_server::watch::{WatchEventStream, WatchHandle}; -use openengine_cluster_server::{BackendError, ClusterBackend, ConnectionContext}; +use openengine_cluster_server::ClusterBackend; use serde_json::{json, Value}; use tokio::io::{AsyncWriteExt, BufReader, DuplexStream}; -use tokio::sync::Notify; use tokio::task::JoinHandle; #[path = "ndjson_test_support/mod.rs"] mod ndjson_test_support; use ndjson_test_support::{read_line, read_value, request_line, write_line}; +#[path = "gated_backend_support/mod.rs"] +mod gated_backend_support; +use gated_backend_support::GatedBackend; + +#[path = "admission_bound_support/mod.rs"] +mod admission_bound_support; +use admission_bound_support::{ + assert_duplicate_in_flight_ids_are_rejected, + assert_excess_requests_are_rejected_with_server_busy, spawn_gated_harness, GatedHarnessSpawn, + RequestChannel, +}; + /// Matches the issue's documented "> 1 MiB" oversized-frame threshold; the exact bound is an /// internal `serve_ndjson` implementation detail, not part of the public contract. const OVERSIZED_LINE_BYTES: usize = 1024 * 1024 + 16; @@ -49,6 +55,22 @@ where } } +impl RequestChannel for Harness { + async fn send_get(&mut self, id: i64) { + write_line(&mut self.write, &request_line(id, "get", json!({}))).await; + } + + async fn recv_value(&mut self) -> Value { + read_value(&mut self.read).await + } +} + +impl GatedHarnessSpawn for Harness { + async fn spawn_gated(backend: GatedBackend) -> Self { + spawn_server(backend) + } +} + fn cancel_line(subscription_id: &str) -> String { json!({ "jsonrpc": "2.0", @@ -130,65 +152,11 @@ async fn oversized_and_malformed_frames_are_deterministic() { shut_down(harness).await; } -/// Wraps [`FixtureBackend`] so a test can hold `get` in flight until it explicitly releases it, -/// making duplicate-in-flight-id detection deterministic instead of racing the backend call. -struct GatedBackend { - inner: FixtureBackend, - gate: Arc, -} - -#[async_trait] -impl ClusterBackend for GatedBackend { - async fn initialize( - &self, - context: &ConnectionContext, - params: InitializeParams, - ) -> Result { - self.inner.initialize(context, params).await - } - - async fn get( - &self, - context: &ConnectionContext, - params: GetParams, - ) -> Result { - self.gate.notified().await; - self.inner.get(context, params).await - } - - async fn watch( - &self, - context: &ConnectionContext, - params: WatchParams, - queue_capacity: usize, - ) -> Result<(WatchResult, WatchEventStream, WatchHandle), BackendError> { - self.inner.watch(context, params, queue_capacity).await - } -} - #[tokio::test] async fn duplicate_request_ids_are_rejected() { - let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); - let gate = Arc::new(Notify::new()); - let mut harness = spawn_server(GatedBackend { - inner: FixtureBackend::new(store), - gate: Arc::clone(&gate), - }); - - write_line(&mut harness.write, &request_line(1, "get", json!({}))).await; - write_line(&mut harness.write, &request_line(1, "get", json!({}))).await; - - // The first request is still blocked on the gate, so the only frame that can possibly exist - // yet is the duplicate rejection for the second. - let duplicate = read_value(&mut harness.read).await; - assert_eq!(duplicate["id"], 1); - assert_eq!(duplicate["error"]["code"], -32600); - assert_eq!(duplicate["error"]["data"]["code"], "DUPLICATE_REQUEST_ID"); - - gate.notify_one(); - let first = read_value(&mut harness.read).await; - assert_eq!(first["id"], 1); - assert!(first.get("result").is_some(), "{first}"); + let (mut harness, gate) = spawn_gated_harness::().await; + + assert_duplicate_in_flight_ids_are_rejected(&mut harness, &gate).await; shut_down(harness).await; } @@ -197,23 +165,9 @@ async fn duplicate_request_ids_are_rejected() { async fn excess_requests_are_rejected_without_unbounded_task_admission() { const MAX_CONNECTION_TASKS: i64 = 256; - let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); - let gate = Arc::new(Notify::new()); - let mut harness = spawn_server(GatedBackend { - inner: FixtureBackend::new(store), - gate, - }); - - for id in 1..=MAX_CONNECTION_TASKS + 1 { - write_line(&mut harness.write, &request_line(id, "get", json!({}))).await; - } + let (mut harness, _gate) = spawn_gated_harness::().await; - let rejected = tokio::time::timeout(Duration::from_secs(1), read_value(&mut harness.read)) - .await - .expect("the bounded admission rejection must not wait for blocked backend calls"); - assert_eq!(rejected["id"], MAX_CONNECTION_TASKS + 1); - assert_eq!(rejected["error"]["code"], -32000); - assert_eq!(rejected["error"]["data"]["code"], "SERVER_BUSY"); + assert_excess_requests_are_rejected_with_server_busy(&mut harness, MAX_CONNECTION_TASKS).await; shut_down(harness).await; } diff --git a/crates/openengine-cluster-server/tests/websocket.rs b/crates/openengine-cluster-server/tests/websocket.rs new file mode 100644 index 00000000..173c19e6 --- /dev/null +++ b/crates/openengine-cluster-server/tests/websocket.rs @@ -0,0 +1,235 @@ +//! WebSocket transport end-to-end coverage: unary dispatch byte-equivalence with the in-process +//! `Dispatcher`, bounded-frame framing (binary rejection, oversize close, malformed-JSON +//! tolerance), `$/cancelRequest` no-op semantics, duplicate in-flight request ids, and bounded +//! task admission -- the WebSocket counterparts of `tests/subscription_ndjson.rs`'s NDJSON +//! coverage, driven over raw tungstenite frames instead of NDJSON lines so binary/oversize/ +//! malformed framing can be crafted directly. + +use std::sync::Arc; +use std::time::Duration; + +use futures_util::{SinkExt, StreamExt}; +use openengine_cluster_protocol::RunId; +use openengine_cluster_server::watch::fixtures::{ + await_websocket_shutdown, spawn_websocket, FixtureBackend, FixtureStore, +}; +use openengine_cluster_server::{ClusterBackend, ConnectionContext, Dispatcher}; +use serde_json::{json, Value}; +use tokio::io::DuplexStream; +use tokio::task::JoinHandle; +use tokio_tungstenite::tungstenite::protocol::frame::coding::CloseCode; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::WebSocketStream; + +#[path = "gated_backend_support/mod.rs"] +mod gated_backend_support; +use gated_backend_support::GatedBackend; + +#[path = "admission_bound_support/mod.rs"] +mod admission_bound_support; +use admission_bound_support::{ + assert_duplicate_in_flight_ids_are_rejected, + assert_excess_requests_are_rejected_with_server_busy, spawn_gated_harness, GatedHarnessSpawn, + RequestChannel, +}; + +/// Matches `serve_websocket`'s documented `MAX_FRAME_BYTES` bound; hardcoded rather than imported +/// since it is an internal implementation detail, not part of the public contract (mirrors +/// `tests/subscription_ndjson.rs`'s identical `OVERSIZED_LINE_BYTES` convention). +const OVERSIZED_TEXT_BYTES: usize = 1024 * 1024 + 16; +const MAX_CONNECTION_TASKS: i64 = 256; + +struct Harness { + client: WebSocketStream, + server: JoinHandle>, +} + +async fn spawn_server(backend: B) -> Harness +where + B: ClusterBackend, +{ + let (client, server) = spawn_websocket(backend).await; + Harness { client, server } +} + +async fn send_text(client: &mut WebSocketStream, text: impl Into) { + client.send(Message::text(text.into())).await.unwrap(); +} + +async fn recv_json(client: &mut WebSocketStream) -> Value { + match client + .next() + .await + .expect("connection closed unexpectedly while awaiting a frame") + { + Ok(Message::Text(text)) => serde_json::from_str(&text).unwrap(), + Ok(other) => panic!("expected a text frame, got {other:?}"), + Err(error) => panic!("websocket read failed: {error}"), + } +} + +fn request_text(id: i64, method: &str, params: Value) -> String { + json!({"jsonrpc": "2.0", "id": id, "method": method, "params": params}).to_string() +} + +async fn shut_down(harness: Harness) { + let Harness { mut client, server } = harness; + // A real close handshake exercises `Message::Close -> break`, the same deterministic shutdown + // path a well-behaved peer takes; already-closed connections (the binary/oversize tests) just + // ignore the resulting error. + let _ = client.close(None).await; + drop(client); + await_websocket_shutdown(server).await; +} + +#[tokio::test] +async fn unary_initialize_and_get_match_in_process_dispatch() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(Arc::clone(&store))).await; + + let init_request = request_text( + 1, + "initialize", + json!({"protocolVersion": "openengine.cluster/v1"}), + ); + send_text(&mut harness.client, init_request.clone()).await; + let ws_init_response = recv_json(&mut harness.client).await; + + let get_request = request_text(2, "get", json!({})); + send_text(&mut harness.client, get_request.clone()).await; + let ws_get_response = recv_json(&mut harness.client).await; + + // AC1: in-process and WebSocket dispatcher results must be byte-equivalent for the same + // request against an equivalently seeded backend. + let in_process = Dispatcher::new(FixtureBackend::new(store), ConnectionContext::default()); + let in_process_init: Value = serde_json::from_str(&in_process.dispatch(&init_request).await) + .expect("in-process initialize response must be valid JSON"); + let in_process_get: Value = serde_json::from_str(&in_process.dispatch(&get_request).await) + .expect("in-process get response must be valid JSON"); + + assert_eq!(ws_init_response, in_process_init); + assert_eq!(ws_get_response, in_process_get); + + shut_down(harness).await; +} + +#[tokio::test] +async fn binary_frame_closes_with_unsupported_data_code() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(store)).await; + + harness + .client + .send(Message::Binary(vec![1, 2, 3].into())) + .await + .unwrap(); + + let close = tokio::time::timeout(Duration::from_secs(1), harness.client.next()) + .await + .expect("server must close promptly on a binary frame"); + match close { + Some(Ok(Message::Close(Some(frame)))) => assert_eq!(frame.code, CloseCode::Unsupported), + other => panic!("expected a close frame with code 1003, got {other:?}"), + } + + shut_down(harness).await; +} + +#[tokio::test] +async fn oversized_text_frame_closes_with_message_too_big_code() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(store)).await; + + // The client has no outgoing size cap of its own; the server's `websocket_config()` bounds the + // *receiving* side, so this exercises that bound deterministically. + let oversized = "a".repeat(OVERSIZED_TEXT_BYTES); + harness.client.send(Message::text(oversized)).await.unwrap(); + + let close = tokio::time::timeout(Duration::from_secs(1), harness.client.next()) + .await + .expect("server must close promptly on an oversized frame"); + match close { + Some(Ok(Message::Close(Some(frame)))) => assert_eq!(frame.code, CloseCode::Size), + other => panic!("expected a close frame with code 1009, got {other:?}"), + } + + shut_down(harness).await; +} + +#[tokio::test] +async fn malformed_json_frame_receives_parse_error_without_closing() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(store)).await; + + send_text(&mut harness.client, "not valid json").await; + let error_response = recv_json(&mut harness.client).await; + assert_eq!(error_response["error"]["code"], -32700); + + // The connection must remain usable after the parse error. + send_text(&mut harness.client, request_text(9, "get", json!({}))).await; + let get_response = recv_json(&mut harness.client).await; + assert_eq!(get_response["id"], 9); + assert!(get_response.get("result").is_some(), "{get_response}"); + + shut_down(harness).await; +} + +#[tokio::test] +async fn cancel_request_for_unknown_id_is_a_silent_no_op() { + let store = Arc::new(FixtureStore::new(RunId::new("run-1"), Vec::new(), 8)); + let mut harness = spawn_server(FixtureBackend::new(store)).await; + + let cancel = json!({ + "jsonrpc": "2.0", + "method": "$/cancelRequest", + "params": {"id": "unknown-id"}, + }) + .to_string(); + send_text(&mut harness.client, cancel).await; + + // No response is emitted for the cancel notification itself; the connection must remain + // usable, so a subsequent unary request completes normally. + send_text(&mut harness.client, request_text(1, "get", json!({}))).await; + let get_response = + tokio::time::timeout(Duration::from_millis(500), recv_json(&mut harness.client)) + .await + .expect("connection must remain usable after an unknown $/cancelRequest id"); + assert_eq!(get_response["id"], 1); + assert!(get_response.get("result").is_some(), "{get_response}"); + + shut_down(harness).await; +} + +impl RequestChannel for Harness { + async fn send_get(&mut self, id: i64) { + send_text(&mut self.client, request_text(id, "get", json!({}))).await; + } + + async fn recv_value(&mut self) -> Value { + recv_json(&mut self.client).await + } +} + +impl GatedHarnessSpawn for Harness { + async fn spawn_gated(backend: GatedBackend) -> Self { + spawn_server(backend).await + } +} + +#[tokio::test] +async fn duplicate_in_flight_request_ids_are_rejected() { + let (mut harness, gate) = spawn_gated_harness::().await; + + assert_duplicate_in_flight_ids_are_rejected(&mut harness, &gate).await; + + shut_down(harness).await; +} + +#[tokio::test] +async fn excess_requests_are_rejected_with_server_busy() { + let (mut harness, _gate) = spawn_gated_harness::().await; + + assert_excess_requests_are_rejected_with_server_busy(&mut harness, MAX_CONNECTION_TASKS).await; + + shut_down(harness).await; +} diff --git a/crates/openengine-cluster-testkit/Cargo.toml b/crates/openengine-cluster-testkit/Cargo.toml index 0b365953..a7ffb217 100644 --- a/crates/openengine-cluster-testkit/Cargo.toml +++ b/crates/openengine-cluster-testkit/Cargo.toml @@ -18,6 +18,8 @@ thiserror.workspace = true tokio.workspace = true [dev-dependencies] +futures-util.workspace = true jsonschema.workspace = true sha2.workspace = true tokio.workspace = true +tokio-tungstenite.workspace = true diff --git a/crates/openengine-cluster-testkit/src/artifacts.rs b/crates/openengine-cluster-testkit/src/artifacts.rs index fc71cc22..25fc3b1f 100644 --- a/crates/openengine-cluster-testkit/src/artifacts.rs +++ b/crates/openengine-cluster-testkit/src/artifacts.rs @@ -3,13 +3,13 @@ use std::path::{Path, PathBuf}; use openengine_cluster_protocol::{ AgentAttachClosedNotification, AgentAttachEventNotification, AgentAttachParams, - AgentAttachResult, ApplyParams, ApplyResult, ArtifactRef, CompiledGraphIr, DeleteParams, - DeleteResult, EventNotification, GetParams, GetResult, GraphDiagnostic, GraphSpec, - InitializeParams, InitializeResult, JsonRpcNotification, JsonRpcRequest, JsonRpcResponse, - LogEventNotification, LogsClosedNotification, LogsParams, LogsResult, PlanParams, PlanResult, - ResubmitParams, ResubmitResult, RetryParams, RetryResult, StopParams, StopResult, - StructuralBounds, SubscriptionCancelParams, SubscriptionClosedNotification, UpdateParams, - UpdateResult, WatchParams, WatchResult, + AgentAttachResult, ApplyParams, ApplyResult, ArtifactRef, CancelRequestParams, CompiledGraphIr, + DeleteParams, DeleteResult, EventNotification, GetParams, GetResult, GraphDiagnostic, + GraphSpec, InitializeParams, InitializeResult, JsonRpcNotification, JsonRpcRequest, + JsonRpcResponse, LogEventNotification, LogsClosedNotification, LogsParams, LogsResult, + PlanParams, PlanResult, ResubmitParams, ResubmitResult, RetryParams, RetryResult, StopParams, + StopResult, StructuralBounds, SubscriptionCancelParams, SubscriptionClosedNotification, + UpdateParams, UpdateResult, WatchParams, WatchResult, }; use openengine_cluster_server::{ConnectionContext, Dispatcher}; use schemars::{schema_for, JsonSchema}; @@ -69,6 +69,7 @@ pub struct ImplementedProtocolSchema { pub event_notification: JsonRpcNotification, pub subscription_cancel_notification: JsonRpcNotification, pub subscription_closed_notification: JsonRpcNotification, + pub cancel_request_notification: JsonRpcNotification, pub logs_request: JsonRpcRequest, pub logs_response: JsonRpcResponse, pub log_event_notification: JsonRpcNotification, diff --git a/crates/openengine-cluster-testkit/src/artifacts/openrpc.rs b/crates/openengine-cluster-testkit/src/artifacts/openrpc.rs index 3ba5ebcf..57b5f72b 100644 --- a/crates/openengine-cluster-testkit/src/artifacts/openrpc.rs +++ b/crates/openengine-cluster-testkit/src/artifacts/openrpc.rs @@ -41,11 +41,14 @@ pub(super) fn document() -> Value { below, shared by every subscription-based method. There is no watch/event, \ watch/cancel, watch/closed, logs/event, logs/cancel, logs/closed, \ agent/attach/event, agent/attach/cancel, or agent/attach/closed method on the \ - wire.", + wire. `$/cancelRequest` is a transport-level best-effort cancellation of any \ + in-flight unary request by its RequestId; it is silently a no-op for an unknown \ + or already-completed id and carries no rollback claim after backend commit.", "notifications": { "event": { "$ref": "schema.json#/$defs/EventNotification" }, "subscription/cancel": { "$ref": "schema.json#/$defs/SubscriptionCancelParams" }, - "subscription/closed": { "$ref": "schema.json#/$defs/SubscriptionClosedNotification" } + "subscription/closed": { "$ref": "schema.json#/$defs/SubscriptionClosedNotification" }, + "$/cancelRequest": { "$ref": "schema.json#/$defs/CancelRequestParams" } } } }) diff --git a/crates/openengine-cluster-testkit/tests/protocol_ndjson.rs b/crates/openengine-cluster-testkit/tests/protocol_ndjson.rs index ba1bb05b..0e89da68 100644 --- a/crates/openengine-cluster-testkit/tests/protocol_ndjson.rs +++ b/crates/openengine-cluster-testkit/tests/protocol_ndjson.rs @@ -5,24 +5,13 @@ //! two-instance comparison pattern as `protocol_v1.rs`'s //! `admission_transcript_matches_in_process_and_stdio`. -use std::sync::Arc; use std::time::Duration; -use openengine_cluster_client::{ - ClusterClient, EventOrClosed, InProcessTransport, NdjsonReconnectingEventStream, - NdjsonTransport, NdjsonWatchClient, WatchClient, -}; -use openengine_cluster_protocol::{Cursor, GetParams, GraphSpec, StopMode, WatchEvent, WatchParams}; -use openengine_cluster_server::admission::AdmissionCoordinator; -use openengine_cluster_server::watch::PublicEventRecord; -use openengine_cluster_server::{ConnectionContext, Dispatcher}; -use openengine_cluster_testkit::admission::{ - compiled_from_graph_fixture, graph_fixture, InMemoryAdmissionStore, ScriptedOutcome, - ScriptedVerifier, -}; +use openengine_cluster_client::{ClusterClient, NdjsonTransport, NdjsonWatchClient}; +use openengine_cluster_protocol::{GetParams, StopMode, WatchParams}; +use openengine_cluster_testkit::admission::graph_fixture; use openengine_cluster_testkit::lifecycle::stop; use serde_json::Value; -use tokio::io::{AsyncRead, AsyncWrite}; #[path = "admission_support/committed.rs"] mod committed_support; @@ -31,109 +20,12 @@ use committed_support::committed; #[path = "stdio_subprocess_support/mod.rs"] mod stdio_subprocess_support; -/// Collects [`EventOrClosed`]s from `stream` until (and including) `Finished`, panicking if the -/// stream closes first. -macro_rules! collect_transcript { - ($stream:expr) => {{ - let mut events = Vec::new(); - loop { - match $stream.next().await.expect("stream ended before Finished") { - EventOrClosed::Event(record) => { - let finished = matches!(record.event, WatchEvent::Finished { .. }); - events.push(record); - if finished { - break; - } - } - EventOrClosed::Closed { reason, .. } => { - panic!("stream closed ({reason:?}) before the Finished event was observed") - } - } - } - events - }}; -} - -/// Runs one apply/get/stop lifecycle against a fresh in-process backend while a `watch` -/// subscription streams on the same dispatcher, returning its collected transcript. -async fn in_process_side_transcript(graph: &GraphSpec) -> Vec { - let compiled = compiled_from_graph_fixture(graph); - let verifier = Arc::new(ScriptedVerifier::new(vec![ScriptedOutcome::approve( - compiled, - vec![], - )])); - let store = Arc::new(InMemoryAdmissionStore::default()); - let backend = AdmissionCoordinator::from_shared(verifier, store); - let dispatcher = Dispatcher::new(backend, ConnectionContext::default()); - let in_process_client = ClusterClient::new(InProcessTransport::new(dispatcher.clone())); - in_process_client.initialize().await.unwrap(); - let in_process_watch = WatchClient::new(dispatcher); - - let (_parked, mut in_process_stream, _handle) = in_process_watch - .watch(WatchParams::default()) - .await - .unwrap(); - - let apply_result = in_process_client - .apply(committed( - graph.clone(), - Value::Null, - 0, - "in-process-create", - )) - .await - .unwrap(); - let generation = apply_result.generation.unwrap().get(); - // AC: a unary request completes correctly while the watch subscription is actively - // streaming on the same connection. - let get_result = in_process_client.get(GetParams::default()).await.unwrap(); - assert_eq!(get_result.spec, Some(graph.clone())); - in_process_client - .stop(stop(StopMode::Drain, generation, "in-process-stop")) - .await - .unwrap(); - collect_transcript!(in_process_stream) -} - -fn assert_transcripts_match(in_process: &[PublicEventRecord], ndjson: &[PublicEventRecord]) { - assert_eq!(in_process.len(), ndjson.len()); - for (in_process, ndjson) in in_process.iter().zip(ndjson.iter()) { - assert_eq!(in_process.cursor, ndjson.cursor); - assert_eq!(in_process.event, ndjson.event); - } -} - -/// Asserts the at-most-one-post-cancel-leak model for a subscription cancelled before its run was -/// ever committed to: the server-side subscription task may already have been parked awaiting the -/// next live event at the moment cancellation was processed, so at most one further event (the -/// commit's own first event, immediately following cancellation) may still leak through before it -/// observes cancellation on its next poll and stops for good. -async fn assert_cancel_probe_leak_model<'a, R, W>( - mut cancel_probe: NdjsonReconnectingEventStream<'a, R, W>, - first_committed_cursor: &Cursor, -) where - R: AsyncRead + Send + Unpin + 'static, - W: AsyncWrite + Send + Unpin + 'static, -{ - let mut leaked = Vec::new(); - loop { - match tokio::time::timeout(Duration::from_millis(300), cancel_probe.next()).await { - Ok(Some(EventOrClosed::Event(record))) => leaked.push(record), - Ok(Some(other)) => panic!("unexpected notification after cancel: {other:?}"), - Ok(None) | Err(_) => break, - } - } - assert!( - leaked.len() <= 1, - "cancelled probe subscription received more than one post-cancel event: {leaked:?}" - ); - if let Some(record) = leaked.first() { - assert_eq!( - record.cursor, *first_committed_cursor, - "cancellation failed to stop delivery before the run's first committed event" - ); - } -} +#[path = "protocol_transcript_support/mod.rs"] +mod protocol_transcript_support; +use protocol_transcript_support::{ + assert_cancel_probe_leak_model, assert_transcripts_match, collect_transcript, + in_process_side_transcript, +}; #[tokio::test] async fn ndjson_watch_transcript_matches_in_process_and_shares_its_connection() { diff --git a/crates/openengine-cluster-testkit/tests/protocol_transcript_support/mod.rs b/crates/openengine-cluster-testkit/tests/protocol_transcript_support/mod.rs new file mode 100644 index 00000000..162f0a8e --- /dev/null +++ b/crates/openengine-cluster-testkit/tests/protocol_transcript_support/mod.rs @@ -0,0 +1,136 @@ +//! Shared cross-transport watch-transcript scenario used identically by `protocol_ndjson.rs` (the +//! NDJSON binding from #745) and `protocol_websocket.rs` (the WebSocket binding from #651) to +//! prove each wire binding reproduces the exact same watch transcript (cursor progression and +//! event algebra) as the in-process `Dispatcher::watch` passthrough from #647, while sharing its +//! connection with ordinary unary traffic and honoring `subscription/cancel`. + +use std::sync::Arc; +use std::time::Duration; + +use openengine_cluster_client::{ + ClusterClient, EventOrClosed, InProcessTransport, SubscriptionTransport, WatchClient, + WatchSubscriptionEventStream, +}; +use openengine_cluster_protocol::{Cursor, GetParams, GraphSpec, StopMode, WatchParams}; +use openengine_cluster_server::admission::AdmissionCoordinator; +use openengine_cluster_server::watch::PublicEventRecord; +use openengine_cluster_server::{ConnectionContext, Dispatcher}; +use openengine_cluster_testkit::admission::{ + compiled_from_graph_fixture, InMemoryAdmissionStore, ScriptedOutcome, ScriptedVerifier, +}; +use openengine_cluster_testkit::lifecycle::stop; +use serde_json::Value; + +use crate::committed_support::committed; + +/// Collects [`EventOrClosed`]s from `stream` until (and including) `Finished`, panicking if the +/// stream closes first. Uses fully-qualified paths (rather than relying on this module's own +/// `use` imports) since `macro_rules!` item-path resolution follows the invocation site, and this +/// macro is invoked from sibling test-binary crates via `pub(crate) use`. +macro_rules! collect_transcript { + ($stream:expr) => {{ + let mut events = Vec::new(); + loop { + match $stream.next().await.expect("stream ended before Finished") { + ::openengine_cluster_client::EventOrClosed::Event(record) => { + let finished = matches!( + record.event, + ::openengine_cluster_protocol::WatchEvent::Finished { .. } + ); + events.push(record); + if finished { + break; + } + } + ::openengine_cluster_client::EventOrClosed::Closed { reason, .. } => { + panic!("stream closed ({reason:?}) before the Finished event was observed") + } + } + } + events + }}; +} +pub(crate) use collect_transcript; + +/// Runs one apply/get/stop lifecycle against a fresh in-process backend while a `watch` +/// subscription streams on the same dispatcher, returning its collected transcript. +pub async fn in_process_side_transcript(graph: &GraphSpec) -> Vec { + let compiled = compiled_from_graph_fixture(graph); + let verifier = Arc::new(ScriptedVerifier::new(vec![ScriptedOutcome::approve( + compiled, + vec![], + )])); + let store = Arc::new(InMemoryAdmissionStore::default()); + let backend = AdmissionCoordinator::from_shared(verifier, store); + let dispatcher = Dispatcher::new(backend, ConnectionContext::default()); + let in_process_client = ClusterClient::new(InProcessTransport::new(dispatcher.clone())); + in_process_client.initialize().await.unwrap(); + let in_process_watch = WatchClient::new(dispatcher); + + let (_parked, mut in_process_stream, _handle) = in_process_watch + .watch(WatchParams::default()) + .await + .unwrap(); + + let apply_result = in_process_client + .apply(committed( + graph.clone(), + Value::Null, + 0, + "in-process-create", + )) + .await + .unwrap(); + let generation = apply_result.generation.unwrap().get(); + // AC: a unary request completes correctly while the watch subscription is actively + // streaming on the same connection. + let get_result = in_process_client.get(GetParams::default()).await.unwrap(); + assert_eq!(get_result.spec, Some(graph.clone())); + in_process_client + .stop(stop(StopMode::Drain, generation, "in-process-stop")) + .await + .unwrap(); + collect_transcript!(in_process_stream) +} + +pub fn assert_transcripts_match(in_process: &[PublicEventRecord], wire: &[PublicEventRecord]) { + assert_eq!(in_process.len(), wire.len()); + for (in_process, wire) in in_process.iter().zip(wire.iter()) { + assert_eq!(in_process.cursor, wire.cursor); + assert_eq!(in_process.event, wire.event); + } +} + +/// Asserts the at-most-one-post-cancel-leak model for a subscription cancelled before its run was +/// ever committed to: the server-side subscription task may already have been parked awaiting the +/// next live event at the moment cancellation was processed, so at most one further event (the +/// commit's own first event, immediately following cancellation) may still leak through before it +/// observes cancellation on its next poll and stops for good. Generic over any +/// [`SubscriptionTransport`]-backed [`WatchSubscriptionEventStream`] (NDJSON's +/// `NdjsonReconnectingEventStream` alias and WebSocket's binding both are this exact generic type), +/// so both wire bindings share this exact scenario logic rather than duplicating it per transport. +pub async fn assert_cancel_probe_leak_model<'a, T>( + mut cancel_probe: WatchSubscriptionEventStream<'a, T>, + first_committed_cursor: &Cursor, +) where + T: SubscriptionTransport, +{ + let mut leaked = Vec::new(); + loop { + match tokio::time::timeout(Duration::from_millis(300), cancel_probe.next()).await { + Ok(Some(EventOrClosed::Event(record))) => leaked.push(record), + Ok(Some(other)) => panic!("unexpected notification after cancel: {other:?}"), + Ok(None) | Err(_) => break, + } + } + assert!( + leaked.len() <= 1, + "cancelled probe subscription received more than one post-cancel event: {leaked:?}" + ); + if let Some(record) = leaked.first() { + assert_eq!( + record.cursor, *first_committed_cursor, + "cancellation failed to stop delivery before the run's first committed event" + ); + } +} diff --git a/crates/openengine-cluster-testkit/tests/protocol_websocket.rs b/crates/openengine-cluster-testkit/tests/protocol_websocket.rs new file mode 100644 index 00000000..2275794a --- /dev/null +++ b/crates/openengine-cluster-testkit/tests/protocol_websocket.rs @@ -0,0 +1,197 @@ +//! Cross-transport equivalence: the WebSocket binding from #651 must reproduce the exact same +//! watch transcript (cursor progression and event algebra) as the in-process `Dispatcher::watch` +//! passthrough from #647 and the NDJSON binding from #745 (see `protocol_ndjson.rs`), while +//! sharing its connection with ordinary unary traffic and honoring `subscription/cancel`. Also +//! proves two independently-authorized WebSocket connections sharing one backend (AC3) preserve +//! CAS/idempotency and cannot observe each other's injected `ConnectionContext`, since +//! `ConnectionContext` is constructed once per connection and never derived from protocol params. + +use std::sync::Arc; +use std::time::Duration; + +use openengine_cluster_client::{ClusterClient, WatchSubscriptionClient, WebSocketTransport}; +use openengine_cluster_protocol::{GetParams, StopMode, WatchParams}; +use openengine_cluster_server::admission::AdmissionCoordinator; +use openengine_cluster_server::watch::fixtures::{await_websocket_shutdown, spawn_websocket}; +use openengine_cluster_server::{ClusterBackend, ConnectionContext, Dispatcher}; +use openengine_cluster_testkit::admission::{ + compiled_from_graph_fixture, graph_fixture, InMemoryAdmissionStore, ScriptedOutcome, + ScriptedVerifier, +}; +use openengine_cluster_testkit::lifecycle::stop; +use serde_json::Value; + +#[path = "admission_support/committed.rs"] +mod committed_support; +use committed_support::committed; + +#[path = "protocol_transcript_support/mod.rs"] +mod protocol_transcript_support; +use protocol_transcript_support::{ + assert_cancel_probe_leak_model, assert_transcripts_match, collect_transcript, + in_process_side_transcript, +}; + +#[tokio::test] +async fn websocket_watch_transcript_matches_in_process_and_shares_its_connection() { + let graph = graph_fixture("worker", serde_json::json!({"kind":"null"})); + + let in_process_events = in_process_side_transcript(&graph).await; + + // WebSocket side, against a fresh in-process duplex-backed `serve_websocket` connection wired + // the same way as `protocol_ndjson.rs`'s subprocess-backed NDJSON side. + let verifier = Arc::new(ScriptedVerifier::new(vec![ScriptedOutcome::approve( + compiled_from_graph_fixture(&graph), + vec![], + )])); + let store = Arc::new(InMemoryAdmissionStore::default()); + let backend = AdmissionCoordinator::from_shared(verifier, store); + let (ws, server) = spawn_websocket(backend).await; + + let transport = WebSocketTransport::new(ws); + let ws_client = ClusterClient::new(&transport); + ws_client.initialize().await.unwrap(); + let ws_watch = WatchSubscriptionClient::new(&transport); + + let (_parked, mut ws_stream) = ws_watch.watch(WatchParams::default()).await.unwrap(); + + // AC: `subscription/cancel` releases only the cancelled subscription. A second, still-parked + // subscription is cancelled immediately; it must observe nothing further even though it would + // otherwise park-attach to the very run committed below. + let (_parked, cancel_probe) = ws_watch.watch(WatchParams::default()).await.unwrap(); + cancel_probe.cancel().await.unwrap(); + tokio::time::sleep(Duration::from_millis(50)).await; + + let apply_result = ws_client + .apply(committed( + graph.clone(), + Value::Null, + 0, + "websocket-wire-create", + )) + .await + .unwrap(); + let generation = apply_result.generation.unwrap().get(); + // AC: a unary request completes correctly while the watch subscription is actively + // streaming on the same connection. + let get_result = ws_client.get(GetParams::default()).await.unwrap(); + assert_eq!(get_result.spec, Some(graph.clone())); + ws_client + .stop(stop(StopMode::Drain, generation, "websocket-wire-stop")) + .await + .unwrap(); + let ws_events = collect_transcript!(ws_stream); + + assert_transcripts_match(&in_process_events, &ws_events); + assert_cancel_probe_leak_model(cancel_probe, &ws_events[0].cursor).await; + + drop(ws_stream); + drop(transport); + await_websocket_shutdown(server).await; +} + +/// Wires `backend` to a fresh [`openengine_cluster_server::websocket::serve_websocket`] task over +/// an in-memory duplex pipe pair, injecting `peer_label` into that connection's +/// [`ConnectionContext`] -- unlike [`spawn_websocket`], which always uses +/// [`Dispatcher::new`]/[`ConnectionContext::default`], this uses +/// [`Dispatcher::from_shared`] so multiple independently-authorized connections can share one +/// backend, exactly as AC3 requires. +async fn spawn_websocket_with_shared_backend( + backend: Arc, + peer_label: &str, +) -> ( + tokio_tungstenite::WebSocketStream, + tokio::task::JoinHandle>, +) +where + B: ClusterBackend, +{ + let (client_io, server_io) = tokio::io::duplex(1 << 16); + let dispatcher = Dispatcher::from_shared( + backend, + ConnectionContext { + peer_label: Some(peer_label.to_owned()), + ..ConnectionContext::default() + }, + ); + let server = tokio::spawn(async move { + let ws = tokio_tungstenite::accept_async_with_config( + server_io, + Some(openengine_cluster_server::websocket::websocket_config()), + ) + .await + .expect("server handshake must succeed"); + openengine_cluster_server::websocket::serve_websocket(dispatcher, ws).await + }); + let (client, _response) = + tokio_tungstenite::client_async("ws://localhost/websocket", client_io) + .await + .expect("client handshake must succeed"); + (client, server) +} + +#[tokio::test] +async fn two_websocket_connections_share_one_backend_with_isolated_context_and_shared_idempotency() +{ + let graph = graph_fixture("worker", serde_json::json!({"kind":"null"})); + let verifier = Arc::new(ScriptedVerifier::new(vec![ScriptedOutcome::approve( + compiled_from_graph_fixture(&graph), + vec![], + )])); + let store = Arc::new(InMemoryAdmissionStore::default()); + let backend = Arc::new(AdmissionCoordinator::from_shared(verifier, store)); + + let (conn_a, server_a) = + spawn_websocket_with_shared_backend(Arc::clone(&backend), "connection-a").await; + let (conn_b, server_b) = + spawn_websocket_with_shared_backend(Arc::clone(&backend), "connection-b").await; + + let transport_a = WebSocketTransport::new(conn_a); + let transport_b = WebSocketTransport::new(conn_b); + let client_a = ClusterClient::new(&transport_a); + let client_b = ClusterClient::new(&transport_b); + + client_a.initialize().await.unwrap(); + client_b.initialize().await.unwrap(); + + // Connection A commits the run. + let apply_a = client_a + .apply(committed( + graph.clone(), + Value::Null, + 0, + "shared-backend-create", + )) + .await + .unwrap(); + assert!(!apply_a.deduped); + + // AC3: connection B replays the exact same idempotency key/params and dedups against A's + // commit, proving CAS/idempotency state is shared across independently-authorized connections + // on one backend rather than partitioned per connection. + let apply_b = client_b + .apply(committed( + graph.clone(), + Value::Null, + 0, + "shared-backend-create", + )) + .await + .unwrap(); + assert!(apply_b.deduped); + assert_eq!(apply_b.generation, apply_a.generation); + assert_eq!(apply_b.run_id, apply_a.run_id); + + // AC3: both connections observe the same committed spec via `get`, independent of which + // connection created it -- `ConnectionContext` carries no route/tenant field that could + // partition visibility between them. + let get_a = client_a.get(GetParams::default()).await.unwrap(); + let get_b = client_b.get(GetParams::default()).await.unwrap(); + assert_eq!(get_a.spec, Some(graph.clone())); + assert_eq!(get_a, get_b); + + drop(transport_a); + drop(transport_b); + await_websocket_shutdown(server_a).await; + await_websocket_shutdown(server_b).await; +} diff --git a/docs/openengine-cluster-protocol/v1/data-plane.md b/docs/openengine-cluster-protocol/v1/data-plane.md new file mode 100644 index 00000000..6eb264b2 --- /dev/null +++ b/docs/openengine-cluster-protocol/v1/data-plane.md @@ -0,0 +1,75 @@ +# Cluster Protocol v1 WebSocket data plane + +This document defines the production Rust WebSocket binding: the wire framing that carries the +backend-neutral `Dispatcher` and the generic subscription framing defined in +[`watch.md`](./watch.md) over one WebSocket connection, plus the boundary between that binding and +everything that hosts it. It adds no protocol method or event semantics of its own beyond +`$/cancelRequest`; every method, event, and generic subscription notification is unchanged from the +in-process and NDJSON stdio bindings. + +## Framing + +One WebSocket text message carries exactly one JSON-RPC object (request, response, or +notification) -- there is no line delimiter and no batching. A binary frame is not a supported +encoding: it closes the connection immediately with code `1003` (unsupported data). A text message +whose UTF-8 byte length exceeds `1,048,576` after frame reassembly closes the connection with code +`1009` (message too big), matching the NDJSON stdio binding's line-length bound. A text message that +fails to parse as JSON-RPC receives a normal `PARSE_ERROR` JSON-RPC error response on the same +connection; unlike the size and encoding violations above, a parse failure does not close the +connection. Ping/pong control frames are answered automatically and otherwise ignored. + +Subscription delivery reuses the exact same generic `event`/`subscription/cancel`/ +`subscription/closed` notification framing described in `watch.md` -- there are no +WebSocket-specific subscription notification names, and results/events/errors are byte-equivalent +between the in-process, NDJSON, and WebSocket bindings for the same request sequence. + +## `$/cancelRequest` + +`$/cancelRequest{id}` is a client notification, best-effort cancelling an in-flight unary request by +its `RequestId` on the same connection. It carries no response of its own. An unknown or +already-completed id is a silent no-op: the connection remains fully usable and no error is +surfaced. Cancelling a request that has already committed backend state leaves that committed state +unchanged -- there is no rollback or compensation, only suppression of the (otherwise still +in-flight) response delivery. `$/cancelRequest` targets unary/passthrough requests only; an +established `watch`/`logs`/`agent/attach` subscription is cancelled with `subscription/cancel` +instead, exactly as over NDJSON. + +## Connection isolation over a shared backend + +One backend instance can be shared by multiple independently authorized WebSocket connections. Each +connection is constructed with its own injected `ConnectionContext`; the wire protocol carries no +auth, tenant, or route field, and no protocol parameter is ever interpreted as one -- that +identity/authorization decision is made once, by whatever hosts this binding, at connection +acceptance time, not by anything this crate parses from a request. CAS and idempotency guarantees +are enforced by the shared backend itself, so two connections against one backend cannot double +commit or observe effects attributable to another connection's context. + +## Capsule data-plane placement + +This binding is the data-plane surface exposed by one capsule: the runtime process or container +that accepts WebSocket connections and serves one backend's `Dispatcher` over them. Everything +upstream of an accepted connection is owned by whatever hosts that capsule, not by this crate: + +- provisioning or scheduling the capsule itself; +- TLS termination (this binding speaks plain WebSocket text frames; TLS, if any, terminates in + front of it); +- resolving tenancy, routing, or authentication/authorization for a connection before it reaches + `serve_websocket`; +- billing and usage metering; +- workspace or secret storage/services; +- artifact bytes (this protocol carries status, events, and control -- never artifact payloads); +- issuing or validating the token/credential that authorized the connection. + +A hosting process is expected to accept the raw connection, resolve and inject that connection's +`ConnectionContext`, and hand the rest to `serve_websocket`; this document defines only what happens +from that handoff onward. + +## Fixture and test boundary + +`crates/openengine-cluster-server/tests/websocket.rs` covers this binding's own framing and +admission behavior directly against raw tungstenite frames. +`crates/openengine-cluster-client/tests/websocket.rs` covers the typed `WebSocketTransport` client +against the same binding. `crates/openengine-cluster-testkit/tests/protocol_websocket.rs` proves +byte-equivalence against the in-process and NDJSON bindings and exercises two independently +authorized connections sharing one backend. All three drive `serve_websocket` over an in-memory +duplex pipe; none of them involve real network I/O, TLS, or the hosting concerns listed above. diff --git a/docs/openengine-cluster-protocol/v1/watch.md b/docs/openengine-cluster-protocol/v1/watch.md index dea15bee..4e64e5e7 100644 --- a/docs/openengine-cluster-protocol/v1/watch.md +++ b/docs/openengine-cluster-protocol/v1/watch.md @@ -2,8 +2,10 @@ This document defines the durable observation contract: retained per-run events, opaque cursors, generic subscription framing, and snapshot-tail reconnect. The Rust wire types are authoritative. -NDJSON/WebSocket line framing for this subscription surface is bound by a later issue; this slice -defines the wire types and an in-process Rust client/server behavior only. +This document defines the wire types and in-process Rust client/server behavior; the NDJSON stdio +binding lives in `openengine-cluster-server`'s `stdio` module, and the production WebSocket +binding -- framing, close codes, `$/cancelRequest`, and the hosted data-plane boundary -- is +defined in [`data-plane.md`](./data-plane.md). ## Method and framing diff --git a/protocol/openengine-cluster/v1/openrpc.json b/protocol/openengine-cluster/v1/openrpc.json index aaad8d0b..221aa99f 100644 --- a/protocol/openengine-cluster/v1/openrpc.json +++ b/protocol/openengine-cluster/v1/openrpc.json @@ -497,8 +497,11 @@ ], "openrpc": "1.3.2", "x-generic-subscription-framing": { - "description": "watch, logs, and agent/attach each establish a subscription via one normal JSON-RPC result; subsequent delivery uses the generic notification methods below, shared by every subscription-based method. There is no watch/event, watch/cancel, watch/closed, logs/event, logs/cancel, logs/closed, agent/attach/event, agent/attach/cancel, or agent/attach/closed method on the wire.", + "description": "watch, logs, and agent/attach each establish a subscription via one normal JSON-RPC result; subsequent delivery uses the generic notification methods below, shared by every subscription-based method. There is no watch/event, watch/cancel, watch/closed, logs/event, logs/cancel, logs/closed, agent/attach/event, agent/attach/cancel, or agent/attach/closed method on the wire. `$/cancelRequest` is a transport-level best-effort cancellation of any in-flight unary request by its RequestId; it is silently a no-op for an unknown or already-completed id and carries no rollback claim after backend commit.", "notifications": { + "$/cancelRequest": { + "$ref": "schema.json#/$defs/CancelRequestParams" + }, "event": { "$ref": "schema.json#/$defs/EventNotification" }, diff --git a/protocol/openengine-cluster/v1/schema.json b/protocol/openengine-cluster/v1/schema.json index c3581702..f99f89da 100644 --- a/protocol/openengine-cluster/v1/schema.json +++ b/protocol/openengine-cluster/v1/schema.json @@ -360,6 +360,19 @@ ], "type": "object" }, + "CancelRequestParams": { + "additionalProperties": false, + "description": "Wire body of the `$/cancelRequest` client notification: best-effort cancellation of an\nin-flight unary request by its `RequestId`. Unknown or already-completed ids are a silent\nno-op; cancelling after the backend has committed leaves that committed state unchanged and\nemits no response or compensation.", + "properties": { + "id": { + "$ref": "#/$defs/RequestId" + } + }, + "required": [ + "id" + ], + "type": "object" + }, "ChoiceBranch": { "additionalProperties": false, "properties": { @@ -1776,7 +1789,7 @@ "type": "string" }, "params": { - "$ref": "#/$defs/LogEventNotification" + "$ref": "#/$defs/CancelRequestParams" } }, "required": [ @@ -1795,7 +1808,7 @@ "type": "string" }, "params": { - "$ref": "#/$defs/LogsClosedNotification" + "$ref": "#/$defs/LogEventNotification" } }, "required": [ @@ -1814,7 +1827,7 @@ "type": "string" }, "params": { - "$ref": "#/$defs/AgentAttachEventNotification" + "$ref": "#/$defs/LogsClosedNotification" } }, "required": [ @@ -1825,6 +1838,25 @@ "type": "object" }, "JsonRpcNotification7": { + "properties": { + "jsonrpc": { + "type": "string" + }, + "method": { + "type": "string" + }, + "params": { + "$ref": "#/$defs/AgentAttachEventNotification" + } + }, + "required": [ + "jsonrpc", + "method", + "params" + ], + "type": "object" + }, + "JsonRpcNotification8": { "properties": { "jsonrpc": { "type": "string" @@ -3790,10 +3822,10 @@ "$schema": "https://json-schema.org/draft/2020-12/schema", "properties": { "agent_attach_closed_notification": { - "$ref": "#/$defs/JsonRpcNotification7" + "$ref": "#/$defs/JsonRpcNotification8" }, "agent_attach_event_notification": { - "$ref": "#/$defs/JsonRpcNotification6" + "$ref": "#/$defs/JsonRpcNotification7" }, "agent_attach_request": { "$ref": "#/$defs/JsonRpcRequest12" @@ -3807,6 +3839,9 @@ "apply_response": { "$ref": "#/$defs/JsonRpcResponse3" }, + "cancel_request_notification": { + "$ref": "#/$defs/JsonRpcNotification4" + }, "delete_request": { "$ref": "#/$defs/JsonRpcRequest9" }, @@ -3829,10 +3864,10 @@ "$ref": "#/$defs/JsonRpcResponse" }, "log_event_notification": { - "$ref": "#/$defs/JsonRpcNotification4" + "$ref": "#/$defs/JsonRpcNotification5" }, "logs_closed_notification": { - "$ref": "#/$defs/JsonRpcNotification5" + "$ref": "#/$defs/JsonRpcNotification6" }, "logs_request": { "$ref": "#/$defs/JsonRpcRequest11" @@ -3907,6 +3942,7 @@ "event_notification", "subscription_cancel_notification", "subscription_closed_notification", + "cancel_request_notification", "logs_request", "logs_response", "log_event_notification",