From 1ce785b4cedf21670a6909281239aacded2ecb90 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 26 Aug 2026 10:51:59 +0200 Subject: [PATCH 01/14] Auto-Reconnect: Connection replacement and batch subscribe Adds the two server-side changes needed for SDK auto-reconnect (proposal 0037). Connection replacement: a client which reconnects can now send a stable, client-generated `session_id` query parameter. Each connection still gets its own ConnectionId and its own lifecycle events. The session id only identifies which earlier connection a new one supersedes. When a connection claims a session held by a live connection of the same identity, that connection's actor is stopped and its module-side disconnect lifecycle is awaited before the new connection's client_connected, so a module never observes two live connections for one session, nor the old connection's client_disconnected after the new client_connected. Batch subscribe: SubscribeBatch and SubscribeBatchApplied register multiple query sets under one subscription-manager lock and evaluate them at a single transaction snapshot. No transaction update can interleave with the responses, so a reconnecting client can replay its subscriptions and reconcile its cache against a consistent snapshot. Per-set errors are reported in the response and the remaining sets still apply --- .../client-api-messages/src/websocket/v2.rs | 78 ++++ crates/client-api/src/routes/subscribe.rs | 79 +++- crates/core/src/client.rs | 2 + crates/core/src/client/client_connection.rs | 34 +- .../src/client/client_connection_index.rs | 13 +- .../core/src/client/client_session_index.rs | 288 ++++++++++++ crates/core/src/client/consume_each_list.rs | 7 + crates/core/src/client/message_handlers_v2.rs | 4 + crates/core/src/client/messages.rs | 11 + crates/core/src/host/module_host.rs | 46 +- .../src/host/wasm_common/module_host_actor.rs | 12 + .../subscription/module_subscription_actor.rs | 424 +++++++++++++++++- crates/core/src/worker_metrics/mod.rs | 6 +- crates/smoketests/Cargo.toml | 6 + crates/smoketests/modules/Cargo.lock | 8 + crates/smoketests/modules/Cargo.toml | 1 + .../modules/connection-session/Cargo.toml | 12 + .../modules/connection-session/src/lib.rs | 24 + crates/smoketests/tests/cluster.rs | 1 + .../tests/cluster/connection_session.rs | 412 +++++++++++++++++ sdks/rust/src/db_connection.rs | 5 + 21 files changed, 1445 insertions(+), 28 deletions(-) create mode 100644 crates/core/src/client/client_session_index.rs create mode 100644 crates/smoketests/modules/connection-session/Cargo.toml create mode 100644 crates/smoketests/modules/connection-session/src/lib.rs create mode 100644 crates/smoketests/tests/cluster/connection_session.rs diff --git a/crates/client-api-messages/src/websocket/v2.rs b/crates/client-api-messages/src/websocket/v2.rs index 734c28fdbe5..c203d4de702 100644 --- a/crates/client-api-messages/src/websocket/v2.rs +++ b/crates/client-api-messages/src/websocket/v2.rs @@ -26,6 +26,8 @@ pub enum ClientMessage { CallReducer(CallReducer), /// Invoke a procedure, a non-transactional side-effecting function which runs in the database. CallProcedure(CallProcedure), + /// Add multiple sets of subscribed queries in one atomic step. + SubscribeBatch(SubscribeBatch), } /// Sent by client to register a subscription to a new query set @@ -92,6 +94,42 @@ pub enum UnsubscribeFlags { SendDroppedRows = 1, } +/// Sent by client to register multiple subscriptions in one atomic step. +/// +/// The server registers every subscription set under a single subscription-manager +/// lock and evaluates all of them at a single transaction snapshot, +/// then responds with one [`SubscribeBatchApplied`] message carrying a result per set. +/// No [`TransactionUpdate`] is delivered between the registration of the first set +/// and the [`SubscribeBatchApplied`] response, +/// and updates for the new sets resume after it. +/// +/// A set whose queries are invalid or fail to compute reports an error in its +/// [`SubscribeSetResult`]. The remaining sets still apply. +#[derive(SpacetimeType)] +#[sats(crate = spacetimedb_lib)] +pub struct SubscribeBatch { + /// An identifier for a client request. + pub request_id: u32, + + /// The subscription sets to register. + /// + /// Each [`QuerySetId`] must be distinct, + /// and must not be used by any other subscription on the same connection. + pub sets: Box<[SubscribeSet]>, +} + +/// One subscription set within a [`SubscribeBatch`]. +#[derive(SpacetimeType)] +#[sats(crate = spacetimedb_lib)] +pub struct SubscribeSet { + /// An identifier for this subscription, + /// which should not be used for any other subscriptions on the same connection. + pub query_set_id: QuerySetId, + + /// A set of queries to subscribe to, each a single SQL `SELECT` statement. + pub query_strings: Box<[Box]>, +} + /// Sent by the client to perform a query at a single point in time. /// /// Unlike subscriptions registered by [`Subscribe`], this query will not receive real-time updates. @@ -193,6 +231,8 @@ pub enum ServerMessage { ReducerResult(ReducerResult), /// Sent in response to a [`CallProcedure`] message, containing the procedure's exit status. ProcedureResult(ProcedureResult), + /// Sent in response to a [`SubscribeBatch`] message, containing a result per query set. + SubscribeBatchApplied(SubscribeBatchApplied), } #[derive(SpacetimeType, Debug)] @@ -290,6 +330,44 @@ pub struct SubscriptionError { pub error: Box, } +/// Response to [`SubscribeBatch`], carrying one result per registered query set. +/// +/// This message's `request_id` matches the one the client provided in the [`SubscribeBatch`] message, +/// and `results` contains exactly one entry per received [`SubscribeSet`], in the same order. +/// +/// Every applied set's rows are evaluated at a single transaction snapshot. +#[derive(SpacetimeType, Debug)] +#[sats(crate = spacetimedb_lib)] +pub struct SubscribeBatchApplied { + /// The request_id of the corresponding [`SubscribeBatch`] message. + pub request_id: u32, + /// One result per query set, in the order the sets appeared in the request. + pub results: Box<[SubscribeSetResult]>, +} + +/// The result for one query set within a [`SubscribeBatchApplied`]. +#[derive(SpacetimeType, Debug)] +#[sats(crate = spacetimedb_lib)] +pub struct SubscribeSetResult { + /// The [`QuerySetId`] the client provided for this set. + pub query_set_id: QuerySetId, + /// The outcome for this set. + pub outcome: SubscribeSetOutcome, +} + +/// The outcome for one query set within a [`SubscribeBatchApplied`]. +#[derive(SpacetimeType, Debug)] +#[sats(crate = spacetimedb_lib)] +pub enum SubscribeSetOutcome { + /// The set was applied; contains its initial matching rows. + /// The set behaves like one registered with an individual [`Subscribe`] afterwards. + Applied(QueryRows), + /// The set failed to compile or compute. + /// The set is not registered; its [`QuerySetId`] may be re-used. + /// The error string follows the conventions of [`SubscriptionError`]'s `error` field. + Error(Box), +} + /// Sent by the server to the client after a transaction runs and commits successfully in the database, /// containing [`QuerySetUpdate`]s for each of the client's subscribed query sets /// whose results were affected by the transaction. diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 165216202e8..5c956fb8937 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -28,7 +28,7 @@ use spacetimedb::client::messages::{ }; use spacetimedb::client::{ ClientActorId, ClientConfig, ClientConnection, ClientConnectionReceiver, DataMessage, MessageExecutionError, - MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, WsVersion, + MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, SessionId, WsVersion, }; use spacetimedb::host::module_host::ClientConnectedError; use spacetimedb::host::NoSuchModule; @@ -86,6 +86,17 @@ pub struct SubscribeParams { #[derive(Deserialize)] pub struct SubscribeQueryParams { pub connection_id: Option, + /// A client-generated identifier for a logical client session, + /// stable across the reconnects of one client connection object. + /// + /// When a connection supplies a session id already held by a live + /// connection of the same identity, the old connection is torn down before + /// this one runs `client_connected`, so the module never observes two live + /// connections for one session. + /// See [`spacetimedb::client::ClientSessionIndex`]. + /// + /// Connections which do not supply one behave exactly as before. + pub session_id: Option, #[serde(default)] pub compression: ws_v1::Compression, /// Whether we want "light" responses, tailored to network bandwidth constrained clients. @@ -100,6 +111,25 @@ pub struct SubscribeQueryParams { pub confirmed: Option, } +/// A [`SessionId`] as supplied in the `session_id` query parameter. +/// Represented by a 32-character hex string. +pub struct SessionIdForUrl(SessionId); + +impl<'de> Deserialize<'de> for SessionIdForUrl { + fn deserialize>(deserializer: D) -> Result { + let hex = >::deserialize(deserializer)?; + let value = u128::from_str_radix(&hex, 16) + .map_err(|_| serde::de::Error::custom("session_id must be a hex-encoded 128-bit value"))?; + Ok(Self(SessionId::from_u128(value))) + } +} + +impl From for SessionId { + fn from(session_id: SessionIdForUrl) -> Self { + session_id.0 + } +} + fn resolve_confirmed_reads_default(version: WsVersion, confirmed: Option) -> bool { if let Some(confirmed) = confirmed { return confirmed; @@ -119,6 +149,7 @@ pub async fn handle_websocket( Path(SubscribeParams { name_or_identity }): Path, Query(SubscribeQueryParams { connection_id, + session_id, compression, light, confirmed, @@ -228,6 +259,8 @@ where connection_id, name: ctx.client_actor_index().next_client_name(), }; + let session_id: Option = session_id.map(Into::into); + let sessions = ctx.client_actor_index().sessions(); let ws_config = WebSocketConfig::default() .max_message_size(Some(0x2000000)) @@ -255,6 +288,29 @@ where log::debug!("websocket: New client connected from {client_log_string}"); + // If this connection resumes a session which a live connection still + // holds, that connection is taken over. Stop its actor and run its + // module-side disconnect to completion, so the module observes + // `client_disconnected` for it strictly before `client_connected` for + // this one, and never two live connections for one session. + if let Some(session_id) = session_id + && let Some(superseded) = sessions.claim_session(client_id, session_id, None) + { + log::debug!( + "websocket: Connection {} supersedes {} for session {session_id}", + client_id.connection_id, + superseded.client_id.connection_id, + ); + if let Some(sender) = &superseded.sender { + sender.kick(ClientDisconnectCause::ConnectionSuperseded); + } + // Awaiting this is what orders the two lifecycle reducers: + // both run on the module's main instance, and this one is + // enqueued first. + let module = module_rx.borrow().clone(); + module.disconnect_client(superseded.client_id).await; + } + let connected = match ClientConnection::call_client_connected_maybe_reject( &mut module_rx, client_id, @@ -284,6 +340,11 @@ where } }; record_client_rejection(db_identity, cause); + // The session claim is only meaningful for a connection which + // exists, so give it up again. + if let Some(session_id) = session_id { + sessions.release_session(client_id, session_id); + } return; } }; @@ -292,7 +353,15 @@ where "websocket: Database accepted connection from {client_log_string}; spawning ws_client_actor and ClientConnection" ); - let actor = |client, receiver| ws_client_actor(ws_opts, client, ws, receiver); + // Release the session claim when the actor ends, including when it is aborted. + let session_guard = session_id.map(|session_id| { + let sessions = sessions.clone(); + scopeguard::guard((), move |()| sessions.release_session(client_id, session_id)) + }); + let actor = |client, receiver| async move { + let _session_guard = session_guard; + ws_client_actor(ws_opts, client, ws, receiver).await; + }; let client = ClientConnection::spawn( client_id, auth.into(), @@ -305,6 +374,12 @@ where ) .await; + // Now that the actor exists, register its sender so that a later + // connection resuming this session can stop it. + if let Some(session_id) = session_id { + sessions.attach_sender(client_id, session_id, &client.sender()); + } + // Send the client their identity token message as the first message // NOTE: We're adding this to the protocol because some client libraries are // unable to access the http response headers. diff --git a/crates/core/src/client.rs b/crates/core/src/client.rs index 812d03c0701..5c90cba6175 100644 --- a/crates/core/src/client.rs +++ b/crates/core/src/client.rs @@ -3,6 +3,7 @@ use std::fmt; mod client_connection; mod client_connection_index; +mod client_session_index; pub mod consume_each_list; mod message_handlers; mod message_handlers_v1; @@ -16,6 +17,7 @@ pub use client_connection::{ WsVersion, }; pub use client_connection_index::ClientActorIndex; +pub use client_session_index::{ClientSessionIndex, SessionId, SupersededConnection}; pub use message_handlers::MessageHandleError; pub use message_handlers_v1::MessageExecutionError; pub use messages::OutboundMessage; diff --git a/crates/core/src/client/client_connection.rs b/crates/core/src/client/client_connection.rs index ed33e29b533..a50805e3854 100644 --- a/crates/core/src/client/client_connection.rs +++ b/crates/core/src/client/client_connection.rs @@ -408,6 +408,24 @@ impl ClientConnectionSender { self.cancelled.load(Ordering::Relaxed) } + /// Stop this connection's websocket actor. + /// + /// Used when a newer connection supersedes this one + /// (see [`super::ClientSessionIndex`]), and when a client exceeds its + /// outgoing queue capacity. + /// + /// This only stops the actor. The module-side disconnect + /// ([`crate::host::ModuleHost::disconnect_client`]) is run separately by + /// the actor's teardown, or by the caller when it needs that teardown to + /// complete before some other work. + pub fn kick(&self, cause: ClientDisconnectCause) { + if let Some(metrics) = &self.metrics { + metrics.disconnect_recorder.record(cause); + } + self.abort_handle.abort(); + self.cancelled.store(true, Ordering::Relaxed); + } + /// Send a message to the client. For data-related messages, you should probably use /// `BroadcastQueue::send` to ensure that the client sees data messages in a consistent order. /// @@ -455,12 +473,8 @@ impl ClientConnectionSender { ); if let Some(metrics) = &self.metrics { metrics.outgoing_queue_disconnects.inc(); - metrics - .disconnect_recorder - .record(ClientDisconnectCause::OutgoingQueueFull); } - self.abort_handle.abort(); - self.cancelled.store(true, Ordering::Relaxed); + self.kick(ClientDisconnectCause::OutgoingQueueFull); return Err(ClientSendError::Cancelled); } Err(mpsc::error::TrySendError::Closed(_)) => return Err(ClientSendError::Disconnected), @@ -1178,6 +1192,16 @@ impl ClientConnection { .call_view_add_v2_subscription(self.sender(), self.auth.clone(), request, timer) .await } + + pub async fn subscribe_batch( + &self, + request: ws_v2::SubscribeBatch, + timer: Instant, + ) -> Result, DBError> { + self.module() + .call_view_add_batch_subscription(self.sender(), self.auth.clone(), request, timer) + .await + } pub async fn subscribe_multi( &self, request: ws_v1::SubscribeMulti, diff --git a/crates/core/src/client/client_connection_index.rs b/crates/core/src/client/client_connection_index.rs index 7ad58ce4738..4ed43b7a61d 100644 --- a/crates/core/src/client/client_connection_index.rs +++ b/crates/core/src/client/client_connection_index.rs @@ -1,10 +1,12 @@ use std::sync::atomic::{AtomicU64, Ordering::Relaxed}; +use std::sync::Arc; -use super::ClientName; +use super::{ClientName, ClientSessionIndex}; #[derive(Default)] pub struct ClientActorIndex { client_name_auto_increment_state: AtomicU64, + sessions: Arc, } impl ClientActorIndex { @@ -14,4 +16,13 @@ impl ClientActorIndex { pub fn next_client_name(&self) -> ClientName { ClientName(self.client_name_auto_increment_state.fetch_add(1, Relaxed)) } + + /// The map of live client sessions, used to replace a connection + /// which a reconnect supersedes. + /// + /// Returns an owned handle, since the websocket handler needs one which + /// outlives the request. + pub fn sessions(&self) -> Arc { + self.sessions.clone() + } } diff --git a/crates/core/src/client/client_session_index.rs b/crates/core/src/client/client_session_index.rs new file mode 100644 index 00000000000..ddc3eae3b7b --- /dev/null +++ b/crates/core/src/client/client_session_index.rs @@ -0,0 +1,288 @@ +//! Tracking of client sessions, used to replace pre-existing connections. +//! +//! A client which reconnects automatically sends the same client-generated +//! session id on every connection attempt. Each connection still receives its +//! own [`ConnectionId`] and its own `client_connected` / `client_disconnected` +//! events. The session id only identifies which earlier connection a new one +//! supersedes. +//! +//! A client frequently notices a dropped connection before the server does +//! as the server needs up to its idle timeout to notice an idle peer. +//! Without this index the module would briefly observe two live +//! connections for the same client, and the old connection's +//! `client_disconnected` could run after the new connection's +//! `client_connected`. + +use std::collections::hash_map::Entry; +use std::collections::HashMap; +use std::sync::{Arc, Mutex, Weak}; + +use spacetimedb_lib::Identity; + +use super::{ClientActorId, ClientConnectionSender}; + +/// A client-generated identifier for a logical client session, +/// stable across the reconnects of one client connection object. +/// +/// Supplied by the client as the `session_id` query parameter. +#[derive(PartialEq, Eq, Clone, Copy, Hash, Debug, PartialOrd, Ord)] +pub struct SessionId(u128); + +impl SessionId { + pub fn from_u128(value: u128) -> Self { + Self(value) + } + + pub fn to_u128(self) -> u128 { + self.0 + } +} + +impl std::fmt::Display for SessionId { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:032x}", self.0) + } +} + +/// A session is identified by the client's identity together with the +/// client-generated session id, so a session can only ever be replaced by the +/// client that owns it. +type SessionKey = (Identity, SessionId); + +/// The connection currently serving a session. +struct SessionEntry { + client_id: ClientActorId, + /// Used to stop the connection's actor when it is superseded. + /// + /// Weak so that a connection whose actor has already ended can be dropped + /// normally rather than being kept alive by this map. + sender: Weak, +} + +/// The map of live sessions for one host. +/// +/// Maps each live session to the connection currently serving it. The entry is +/// removed when that connection ends, so a session never outlives its +/// connection. +#[derive(Default)] +pub struct ClientSessionIndex { + sessions: Mutex>, +} + +/// A connection which a newly arriving connection supersedes. +/// Returned by `ClientSessionIndex::claim_session`. +pub struct SupersededConnection { + /// The client actor matching a session id and identity. + /// Used by the caller to run the module side disconnect + /// lifecyle before allowing the new connection's `client_connected` to run. + pub client_id: ClientActorId, + /// The superseded connection's sender, if its actor is still alive. + /// Used to terminate the connection's websocket actor. + pub sender: Option>, +} + +impl ClientSessionIndex { + pub fn new() -> Self { + Self::default() + } + + /// Claim `session_id` for `client`, returning the connection it supersedes, + /// if any. + /// + /// The caller must tear that connection down before allowing `client`'s + /// `client_connected` to run, so that the module never observes two live + /// connections for one session. The claim takes effect immediately, so a + /// third connection racing for the same session supersedes `client` rather + /// than the connection returned here. + /// + /// `sender` is registered so that a later connection can stop this + /// connection's actor. It is `None` before the connection's actor exists, + /// in which case the entry is registered without one. + pub fn claim_session( + &self, + client: ClientActorId, + session_id: SessionId, + sender: Option<&Arc>, + ) -> Option { + let key = (client.identity, session_id); + let entry = SessionEntry { + client_id: client, + sender: sender.map(Arc::downgrade).unwrap_or_default(), + }; + let mut sessions = self.sessions.lock().expect("session index poisoned"); + match sessions.entry(key) { + Entry::Occupied(mut occupied) => { + let superseded = occupied.insert(entry); + (superseded.client_id.connection_id != client.connection_id).then(|| SupersededConnection { + client_id: superseded.client_id, + sender: superseded.sender.upgrade(), + }) + } + Entry::Vacant(vacant) => { + vacant.insert(entry); + None + } + } + } + + /// Record the sender for the connection currently holding `session_id`. + /// + /// Called once the connection's actor exists. Does nothing if the session + /// has already been claimed by a newer connection. + pub fn attach_sender(&self, client: ClientActorId, session_id: SessionId, sender: &Arc) { + let key = (client.identity, session_id); + let mut sessions = self.sessions.lock().expect("session index poisoned"); + if let Some(entry) = sessions.get_mut(&key) + && entry.client_id.connection_id == client.connection_id + { + entry.sender = Arc::downgrade(sender); + } + } + + /// Release `session_id` if it is still held by `client`. + /// + /// Called when a connection ends. A connection which has already been + /// superseded no longer holds the session, so it leaves the entry alone: + /// otherwise a slow teardown would evict its own replacement. + pub fn release_session(&self, client: ClientActorId, session_id: SessionId) { + let key = (client.identity, session_id); + let mut sessions = self.sessions.lock().expect("session index poisoned"); + if let Entry::Occupied(entry) = sessions.entry(key) + && entry.get().client_id.connection_id == client.connection_id + { + entry.remove(); + } + } + + /// The number of live sessions. Intended for tests and diagnostics. + pub fn len(&self) -> usize { + self.sessions.lock().expect("session index poisoned").len() + } + + pub fn is_empty(&self) -> bool { + self.len() == 0 + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::ClientName; + use spacetimedb_lib::ConnectionId; + + fn client(identity: Identity, connection_id: u128) -> ClientActorId { + ClientActorId { + identity, + connection_id: ConnectionId::from_u128(connection_id), + name: ClientName(0), + } + } + + fn an_identity() -> Identity { + Identity::from_byte_array([1; 32]) + } + + fn another_identity() -> Identity { + Identity::from_byte_array([2; 32]) + } + + #[test] + fn first_connection_supersedes_nothing() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + assert!(index.claim_session(client(an_identity(), 1), session, None).is_none()); + assert_eq!(index.len(), 1); + } + + #[test] + fn reconnect_supersedes_previous_connection() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + index.claim_session(client(an_identity(), 1), session, None); + + let superseded = index.claim_session(client(an_identity(), 2), session, None); + + assert_eq!( + superseded.map(|s| s.client_id.connection_id), + Some(ConnectionId::from_u128(1)) + ); + // The session is now held by the new connection, not the old one. + assert_eq!(index.len(), 1); + } + + #[test] + fn different_identity_does_not_supersede() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + index.claim_session(client(an_identity(), 1), session, None); + + let superseded = index.claim_session(client(another_identity(), 2), session, None); + + assert!(superseded.is_none()); + assert_eq!(index.len(), 2); + } + + #[test] + fn different_session_does_not_supersede() { + let index = ClientSessionIndex::new(); + index.claim_session(client(an_identity(), 1), SessionId::from_u128(7), None); + + let superseded = index.claim_session(client(an_identity(), 2), SessionId::from_u128(8), None); + + assert!(superseded.is_none()); + assert_eq!(index.len(), 2); + } + + #[test] + fn release_removes_the_session() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + let connection = client(an_identity(), 1); + index.claim_session(connection, session, None); + + index.release_session(connection, session); + + assert!(index.is_empty()); + } + + #[test] + fn superseded_connection_release_does_not_evict_its_replacement() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + let old = client(an_identity(), 1); + let new = client(an_identity(), 2); + index.claim_session(old, session, None); + index.claim_session(new, session, None); + + // The old connection tears down after being superseded. + index.release_session(old, session); + + // The replacement still holds the session. + assert_eq!(index.len(), 1); + assert_eq!( + index + .claim_session(client(an_identity(), 3), session, None) + .map(|s| s.client_id.connection_id), + Some(new.connection_id) + ); + } + + #[test] + fn three_way_race_supersedes_the_most_recent_connection() { + let index = ClientSessionIndex::new(); + let session = SessionId::from_u128(7); + index.claim_session(client(an_identity(), 1), session, None); + + let second = index.claim_session(client(an_identity(), 2), session, None); + let third = index.claim_session(client(an_identity(), 3), session, None); + + assert_eq!( + second.map(|s| s.client_id.connection_id), + Some(ConnectionId::from_u128(1)) + ); + assert_eq!( + third.map(|s| s.client_id.connection_id), + Some(ConnectionId::from_u128(2)) + ); + } +} diff --git a/crates/core/src/client/consume_each_list.rs b/crates/core/src/client/consume_each_list.rs index 96fc4fe3414..5a382b7d2f4 100644 --- a/crates/core/src/client/consume_each_list.rs +++ b/crates/core/src/client/consume_each_list.rs @@ -52,6 +52,13 @@ impl ConsumeEachBuffer for ws_v2::ServerMessage { use ws_v2::ServerMessage::*; match self { SubscribeApplied(x) => x.rows.consume_each_list(each), + SubscribeBatchApplied(x) => { + for result in x.results { + if let ws_v2::SubscribeSetOutcome::Applied(rows) = result.outcome { + rows.consume_each_list(each); + } + } + } OneOffQueryResult(x) => x.result.ok().consume_each_list(each), UnsubscribeApplied(x) => x.rows.consume_each_list(each), SubscriptionError(_) | InitialConnection(_) | ProcedureResult(_) => {} diff --git a/crates/core/src/client/message_handlers_v2.rs b/crates/core/src/client/message_handlers_v2.rs index d228fda9fcd..f112162dcd6 100644 --- a/crates/core/src/client/message_handlers_v2.rs +++ b/crates/core/src/client/message_handlers_v2.rs @@ -32,6 +32,10 @@ pub(super) async fn handle_decoded_message( let res = client.subscribe_v2(subscribe, timer).await; res.map(drop).map_err(|e| (None, None, e.into())) } + ws_v2::ClientMessage::SubscribeBatch(subscribe_batch) => { + let res = client.subscribe_batch(subscribe_batch, timer).await; + res.map(drop).map_err(|e| (None, None, e.into())) + } ws_v2::ClientMessage::Unsubscribe(unsubscribe) => { let res = client.unsubscribe_v2(unsubscribe, timer).await; res.map(drop).map_err(|e| (None, None, e.into())) diff --git a/crates/core/src/client/messages.rs b/crates/core/src/client/messages.rs index 2de3a676bc0..6999a3cd0da 100644 --- a/crates/core/src/client/messages.rs +++ b/crates/core/src/client/messages.rs @@ -316,6 +316,7 @@ impl OutboundMessage { Self::V2(message) => match message { ws_v2::ServerMessage::InitialConnection(_) => None, ws_v2::ServerMessage::SubscribeApplied(_) => Some(WorkloadType::Subscribe), + ws_v2::ServerMessage::SubscribeBatchApplied(_) => Some(WorkloadType::Subscribe), ws_v2::ServerMessage::UnsubscribeApplied(_) => Some(WorkloadType::Unsubscribe), ws_v2::ServerMessage::SubscriptionError(_) => None, ws_v2::ServerMessage::TransactionUpdate(_) => Some(WorkloadType::Update), @@ -331,6 +332,16 @@ fn v2_message_num_rows(message: &ws_v2::ServerMessage) -> Option { match message { ws_v2::ServerMessage::InitialConnection(_) => None, ws_v2::ServerMessage::SubscribeApplied(message) => Some(count_query_rows(&message.rows)), + ws_v2::ServerMessage::SubscribeBatchApplied(message) => Some( + message + .results + .iter() + .map(|result| match &result.outcome { + ws_v2::SubscribeSetOutcome::Applied(rows) => count_query_rows(rows), + ws_v2::SubscribeSetOutcome::Error(_) => 0, + }) + .sum(), + ), ws_v2::ServerMessage::UnsubscribeApplied(message) => { Some(message.rows.as_ref().map(count_query_rows).unwrap_or_default()) } diff --git a/crates/core/src/host/module_host.rs b/crates/core/src/host/module_host.rs index eb15f619e2a..e9235a4da1c 100644 --- a/crates/core/src/host/module_host.rs +++ b/crates/core/src/host/module_host.rs @@ -876,6 +876,12 @@ pub enum ViewCommand { request: ws_v2::Subscribe, _timer: Instant, }, + AddBatchSubscription { + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + _timer: Instant, + }, RemoveSingleSubscription { sender: Arc, auth: AuthCtx, @@ -916,6 +922,13 @@ pub(in crate::host) enum ViewCommandErrorTarget { request_id: Option, query_set_id: ws_v2::QuerySetId, }, + /// A [`ViewCommand::AddBatchSubscription`] which failed as a whole. + /// Every set in the batch is reported as failed with the same error. + Batch { + sender: Arc, + request_id: RequestId, + query_set_ids: Box<[ws_v2::QuerySetId]>, + }, } impl ViewCommand { @@ -924,7 +937,8 @@ impl ViewCommand { Self::AddSingleSubscription { _timer, .. } | Self::AddMultiSubscription { _timer, .. } | Self::AddLegacySubscription { _timer, .. } - | Self::AddSubscriptionV2 { _timer, .. } => ViewCommandMetric { + | Self::AddSubscriptionV2 { _timer, .. } + | Self::AddBatchSubscription { _timer, .. } => ViewCommandMetric { workload: WorkloadType::Subscribe, timer: *_timer, }, @@ -998,6 +1012,11 @@ impl ViewCommand { request_id: Some(request.request_id), query_set_id: request.query_set_id, }, + Self::AddBatchSubscription { sender, request, .. } => ViewCommandErrorTarget::Batch { + sender: sender.clone(), + request_id: request.request_id, + query_set_ids: request.sets.iter().map(|set| set.query_set_id).collect(), + }, } } } @@ -1027,6 +1046,16 @@ impl ViewCommandErrorTarget { *query_set_id, err.to_string().into(), ), + Self::Batch { + sender, + request_id, + query_set_ids, + } => subscriptions.send_batch_subscription_error( + sender.clone(), + *request_id, + query_set_ids, + err.to_string().into(), + ), }; if let Err(send_err) = res { log::warn!("failed to send subscription error: {send_err:#}"); @@ -2502,6 +2531,21 @@ impl ModuleHost { } } + call_view_command_method! { + pub async fn call_view_add_batch_subscription( + &self, + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + timer: Instant, + ) -> "call_view_add_batch_subscription" => AddBatchSubscription { + sender, + auth, + request, + _timer: timer, + } + } + call_view_command_method! { pub async fn call_view_remove_single_subscription( &self, diff --git a/crates/core/src/host/wasm_common/module_host_actor.rs b/crates/core/src/host/wasm_common/module_host_actor.rs index 1d6db3763fd..a0f0c070d28 100644 --- a/crates/core/src/host/wasm_common/module_host_actor.rs +++ b/crates/core/src/host/wasm_common/module_host_actor.rs @@ -1204,6 +1204,18 @@ impl InstanceCommon { Ok((metrics, trapped)) => (Ok(metrics), trapped), Err(err) => (Err(err), false), }, + ViewCommand::AddBatchSubscription { + sender, + auth, + request, + _timer: timer, + } => match info + .subscriptions + .add_batch_subscription_with_instance(&mut inst, sender, auth, request, timer, None) + { + Ok((metrics, trapped)) => (Ok(metrics), trapped), + Err(err) => (Err(err), false), + }, ViewCommand::RemoveSingleSubscription { sender, auth, diff --git a/crates/core/src/subscription/module_subscription_actor.rs b/crates/core/src/subscription/module_subscription_actor.rs index a3f1058308d..9265796b971 100644 --- a/crates/core/src/subscription/module_subscription_actor.rs +++ b/crates/core/src/subscription/module_subscription_actor.rs @@ -218,6 +218,53 @@ struct CompiledQueryBatch { compile_timer: HistogramTimer, } +/// Like [`CompiledQueryBatch`], but without mut_tx. +/// The queries were compiled under a tx owned by the caller. +/// Returned by [`ModuleSubscriptions::compile_hashed_queries`]. +struct CompiledQueries { + queries: Vec>, + physical_plans: HashMap>, + auth: AuthCtx, + compile_timer: HistogramTimer, +} + +/// The queries of a subscribe message, hashed by [`hash_queries`] +/// for compilation cache lookup. +struct HashedQueries<'a> { + subscribe_to_all_tables: bool, + /// Each query's SQL along with its unparameterized and parameterized hashes. + query_hashes: Vec<(&'a str, QueryHash, QueryHash)>, + /// The number of queries in the message, for allocation sizing. + num_queries: usize, +} + +/// Hashes the queries in a subscribe message for compilation cache lookup. +/// +/// This requires only the query strings, and should be called +/// before taking the db lock. +/// See doc comment on [`ModuleSubscriptions::compile_queries`]. +fn hash_queries<'a>(sender: Identity, queries: &'a [Box], num_queries: usize) -> HashedQueries<'a> { + let mut subscribe_to_all_tables = false; + let mut query_hashes = Vec::with_capacity(num_queries); + + for sql in queries { + let sql = sql.trim(); + if is_subscribe_to_all_tables(sql) { + subscribe_to_all_tables = true; + continue; + } + let hash = QueryHash::from_string(sql, sender, false); + let hash_with_param = QueryHash::from_string(sql, sender, true); + query_hashes.push((sql, hash, hash_with_param)); + } + + HashedQueries { + subscribe_to_all_tables, + query_hashes, + num_queries, + } +} + #[derive(Clone, Copy)] enum FailedSubscription { V1(ws_v1::QueryId), @@ -1069,24 +1116,42 @@ impl ModuleSubscriptions { num_queries: usize, metrics: &SubscriptionMetrics, ) -> Result { - let mut subscribe_to_all_tables = false; - let mut plans = Vec::with_capacity(num_queries); - let mut query_hashes = Vec::with_capacity(num_queries); - - for sql in queries { - let sql = sql.trim(); - if is_subscribe_to_all_tables(sql) { - subscribe_to_all_tables = true; - continue; - } - let hash = QueryHash::from_string(sql, sender, false); - let hash_with_param = QueryHash::from_string(sql, sender, true); - query_hashes.push((sql, hash, hash_with_param)); - } + let hashed = hash_queries(sender, queries, num_queries); // We always get the db lock before the subscription lock to avoid deadlocks. let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe); + let CompiledQueries { + queries, + physical_plans, + auth, + compile_timer, + } = self.compile_hashed_queries(hashed, auth, metrics, &mut_tx)?; + + Ok(CompiledQueryBatch { + queries, + physical_plans, + auth, + mut_tx: ScopeGuard::::into_inner(mut_tx), + compile_timer, + }) + } + + /// Compiles the queries hashed by [`hash_queries`] under `mut_tx`. + fn compile_hashed_queries( + &self, + hashed: HashedQueries<'_>, + auth: AuthCtx, + metrics: &SubscriptionMetrics, + mut_tx: &MutTxId, + ) -> Result { + let HashedQueries { + subscribe_to_all_tables, + query_hashes, + num_queries, + } = hashed; + let mut plans = Vec::with_capacity(num_queries); + let compile_timer = metrics.compilation_time.start_timer(); let guard = { @@ -1104,7 +1169,7 @@ impl ModuleSubscriptions { for compiled in super::subscription::get_all( |relational_db, tx| relational_db.get_all_tables_mut(tx).map(|schemas| schemas.into_iter()), &self.relational_db, - &*mut_tx, + mut_tx, &auth, )? { add_compiled_query( @@ -1133,7 +1198,7 @@ impl ModuleSubscriptions { plans.push(unit); } _ => { - let compiled = compile_query_with_hashes(&auth, &*mut_tx, sql, hash, hash_with_param) + let compiled = compile_query_with_hashes(&auth, mut_tx, sql, hash, hash_with_param) .map_err(|err| DBError::WithSql { error: Box::new(DBError::Other(err.into())), sql: sql.into(), @@ -1155,11 +1220,10 @@ impl ModuleSubscriptions { // How many queries in this subscription are not cached? metrics.num_new_queries_subscribed.inc_by(new_queries); - Ok(CompiledQueryBatch { + Ok(CompiledQueries { queries: plans, physical_plans, auth, - mut_tx: ScopeGuard::::into_inner(mut_tx), compile_timer, }) } @@ -1251,6 +1315,32 @@ impl ModuleSubscriptions { ) } + /// Report a whole-batch failure, marking every set in the batch as failed. + /// + /// Used when a [`ws_v2::SubscribeBatch`] fails before per-set outcomes + /// could be determined. + pub fn send_batch_subscription_error( + &self, + recipient: Arc, + request_id: RequestId, + query_set_ids: &[ws_v2::QuerySetId], + message: Box, + ) -> Result<(), BroadcastError> { + let results = query_set_ids + .iter() + .map(|query_set_id| ws_v2::SubscribeSetResult { + query_set_id: *query_set_id, + outcome: ws_v2::SubscribeSetOutcome::Error(message.clone()), + }) + .collect::>() + .into_boxed_slice(); + self.broadcast_queue.send_client_message_v2( + recipient, + None, + ws_v2::SubscribeBatchApplied { request_id, results }, + ) + } + /// Add a subscription consisting of multiple queries. /// /// Read more in [`Self::add_single_subscription`]. @@ -1269,6 +1359,36 @@ impl ModuleSubscriptions { None => panic!("v2 subscriptions without a module host are not supported yet"), } } + + /// Add multiple query sets in one atomic step, in response to a + /// [`ws_v2::SubscribeBatch`] message. + /// + /// Every set is registered under a single subscription-manager lock and + /// evaluated at a single transaction snapshot, so no transaction update + /// for any of the new sets can precede the [`ws_v2::SubscribeBatchApplied`] + /// response, and updates resume after it, all relative to the same snapshot. + /// + /// A set which fails to compile or evaluate reports a per-set error in the + /// response while the remaining sets still apply. + #[tracing::instrument(level = "trace", skip_all)] + pub async fn add_batch_subscription( + &self, + host: Option<&ModuleHost>, + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + timer: Instant, + _assert: Option, + ) -> Result, DBError> { + match host { + Some(host) => { + host.call_view_add_batch_subscription(sender, auth, request, timer) + .await + } + None => panic!("batch subscriptions without a module host are not supported yet"), + } + } + /// Add a subscription consisting of multiple queries. /// /// Read more in [`Self::add_single_subscription`]. @@ -1308,6 +1428,21 @@ impl ModuleSubscriptions { ) -> Result<(Option, bool), DBError> { self.add_v2_subscription_inner(Some(instance), sender, auth, request, timer, _assert) } + + /// Similar to [`Self::add_v2_subscription_with_instance`], + /// but registers every query set of a batch atomically. + pub(crate) fn add_batch_subscription_with_instance( + &self, + instance: &mut RefInstance, + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + timer: Instant, + _assert: Option, + ) -> Result<(Option, bool), DBError> { + self.add_batch_subscription_inner(Some(instance), sender, auth, request, timer, _assert) + } + /// Similar to [`Self::add_single_subscription_with_instance`], /// but for multiple queries. pub(crate) fn add_multi_subscription_with_instance( @@ -1425,6 +1560,166 @@ impl ModuleSubscriptions { Ok((Some(metrics), trapped)) } + + /// Implementation of [`Self::add_batch_subscription`]. + /// + /// Each set is compiled and evaluated in the same way as an individual + /// `Subscribe` ([`Self::compile_queries`], [`Self::check_new_query_row_limit`], + /// [`Self::evaluate_queries`]). + /// All subscription queries are evaluated at a single transaction snapshot + /// and answered in one response. A set which fails compilation, + /// the row limit, or evaluation is reported as an error in the response and + /// is not registered while the remaining sets still apply. + fn add_batch_subscription_inner( + &self, + instance: Option<&mut RefInstance<'_, I>>, + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + _timer: Instant, + _assert: Option, + ) -> Result<(Option, bool), DBError> { + let subscription_metrics = &self.metrics.subscribe; + + // The per-set outcome, in request order. Sets which fail before + // evaluation are filled in here and skipped later. + let mut outcomes: Vec> = (0..request.sets.len()).map(|_| None).collect(); + // The sets which compiled, as (request index, query set id, plans). + let mut compiled_sets: Vec<(usize, ws_v2::QuerySetId, Vec>)> = Vec::new(); + let mut physical_plans: HashMap> = HashMap::default(); + + let num_queries: usize = request.sets.iter().map(|set| set.query_strings.len()).sum(); + subscription_metrics.num_queries_subscribed.inc_by(num_queries as _); + + // Hash every set's queries before taking the db lock, + // for the reasons described in [`Self::compile_queries`]. + let hashed_sets: Vec> = request + .sets + .iter() + .map(|set| hash_queries(sender.id.identity, &set.query_strings, set.query_strings.len())) + .collect(); + + // We always get the db lock before the subscription lock to avoid deadlocks. + // + // A single transaction spans the compilation, registration, and + // evaluation of every set, exactly as an individual `Subscribe` holds + // one transaction end-to-end. This means no schema change can + // invalidate a compiled set before it is registered, at the cost of + // holding the db lock across all compilations. + let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe); + + // Compile every set. A set which fails to compile is reported in its + // outcome and does not fail the batch. + for (index, hashed) in hashed_sets.into_iter().enumerate() { + match self.compile_hashed_queries(hashed, auth.clone(), subscription_metrics, &mut_tx) { + Ok(CompiledQueries { + queries, + physical_plans: set_physical_plans, + auth: _, + compile_timer: _compile_timer, + }) => { + physical_plans.extend(set_physical_plans); + compiled_sets.push((index, request.sets[index].query_set_id, queries)); + } + Err(err) => { + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + } + } + } + + // Register every compiled set under a single write lock, so no + // transaction committed between registrations can be observed by some + // sets but not others. + let mut registered: Vec<(usize, ws_v2::QuerySetId, Vec>)> = Vec::new(); + { + let mut subscriptions = { + let _wait_guard = subscription_metrics.lock_waiters.inc_scope(); + let _wait_timer = subscription_metrics.lock_wait_time.start_timer(); + self.subscriptions.write() + }; + for (index, query_set_id, queries) in compiled_sets { + match subscriptions.add_subscription_v2(sender.clone(), queries.clone(), query_set_id) { + Ok(_) => registered.push((index, query_set_id, queries)), + Err(err) => { + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + } + } + } + } + + let mut_tx = ScopeGuard::::into_inner(mut_tx); + let all_queries: Vec> = registered + .iter() + .flat_map(|(_, _, queries)| queries.iter().cloned()) + .collect(); + let (mut tx, tx_offset, trapped) = + self.materialize_views_and_downgrade_tx(mut_tx, instance, &all_queries, auth.caller())?; + + // Evaluate every registered set against the single snapshot above. + // This will do the same row-limit check and evaluation of an individual + // `Subscribe` (see [`Self::add_v2_subscription_inner`]), except that a + // failure is recorded as the set's outcome instead of aborting the + // request, and the registration is removed so the failed set never + // receives transaction updates. + let mut total_metrics = ExecutionMetrics::default(); + for (index, query_set_id, queries) in registered { + let failed_subscription = FailedSubscription::V2(query_set_id); + if let Err(err) = self.check_new_query_row_limit(&queries, &physical_plans, &tx, &auth) { + self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + continue; + } + + let Ok((update, metrics)) = + self.evaluate_queries(sender.clone(), &queries, &tx, TableUpdateType::Subscribe) + else { + self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error( + "Internal error evaluating queries".into(), + )); + continue; + }; + tx.metrics.merge(metrics); + + subscription_metrics.num_queries_evaluated.inc_by(queries.len() as _); + + let rows = match update { + ws_v1::FormatSwitch::Bsatn(update) => query_rows_from_update(update, false)?, + ws_v1::FormatSwitch::Json(_) => { + return Err(DBError::Other(anyhow::anyhow!( + "v2 subscriptions require binary protocol" + ))) + } + }; + total_metrics.merge(metrics); + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Applied(rows)); + } + + // One response for the whole batch, so nothing interleaves with it. + let results = outcomes + .into_iter() + .zip(request.sets.iter()) + .map(|(outcome, set)| ws_v2::SubscribeSetResult { + query_set_id: set.query_set_id, + outcome: outcome.unwrap_or_else(|| { + ws_v2::SubscribeSetOutcome::Error("Internal error registering query set".into()) + }), + }) + .collect::>() + .into_boxed_slice(); + + let _ = self.broadcast_queue.send_client_message_v2( + sender.clone(), + Some(tx_offset), + ws_v2::SubscribeBatchApplied { + request_id: request.request_id, + results, + }, + ); + + Ok((Some(total_metrics), trapped)) + } + fn add_multi_subscription_inner( &self, instance: Option<&mut RefInstance>, @@ -2483,6 +2778,99 @@ mod tests { Ok(()) } + /// Test that a batch subscription answers all sets in one message, + /// applies and registers the valid sets, + /// and reports an invalid set's error without failing the batch. + #[tokio::test] + async fn subscribe_batch_applies_sets_and_reports_errors() -> anyhow::Result<()> { + let db = relational_db()?; + + let client_id = client_id_from_u8(1); + let (sender, mut rx) = v2_client_connection(client_id, &db); + + let auth = AuthCtx::new(db.owner_identity(), client_id.identity); + let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone()); + + let t_id = db.create_table_for_test("t", &[("x", AlgebraicType::U8)], &[])?; + db.create_table_for_test("s", &[("x", AlgebraicType::U8)], &[])?; + with_auto_commit(&db, |tx| -> anyhow::Result<_> { + db.insert(tx, t_id, &bsatn::to_vec(&product![1_u8])?)?; + Ok(()) + })?; + + subs.add_batch_subscription_inner::( + None, + sender.clone(), + auth, + ws_v2::SubscribeBatch { + request_id: 1, + sets: [ + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(1), + query_strings: ["select * from t".into()].into(), + }, + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(2), + query_strings: ["select * from no_such_table".into()].into(), + }, + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(3), + query_strings: ["select * from s".into()].into(), + }, + ] + .into(), + }, + Instant::now(), + None, + )?; + + // The whole batch is answered by a single message, + // with one result per set in request order. + let results = match rx.recv().await { + Some(OutboundMessage::V2(ws_v2::ServerMessage::SubscribeBatchApplied(msg))) => { + assert_eq!(msg.request_id, 1); + msg.results + } + other => panic!("Expected v2 SubscribeBatchApplied, got: {other:?}"), + }; + let [first, second, third] = &*results else { + panic!("Expected one result per set, got: {results:?}"); + }; + + // The first set is applied with the initial row of `t`. + assert_eq!(first.query_set_id, ws_v2::QuerySetId::new(1)); + match &first.outcome { + ws_v2::SubscribeSetOutcome::Applied(rows) => { + assert_eq!(rows.tables.len(), 1); + assert_eq!(rows.tables[0].rows.len(), 1); + } + other => panic!("Expected the first set to be applied, got: {other:?}"), + } + + // The second set fails to compile, but does not fail the batch. + assert_eq!(second.query_set_id, ws_v2::QuerySetId::new(2)); + assert!( + matches!(&second.outcome, ws_v2::SubscribeSetOutcome::Error(_)), + "Expected the second set to error, got: {second:?}" + ); + + // The third set is applied even though the second errored. + assert_eq!(third.query_set_id, ws_v2::QuerySetId::new(3)); + assert!( + matches!(&third.outcome, ws_v2::SubscribeSetOutcome::Applied(_)), + "Expected the third set to be applied, got: {third:?}" + ); + + // The applied sets are registered for updates, + // and the failed set is not. + commit_tx(&db, &subs, [], [(t_id, product![2_u8])])?; + + let schema = ProductType::from([AlgebraicType::U8]); + assert_v2_tx_update_for_table(rx.recv(), ws_v2::QuerySetId::new(1), "t", &schema, [product![2_u8]], []).await; + + Ok(()) + } + #[tokio::test] async fn unsubscribe_v2_other_clients_receive_sender_view_updates() -> anyhow::Result<()> { let db = relational_db()?; diff --git a/crates/core/src/worker_metrics/mod.rs b/crates/core/src/worker_metrics/mod.rs index 5cecd3c40ab..f46206bc3f0 100644 --- a/crates/core/src/worker_metrics/mod.rs +++ b/crates/core/src/worker_metrics/mod.rs @@ -67,12 +67,14 @@ pub enum ClientDisconnectCause { WebsocketSendError, /// The websocket receive stream ended without a more specific cause. WebsocketStreamEnded, + /// A newer connection for the same client session superseded this one. + ConnectionSuperseded, /// The accepted websocket actor ended without a more specific recorded cause. Unknown, } impl ClientDisconnectCause { - pub const ALL: [Self; 22] = [ + pub const ALL: [Self; 23] = [ Self::ClientClose, Self::IdleTimeout, Self::IncomingQueueFull, @@ -94,6 +96,7 @@ impl ClientDisconnectCause { Self::WebsocketReceiveHttpFormat, Self::WebsocketSendError, Self::WebsocketStreamEnded, + Self::ConnectionSuperseded, Self::Unknown, ]; @@ -120,6 +123,7 @@ impl ClientDisconnectCause { Self::WebsocketReceiveHttpFormat => "websocket_receive_http_format", Self::WebsocketSendError => "websocket_send_error", Self::WebsocketStreamEnded => "websocket_stream_ended", + Self::ConnectionSuperseded => "connection_superseded", Self::Unknown => "unknown", } } diff --git a/crates/smoketests/Cargo.toml b/crates/smoketests/Cargo.toml index 90ad676634d..be6781e9251 100644 --- a/crates/smoketests/Cargo.toml +++ b/crates/smoketests/Cargo.toml @@ -17,11 +17,17 @@ reqwest = { workspace = true, features = ["blocking"] } which = "8.0.0" [dev-dependencies] +spacetimedb-core.workspace = true +spacetimedb-client-api-messages.workspace = true +spacetimedb-lib.workspace = true cargo_metadata.workspace = true +assert_cmd = "2" +futures.workspace = true predicates = "3" socket2.workspace = true tokio.workspace = true tokio-postgres.workspace = true +tokio-tungstenite.workspace = true xmltree.workspace = true [lints] diff --git a/crates/smoketests/modules/Cargo.lock b/crates/smoketests/modules/Cargo.lock index 040c9ee3ffe..ef2d8e49aff 100644 --- a/crates/smoketests/modules/Cargo.lock +++ b/crates/smoketests/modules/Cargo.lock @@ -734,6 +734,14 @@ dependencies = [ "spacetimedb", ] +[[package]] +name = "smoketest-module-connection-session" +version = "0.1.0" +dependencies = [ + "log", + "spacetimedb", +] + [[package]] name = "smoketest-module-delete-database" version = "0.1.0" diff --git a/crates/smoketests/modules/Cargo.toml b/crates/smoketests/modules/Cargo.toml index 63dc67687eb..d92ff936ba0 100644 --- a/crates/smoketests/modules/Cargo.toml +++ b/crates/smoketests/modules/Cargo.toml @@ -104,6 +104,7 @@ members = [ # Connection tests "connect-disconnect", + "connection-session", "confirmed-reads", "delete-database", "client-connection-reject", diff --git a/crates/smoketests/modules/connection-session/Cargo.toml b/crates/smoketests/modules/connection-session/Cargo.toml new file mode 100644 index 00000000000..26b7e1021cd --- /dev/null +++ b/crates/smoketests/modules/connection-session/Cargo.toml @@ -0,0 +1,12 @@ +[package] +name = "smoketest-module-connection-session" +version = "0.1.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["cdylib"] + +[dependencies] +spacetimedb.workspace = true +log.workspace = true diff --git a/crates/smoketests/modules/connection-session/src/lib.rs b/crates/smoketests/modules/connection-session/src/lib.rs new file mode 100644 index 00000000000..2b3daa189a6 --- /dev/null +++ b/crates/smoketests/modules/connection-session/src/lib.rs @@ -0,0 +1,24 @@ +//! Logs the lifecycle reducers with their connection ids, so tests can assert +//! the order in which connections are established and torn down. + +use spacetimedb::{log, ReducerContext}; + +#[spacetimedb::reducer(client_connected)] +pub fn connected(ctx: &ReducerContext) { + log::info!( + "connected {}", + ctx.connection_id() + .map(|id| id.to_hex().to_string()) + .unwrap_or_default() + ); +} + +#[spacetimedb::reducer(client_disconnected)] +pub fn disconnected(ctx: &ReducerContext) { + log::info!( + "disconnected {}", + ctx.connection_id() + .map(|id| id.to_hex().to_string()) + .unwrap_or_default() + ); +} diff --git a/crates/smoketests/tests/cluster.rs b/crates/smoketests/tests/cluster.rs index b4cb1865a82..f0aa305b01a 100644 --- a/crates/smoketests/tests/cluster.rs +++ b/crates/smoketests/tests/cluster.rs @@ -13,6 +13,7 @@ mod cluster { mod column_defaults; mod confirmed_reads; mod connect_disconnect_from_cli; + mod connection_session; mod database_lock; mod delete_database; mod describe; diff --git a/crates/smoketests/tests/cluster/connection_session.rs b/crates/smoketests/tests/cluster/connection_session.rs new file mode 100644 index 00000000000..d6048b2868d --- /dev/null +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -0,0 +1,412 @@ +//! Tests for connection replacement, the server side of SDK auto-reconnect. +//! +//! A reconnecting client supplies a stable `session_id`. When it reconnects +//! before the server has noticed the old socket died, the new connection +//! supersedes the old one. The old connection is torn down through the normal +//! disconnect sequence before the new connection's `client_connected` runs. + +use anyhow::{bail, Context, Result}; +use futures::{SinkExt, StreamExt}; +use spacetimedb_client_api_messages::websocket::{common as ws_common, v2 as ws_v2, v3 as ws_v3}; +use spacetimedb_lib::bsatn; +use spacetimedb_smoketests::Smoketest; +use tokio_tungstenite::tungstenite::client::IntoClientRequest; +use tokio_tungstenite::tungstenite::http::header::SEC_WEBSOCKET_PROTOCOL; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{connect_async, MaybeTlsStream, WebSocketStream}; + +type Socket = WebSocketStream>; + +/// A raw v3 websocket connection to a database, bypassing the SDKs so that a +/// test controls exactly which query parameters are sent. +struct TestConnection { + socket: Socket, + connection_id: String, +} + +impl TestConnection { + /// Open a connection, optionally supplying a `session_id`, and wait for the + /// server's `InitialConnection` message. + async fn open(test: &Smoketest, connection_id: &str, session_id: Option<&str>) -> Result { + let token = test.read_token()?; + let host = test.server_host(); + let database = test + .database_identity + .as_deref() + .context("test database has not been published")?; + + // Uncompressed, so the test can decode payloads with plain BSATN. + let mut url = + format!("ws://{host}/v1/database/{database}/subscribe?compression=None&connection_id={connection_id}"); + if let Some(session_id) = session_id { + url.push_str(&format!("&session_id={session_id}")); + } + + let mut request = url.into_client_request()?; + request + .headers_mut() + .insert(SEC_WEBSOCKET_PROTOCOL, ws_v3::BIN_PROTOCOL.parse()?); + request + .headers_mut() + .insert("Authorization", format!("Bearer {token}").parse()?); + + let (socket, response) = connect_async(request).await?; + let negotiated = response + .headers() + .get(SEC_WEBSOCKET_PROTOCOL) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default(); + if negotiated != ws_v3::BIN_PROTOCOL { + bail!("server negotiated {negotiated:?}, expected {}", ws_v3::BIN_PROTOCOL); + } + + let mut connection = Self { + socket, + connection_id: connection_id.to_string(), + }; + match connection.next_message().await? { + ws_v2::ServerMessage::InitialConnection(initial) => { + let established = initial.connection_id.to_hex().to_string(); + if established != connection.connection_id { + bail!( + "server established connection id {established}, expected {}", + connection.connection_id + ); + } + } + other => bail!("expected InitialConnection, got {other:?}"), + } + Ok(connection) + } + + /// Read the next server message, decoding the v3 framing, which packs one + /// or more messages into a single binary payload. + async fn next_message(&mut self) -> Result { + loop { + let message = self + .socket + .next() + .await + .context("websocket closed while awaiting a message")??; + match message { + Message::Binary(payload) => { + // Binary payloads start with a compression tag; the rest is + // one or more BSATN server messages back to back. + let (tag, mut body) = payload.split_first().context("empty binary websocket payload")?; + if *tag != ws_common::SERVER_MSG_COMPRESSION_TAG_NONE { + bail!("expected an uncompressed payload, got compression tag {tag}"); + } + return Ok(bsatn::from_reader(&mut body)?); + } + Message::Ping(_) | Message::Pong(_) => continue, + Message::Close(frame) => bail!("websocket closed: {frame:?}"), + other => bail!("unexpected websocket message: {other:?}"), + } + } + } + + async fn send(&mut self, message: ws_v2::ClientMessage) -> Result<()> { + let payload = bsatn::to_vec(&message)?; + self.socket.send(Message::Binary(payload.into())).await?; + Ok(()) + } + + /// Whether the server still serves this connection. + /// + /// A superseded connection's actor is stopped, so a request on it is never + /// answered. Note the server does not send a close frame. The peer's + /// socket stays half-open until it writes, which is what this does. + async fn is_still_served(&mut self) -> bool { + if self + .send(ws_v2::ClientMessage::Subscribe(ws_v2::Subscribe { + request_id: 999, + query_set_id: ws_v2::QuerySetId::new(999), + query_strings: vec!["SELECT * FROM st_client".into()].into_boxed_slice(), + })) + .await + .is_err() + { + return false; + } + matches!( + tokio::time::timeout(std::time::Duration::from_secs(10), self.next_message()).await, + Ok(Ok(_)) + ) + } +} + +fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .expect("failed to build tokio runtime") +} + +/// The order of lifecycle log lines for the given connection ids. +fn lifecycle_log(test: &Smoketest) -> Vec { + test.logs(200) + .unwrap_or_default() + .into_iter() + .filter(|line| line.contains("connected ") || line.contains("disconnected ")) + .collect() +} + +fn position_of(lines: &[String], event: &str, connection_id: &str) -> Option { + lines + .iter() + .position(|line| line.contains(&format!("{event} {connection_id}"))) +} + +/// Wait for a log line to appear, since the lifecycle reducers run +/// asynchronously with respect to the websocket handshake. +fn wait_for_log(test: &Smoketest, event: &str, connection_id: &str) -> Vec { + let deadline = std::time::Instant::now() + std::time::Duration::from_secs(10); + loop { + let lines = lifecycle_log(test); + if position_of(&lines, event, connection_id).is_some() { + return lines; + } + if std::time::Instant::now() > deadline { + panic!("timed out waiting for `{event} {connection_id}` in logs: {lines:?}"); + } + std::thread::sleep(std::time::Duration::from_millis(200)); + } +} + +const CONNECTION_A: &str = "00000000000000000000000000000a11"; +const CONNECTION_B: &str = "00000000000000000000000000000b22"; +const CONNECTION_C: &str = "00000000000000000000000000000c33"; +const SESSION: &str = "0000000000000000000000000000dead"; +const OTHER_SESSION: &str = "0000000000000000000000000000beef"; + +/// A second connection with the same session id supersedes the first: the old +/// connection is disconnected, and its `client_disconnected` runs strictly +/// before the new connection's `client_connected`. +#[test] +fn test_reconnect_with_same_session_replaces_connection() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let mut first = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("first connection failed"); + wait_for_log(&test, "connected", CONNECTION_A); + + // Reconnect with the same session before the server notices the drop. + let _second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("second connection failed"); + + assert!( + !first.is_still_served().await, + "the superseded connection should no longer be served" + ); + + let lines = wait_for_log(&test, "connected", CONNECTION_B); + let connected_a = position_of(&lines, "connected", CONNECTION_A).expect("A never connected"); + let disconnected_a = position_of(&lines, "disconnected", CONNECTION_A).expect("A never disconnected"); + let connected_b = position_of(&lines, "connected", CONNECTION_B).expect("B never connected"); + + assert!( + connected_a < disconnected_a, + "expected A to connect before disconnecting: {lines:?}" + ); + assert!( + disconnected_a < connected_b, + "expected A's client_disconnected to run before B's client_connected: {lines:?}" + ); + }); +} + +/// A connection with a different session id does not supersede: both stay live. +#[test] +fn test_different_session_does_not_replace_connection() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let _first = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("first connection failed"); + wait_for_log(&test, "connected", CONNECTION_A); + + let _second = TestConnection::open(&test, CONNECTION_B, Some(OTHER_SESSION)) + .await + .expect("second connection failed"); + let lines = wait_for_log(&test, "connected", CONNECTION_B); + + assert!( + position_of(&lines, "disconnected", CONNECTION_A).is_none(), + "the first connection should still be live: {lines:?}" + ); + }); +} + +/// A connection which supplies no session id behaves exactly as before: it +/// neither supersedes nor is superseded. +#[test] +fn test_connection_without_session_is_not_replaced() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let _first = TestConnection::open(&test, CONNECTION_A, None) + .await + .expect("first connection failed"); + wait_for_log(&test, "connected", CONNECTION_A); + + let _second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("second connection failed"); + let lines = wait_for_log(&test, "connected", CONNECTION_B); + + assert!( + position_of(&lines, "disconnected", CONNECTION_A).is_none(), + "a connection without a session id should not be superseded: {lines:?}" + ); + }); +} + +/// Repeated reconnects each supersede only the connection immediately before +/// them, leaving exactly one live connection for the session. +#[test] +fn test_repeated_reconnects_leave_one_live_connection() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let mut first = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("first connection failed"); + wait_for_log(&test, "connected", CONNECTION_A); + + let mut second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("second connection failed"); + wait_for_log(&test, "connected", CONNECTION_B); + assert!(!first.is_still_served().await, "A should have been superseded"); + + let mut third = TestConnection::open(&test, CONNECTION_C, Some(SESSION)) + .await + .expect("third connection failed"); + wait_for_log(&test, "connected", CONNECTION_C); + assert!(!second.is_still_served().await, "B should have been superseded"); + + let lines = lifecycle_log(&test); + assert!( + position_of(&lines, "disconnected", CONNECTION_B).is_some(), + "B should have been superseded by C: {lines:?}" + ); + assert!( + position_of(&lines, "disconnected", CONNECTION_C).is_none(), + "C should still be live: {lines:?}" + ); + + assert!( + third.is_still_served().await, + "the newest connection should still be served" + ); + + // Exactly one websocket client row remains for the session. The SQL + // query itself opens a short-lived connection, so allow for one extra. + let sql_out = test.sql("SELECT * FROM st_client").unwrap(); + let row_count = sql_out.lines().filter(|line| line.contains("0x")).count(); + assert!( + row_count <= 2, + "expected at most 2 st_client rows (the live connection and the SQL query's own), got {row_count}: {sql_out}" + ); + }); +} + +/// A batch subscribe registers every query set atomically and answers with one +/// result per set, in request order. +#[test] +fn test_batch_subscribe_applies_all_sets() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let mut connection = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("connection failed"); + + connection + .send(ws_v2::ClientMessage::SubscribeBatch(ws_v2::SubscribeBatch { + request_id: 1, + sets: vec![ + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(1), + query_strings: vec!["SELECT * FROM st_client".into()].into_boxed_slice(), + }, + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(2), + query_strings: vec!["SELECT * FROM st_table".into()].into_boxed_slice(), + }, + ] + .into_boxed_slice(), + })) + .await + .expect("failed to send SubscribeBatch"); + + match connection.next_message().await.expect("no response") { + ws_v2::ServerMessage::SubscribeBatchApplied(applied) => { + assert_eq!(applied.request_id, 1); + assert_eq!(applied.results.len(), 2, "expected one result per set"); + assert_eq!(applied.results[0].query_set_id, ws_v2::QuerySetId::new(1)); + assert_eq!(applied.results[1].query_set_id, ws_v2::QuerySetId::new(2)); + for result in applied.results.iter() { + assert!( + matches!(result.outcome, ws_v2::SubscribeSetOutcome::Applied(_)), + "expected every set to apply, got {:?}", + result.outcome + ); + } + } + other => panic!("expected SubscribeBatchApplied, got {other:?}"), + } + }); +} + +/// A batch subscribe with one invalid query reports that set's error while the +/// other sets still apply. +#[test] +fn test_batch_subscribe_reports_per_set_errors() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let mut connection = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("connection failed"); + + connection + .send(ws_v2::ClientMessage::SubscribeBatch(ws_v2::SubscribeBatch { + request_id: 7, + sets: vec![ + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(1), + query_strings: vec!["SELECT * FROM st_client".into()].into_boxed_slice(), + }, + ws_v2::SubscribeSet { + query_set_id: ws_v2::QuerySetId::new(2), + query_strings: vec!["SELECT * FROM no_such_table".into()].into_boxed_slice(), + }, + ] + .into_boxed_slice(), + })) + .await + .expect("failed to send SubscribeBatch"); + + match connection.next_message().await.expect("no response") { + ws_v2::ServerMessage::SubscribeBatchApplied(applied) => { + assert_eq!(applied.request_id, 7); + assert!( + matches!(applied.results[0].outcome, ws_v2::SubscribeSetOutcome::Applied(_)), + "the valid set should apply, got {:?}", + applied.results[0].outcome + ); + assert!( + matches!(applied.results[1].outcome, ws_v2::SubscribeSetOutcome::Error(_)), + "the invalid set should report an error, got {:?}", + applied.results[1].outcome + ); + } + other => panic!("expected SubscribeBatchApplied, got {other:?}"), + } + }); +} diff --git a/sdks/rust/src/db_connection.rs b/sdks/rust/src/db_connection.rs index 332aac1b322..137166ae614 100644 --- a/sdks/rust/src/db_connection.rs +++ b/sdks/rust/src/db_connection.rs @@ -1478,6 +1478,11 @@ async fn parse_loop( query_set_id: e.query_set_id, error: e.error.to_string(), }, + // This SDK negotiates v2 and never sends `SubscribeBatch`, + // so the server should never send this response. + ws::v2::ServerMessage::SubscribeBatchApplied(_) => ParsedMessage::Error( + InternalError::new("Received SubscribeBatchApplied, which this client never requests").into(), + ), ws::v2::ServerMessage::ProcedureResult(procedure_result) => ParsedMessage::ProcedureResult { request_id: procedure_result.request_id, result: match procedure_result.status { From 7803b6a5658f848de53575448cd8ab2a9f9a391b Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 26 Aug 2026 10:51:59 +0200 Subject: [PATCH 02/14] Cargo.lock --- Cargo.lock | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/Cargo.lock b/Cargo.lock index 22cd022fb09..0337962e4a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -210,6 +210,21 @@ version = "0.7.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" +[[package]] +name = "assert_cmd" +version = "2.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2aa3a22042e45de04255c7bf3626e239f450200fd0493c1e382263544b20aea6" +dependencies = [ + "anstyle", + "bstr", + "libc", + "predicates", + "predicates-core", + "predicates-tree", + "wait-timeout", +] + [[package]] name = "async-channel" version = "2.5.0" @@ -619,6 +634,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab" dependencies = [ "memchr", + "regex-automata", "serde", ] @@ -5614,6 +5630,16 @@ version = "1.0.9" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "727e462b119fe9c93fd0eb1429a5f7647394014cf3c04ab2c0350eeb09095ffa" +[[package]] +name = "predicates-tree" +version = "1.0.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0de1b847b39c8131db0467e9df1ff60e6d0562ab8e9a16e568ad0fdb372e2f2" +dependencies = [ + "predicates-core", + "termtree", +] + [[package]] name = "pretty_assertions" version = "1.4.1" @@ -8453,17 +8479,23 @@ name = "spacetimedb-smoketests" version = "2.8.3" dependencies = [ "anyhow", + "assert_cmd", "cargo_metadata", "fs_extra", + "futures", "predicates", "regex", "reqwest 0.12.24", "serde_json", "socket2 0.5.10", + "spacetimedb-client-api-messages", + "spacetimedb-core", "spacetimedb-guard", + "spacetimedb-lib", "tempfile", "tokio", "tokio-postgres", + "tokio-tungstenite 0.27.0", "toml 0.8.23", "which 8.0.0", "xmltree", From 508db01025c49c450714cd7dc14164a431b44121 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Tue, 1 Sep 2026 11:44:52 +0200 Subject: [PATCH 03/14] Refactor to use the same function between subscribe and subscribe batch --- .../subscription/module_subscription_actor.rs | 437 ++++++++++-------- 1 file changed, 257 insertions(+), 180 deletions(-) diff --git a/crates/core/src/subscription/module_subscription_actor.rs b/crates/core/src/subscription/module_subscription_actor.rs index 88ab8b22b5e..3df709378ef 100644 --- a/crates/core/src/subscription/module_subscription_actor.rs +++ b/crates/core/src/subscription/module_subscription_actor.rs @@ -45,6 +45,7 @@ use spacetimedb_physical_plan::plan::ProjectPlan; use spacetimedb_schema::def::RawModuleDefVersion; use spacetimedb_table::static_assert_size; use std::{ + ops::Range, sync::{ atomic::{AtomicU8, Ordering}, Arc, @@ -224,10 +225,26 @@ struct CompiledQueryBatch { struct CompiledQueries { queries: Vec>, physical_plans: HashMap>, - auth: AuthCtx, compile_timer: HistogramTimer, } +/// The result of [`ModuleSubscriptions::subscribe_query_sets`]. +struct SubscribedQuerySets { + /// The outcome of each query set, in the order the sets were requested. + outcomes: Vec, + /// The transaction the applied sets were evaluated at, or `None` if no set + /// applied. The caller is expected to hold this until it has enqueued its + /// response. + tx: Option>, + /// The offset of the transaction the applied sets were evaluated at, + /// or `None` if no set applied. + tx_offset: Option, + /// The metrics of evaluating every applied set. + metrics: ExecutionMetrics, + /// Whether materializing the subscribed views trapped. + trapped: bool, +} + /// The queries of a subscribe message, hashed by [`hash_queries`] /// for compilation cache lookup. struct HashedQueries<'a> { @@ -1124,9 +1141,8 @@ impl ModuleSubscriptions { let CompiledQueries { queries, physical_plans, - auth, compile_timer, - } = self.compile_hashed_queries(hashed, auth, metrics, &mut_tx)?; + } = self.compile_hashed_queries(hashed, &auth, metrics, &mut_tx)?; Ok(CompiledQueryBatch { queries, @@ -1141,7 +1157,7 @@ impl ModuleSubscriptions { fn compile_hashed_queries( &self, hashed: HashedQueries<'_>, - auth: AuthCtx, + auth: &AuthCtx, metrics: &SubscriptionMetrics, mut_tx: &MutTxId, ) -> Result { @@ -1170,7 +1186,7 @@ impl ModuleSubscriptions { |relational_db, tx| relational_db.get_all_tables_mut(tx).map(|schemas| schemas.into_iter()), &self.relational_db, mut_tx, - &auth, + auth, )? { add_compiled_query( compiled, @@ -1198,7 +1214,7 @@ impl ModuleSubscriptions { plans.push(unit); } _ => { - let compiled = compile_query_with_hashes(&auth, mut_tx, sql, hash, hash_with_param) + let compiled = compile_query_with_hashes(auth, mut_tx, sql, hash, hash_with_param) .map_err(|err| DBError::WithSql { error: Box::new(DBError::Other(err.into())), sql: sql.into(), @@ -1223,7 +1239,6 @@ impl ModuleSubscriptions { Ok(CompiledQueries { queries: plans, physical_plans, - auth, compile_timer, }) } @@ -1466,95 +1481,51 @@ impl ModuleSubscriptions { _timer: Instant, _assert: Option, ) -> Result<(Option, bool), DBError> { - // Send an error message to the client - // TODO: update for v2 - let send_err_msg = |message| { - let _ = self.broadcast_queue.send_client_message_v2( - sender.clone(), - None, - ws_v2::SubscriptionError { - request_id: Some(request.request_id), - query_set_id: request.query_set_id, - error: message, - }, - ); - }; - let subscription_metrics = &self.metrics.subscribe; - let num_queries = request.query_strings.len(); - subscription_metrics.num_queries_subscribed.inc_by(num_queries as _); - - let CompiledQueryBatch { - queries, - physical_plans, - auth, - mut_tx, - compile_timer: _compile_timer, - } = return_on_err!( - self.compile_queries( - sender.id.identity, - auth, - &request.query_strings, - num_queries, - subscription_metrics - ), - send_err_msg, - (None, false) - ); - let (mut_tx, _) = self.guard_mut_tx(mut_tx, <_>::default()); - - // We minimize locking so that other clients can add subscriptions concurrently. - // We are protected from race conditions with broadcasts, because we have the db lock, - // an `commit_and_broadcast_event` grabs a read lock on `subscriptions` while it still has a - // write lock on the db. - let queries = { - let mut subscriptions = { - // How contended is the lock? - let _wait_guard = subscription_metrics.lock_waiters.inc_scope(); - let _wait_timer = subscription_metrics.lock_wait_time.start_timer(); - self.subscriptions.write() - }; - - subscriptions.add_subscription_v2(sender.clone(), queries, request.query_set_id)? - }; - - let mut_tx = ScopeGuard::::into_inner(mut_tx); - - let (mut tx, tx_offset, trapped) = - self.materialize_views_and_downgrade_tx(mut_tx, instance, &queries, auth.caller())?; - - let failed_subscription = FailedSubscription::V2(request.query_set_id); - if let Err(err) = self.check_new_query_row_limit(&queries, &physical_plans, &tx, &auth) { - self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; - send_err_msg(err.to_string().into()); - return Ok((None, trapped)); - } - - let Ok((update, metrics)) = self.evaluate_queries(sender.clone(), &queries, &tx, TableUpdateType::Subscribe) - else { - self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; - send_err_msg("Internal error evaluating queries".into()); - return Ok((None, trapped)); - }; - tx.metrics.merge(metrics); - - subscription_metrics.num_queries_evaluated.inc_by(queries.len() as _); - - let ws_v2::QueryRows { tables } = match update { - ws_v1::FormatSwitch::Bsatn(update) => query_rows_from_update(update, false)?, - ws_v1::FormatSwitch::Json(_) => { - return Err(DBError::Other(anyhow::anyhow!( - "v2 subscriptions require binary protocol" - ))) + let ws_v2::Subscribe { + request_id, + query_set_id, + query_strings, + } = request; + let sets = [ws_v2::SubscribeSet { + query_set_id, + query_strings, + }]; + + let SubscribedQuerySets { + outcomes, + // Held until the response below has been enqueued. + tx: _tx, + tx_offset, + metrics, + trapped, + } = self.subscribe_query_sets(instance, &sender, auth, &sets)?; + let outcome = outcomes.into_iter().next().expect("one outcome for the set"); + + let rows = match outcome { + ws_v2::SubscribeSetOutcome::Applied(rows) => rows, + // Send an error message to the client + // TODO: update for v2 + ws_v2::SubscribeSetOutcome::Error(error) => { + let _ = self.broadcast_queue.send_client_message_v2( + sender.clone(), + None, + ws_v2::SubscriptionError { + request_id: Some(request_id), + query_set_id, + error, + }, + ); + return Ok((None, trapped)); } }; let _ = self.broadcast_queue.send_client_message_v2( sender.clone(), - Some(tx_offset), + tx_offset, ws_v2::SubscribeApplied { - request_id: request.request_id, - query_set_id: request.query_set_id, - rows: ws_v2::QueryRows { tables }, + request_id, + query_set_id, + rows, }, ); @@ -1563,13 +1534,9 @@ impl ModuleSubscriptions { /// Implementation of [`Self::add_batch_subscription`]. /// - /// Each set is compiled and evaluated in the same way as an individual - /// `Subscribe` ([`Self::compile_queries`], [`Self::check_new_query_row_limit`], - /// [`Self::evaluate_queries`]). - /// All subscription queries are evaluated at a single transaction snapshot - /// and answered in one response. A set which fails compilation, - /// the row limit, or evaluation is reported as an error in the response and - /// is not registered while the remaining sets still apply. + /// The whole batch is subscribed to by [`Self::subscribe_query_sets`], + /// and answered by a single [`ws_v2::SubscribeBatchApplied`], + /// so that nothing interleaves with the response. fn add_batch_subscription_inner( &self, instance: Option<&mut RefInstance<'_, I>>, @@ -1579,107 +1546,170 @@ impl ModuleSubscriptions { _timer: Instant, _assert: Option, ) -> Result<(Option, bool), DBError> { - let subscription_metrics = &self.metrics.subscribe; + let ws_v2::SubscribeBatch { request_id, sets } = request; - // The per-set outcome, in request order. Sets which fail before - // evaluation are filled in here and skipped later. - let mut outcomes: Vec> = (0..request.sets.len()).map(|_| None).collect(); - // The sets which compiled, as (request index, query set id, plans). - let mut compiled_sets: Vec<(usize, ws_v2::QuerySetId, Vec>)> = Vec::new(); - let mut physical_plans: HashMap> = HashMap::default(); + let SubscribedQuerySets { + outcomes, + // Held until the response below has been enqueued. + tx: _tx, + tx_offset, + metrics, + trapped, + } = self.subscribe_query_sets(instance, &sender, auth, &sets)?; + + let results = sets + .iter() + .zip(outcomes) + .map(|(set, outcome)| ws_v2::SubscribeSetResult { + query_set_id: set.query_set_id, + outcome, + }) + .collect(); - let num_queries: usize = request.sets.iter().map(|set| set.query_strings.len()).sum(); + let _ = self.broadcast_queue.send_client_message_v2( + sender.clone(), + tx_offset, + ws_v2::SubscribeBatchApplied { request_id, results }, + ); + + Ok((Some(metrics), trapped)) + } + + /// Subscribe `sender` to each of `sets`, returning an outcome per set. + /// + /// Every set is compiled, registered and evaluated within a single + /// transaction, and all of them are registered under a single tx lock. + /// + /// A set which fails to compile, fails to register, exceeds the row limit, + /// or fails to evaluate is reported as a [`ws_v2::SubscribeSetOutcome::Error`] + /// and is not registered, while the remaining sets still apply. + /// An `Err` is only returned for a failure of the request as a whole. + fn subscribe_query_sets( + &self, + instance: Option<&mut RefInstance<'_, I>>, + sender: &Arc, + auth: AuthCtx, + sets: &[ws_v2::SubscribeSet], + ) -> Result, DBError> { + let subscription_metrics = &self.metrics.subscribe; + + let num_queries: usize = sets.iter().map(|set| set.query_strings.len()).sum(); subscription_metrics.num_queries_subscribed.inc_by(num_queries as _); - // Hash every set's queries before taking the db lock, - // for the reasons described in [`Self::compile_queries`]. - let hashed_sets: Vec> = request - .sets + // We hash queries to avoid recompilation + let hashed_sets = sets .iter() .map(|set| hash_queries(sender.id.identity, &set.query_strings, set.query_strings.len())) - .collect(); + .collect::>(); + + // The outcome of each set, in the order of `sets`. + // Each stage below records the outcome of the sets which fail in it, + // and those sets are skipped by the later stages. + // The initial value is only ever observed if a stage fails to do so. + let mut outcomes = sets + .iter() + .map(|_| ws_v2::SubscribeSetOutcome::Error("Internal error subscribing to query set".into())) + .collect::>(); // We always get the db lock before the subscription lock to avoid deadlocks. // - // A single transaction spans the compilation, registration, and - // evaluation of every set, exactly as an individual `Subscribe` holds - // one transaction end-to-end. This means no schema change can - // invalidate a compiled set before it is registered, at the cost of - // holding the db lock across all compilations. + // A single transaction spans the compilation, registration and evaluation + // of every set, so no schema change can invalidate a compiled set before + // it is registered, at the cost of holding the db lock across all compilations. let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe); - // Compile every set. A set which fails to compile is reported in its - // outcome and does not fail the batch. + // Compile every set. A set which fails to compile does not fail the others. + let mut physical_plans: HashMap> = HashMap::default(); + let mut compiled_sets = Vec::with_capacity(sets.len()); for (index, hashed) in hashed_sets.into_iter().enumerate() { - match self.compile_hashed_queries(hashed, auth.clone(), subscription_metrics, &mut_tx) { + match self.compile_hashed_queries(hashed, &auth, subscription_metrics, &mut_tx) { Ok(CompiledQueries { queries, physical_plans: set_physical_plans, - auth: _, compile_timer: _compile_timer, }) => { physical_plans.extend(set_physical_plans); - compiled_sets.push((index, request.sets[index].query_set_id, queries)); - } - Err(err) => { - outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + compiled_sets.push((index, queries)); } + Err(err) => outcomes[index] = ws_v2::SubscribeSetOutcome::Error(err.to_string().into()), } } - // Register every compiled set under a single write lock, so no - // transaction committed between registrations can be observed by some - // sets but not others. - let mut registered: Vec<(usize, ws_v2::QuerySetId, Vec>)> = Vec::new(); + // Register every compiled set under a single write lock, so that no + // transaction committed between two registrations can be observed by + // some sets but not others. + // + // We minimize locking so that other clients can add subscriptions concurrently. + // We are protected from race conditions with broadcasts, because we have the db lock, + // and `commit_and_broadcast_event` grabs a read lock on `subscriptions` while it still + // has a write lock on the db. + // + // The registered queries of all sets are stored contiguously, + // each set holding the range of `registered_queries` which is its own. + let mut registered_queries: Vec> = Vec::with_capacity(num_queries); + let mut registered: Vec<(usize, Range)> = Vec::with_capacity(compiled_sets.len()); { let mut subscriptions = { + // How contended is the lock? let _wait_guard = subscription_metrics.lock_waiters.inc_scope(); let _wait_timer = subscription_metrics.lock_wait_time.start_timer(); self.subscriptions.write() }; - for (index, query_set_id, queries) in compiled_sets { - match subscriptions.add_subscription_v2(sender.clone(), queries.clone(), query_set_id) { - Ok(_) => registered.push((index, query_set_id, queries)), - Err(err) => { - outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + for (index, queries) in compiled_sets { + match subscriptions.add_subscription_v2(sender.clone(), queries, sets[index].query_set_id) { + // Note that we evaluate the queries returned by the subscription manager, + // as those are the ones it deduplicated and registered. + Ok(queries) => { + let start = registered_queries.len(); + registered_queries.extend(queries); + registered.push((index, start..registered_queries.len())); } + Err(err) => outcomes[index] = ws_v2::SubscribeSetOutcome::Error(err.to_string().into()), } } } + if registered.is_empty() { + // No set was registered, so there is nothing to evaluate, + // and no snapshot to evaluate it at. + // No update can concern a set which is not registered, + // so the caller's response needs no transaction to order it. + // The mutable transaction is committed when `mut_tx` is dropped. + return Ok(SubscribedQuerySets { + outcomes, + tx: None, + tx_offset: None, + metrics: ExecutionMetrics::default(), + trapped: false, + }); + } + let mut_tx = ScopeGuard::::into_inner(mut_tx); - let all_queries: Vec> = registered - .iter() - .flat_map(|(_, _, queries)| queries.iter().cloned()) - .collect(); let (mut tx, tx_offset, trapped) = - self.materialize_views_and_downgrade_tx(mut_tx, instance, &all_queries, auth.caller())?; + self.materialize_views_and_downgrade_tx(mut_tx, instance, ®istered_queries, auth.caller())?; // Evaluate every registered set against the single snapshot above. - // This will do the same row-limit check and evaluation of an individual - // `Subscribe` (see [`Self::add_v2_subscription_inner`]), except that a - // failure is recorded as the set's outcome instead of aborting the - // request, and the registration is removed so the failed set never - // receives transaction updates. + // A set which fails has its registration removed, + // so that it never receives transaction updates. let mut total_metrics = ExecutionMetrics::default(); - for (index, query_set_id, queries) in registered { - let failed_subscription = FailedSubscription::V2(query_set_id); - if let Err(err) = self.check_new_query_row_limit(&queries, &physical_plans, &tx, &auth) { + for (index, range) in registered { + let queries = ®istered_queries[range]; + let failed_subscription = FailedSubscription::V2(sets[index].query_set_id); + + if let Err(err) = self.check_new_query_row_limit(queries, &physical_plans, &tx, &auth) { self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; - outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + outcomes[index] = ws_v2::SubscribeSetOutcome::Error(err.to_string().into()); continue; } - let Ok((update, metrics)) = - self.evaluate_queries(sender.clone(), &queries, &tx, TableUpdateType::Subscribe) + let Ok((update, metrics)) = self.evaluate_queries(sender.clone(), queries, &tx, TableUpdateType::Subscribe) else { self.remove_failed_subscription(subscription_metrics, sender.id, failed_subscription)?; - outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error( - "Internal error evaluating queries".into(), - )); + outcomes[index] = ws_v2::SubscribeSetOutcome::Error("Internal error evaluating queries".into()); continue; }; tx.metrics.merge(metrics); + total_metrics.merge(metrics); subscription_metrics.num_queries_evaluated.inc_by(queries.len() as _); @@ -1691,33 +1721,16 @@ impl ModuleSubscriptions { ))) } }; - total_metrics.merge(metrics); - outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Applied(rows)); + outcomes[index] = ws_v2::SubscribeSetOutcome::Applied(rows); } - // One response for the whole batch, so nothing interleaves with it. - let results = outcomes - .into_iter() - .zip(request.sets.iter()) - .map(|(outcome, set)| ws_v2::SubscribeSetResult { - query_set_id: set.query_set_id, - outcome: outcome.unwrap_or_else(|| { - ws_v2::SubscribeSetOutcome::Error("Internal error registering query set".into()) - }), - }) - .collect::>() - .into_boxed_slice(); - - let _ = self.broadcast_queue.send_client_message_v2( - sender.clone(), - Some(tx_offset), - ws_v2::SubscribeBatchApplied { - request_id: request.request_id, - results, - }, - ); - - Ok((Some(total_metrics), trapped)) + Ok(SubscribedQuerySets { + outcomes, + tx: Some(tx), + tx_offset: Some(tx_offset), + metrics: total_metrics, + trapped, + }) } fn add_multi_subscription_inner( @@ -2189,13 +2202,15 @@ impl ModuleSubscriptions { /// Materialize the views returned by the `view_collector`, if not already materialized, /// and subsequently downgrade to a read-only transaction. #[allow(clippy::type_complexity)] - fn materialize_views_and_downgrade_tx( - &self, + // The returned guard only borrows `self`, so it may outlive the borrows of + // `instance` and `view_collector`, which `use<..>` keeps out of its type. + fn materialize_views_and_downgrade_tx<'a, I: WasmInstance, V: CollectViews>( + &'a self, mut tx: MutTxId, instance: Option<&mut RefInstance<'_, I>>, - view_collector: &impl CollectViews, + view_collector: &V, sender: Identity, - ) -> Result<(TxGuard, TransactionOffset, bool), DBError> { + ) -> Result<(TxGuard>, TransactionOffset, bool), DBError> { let mut trapped = false; if let Some(instance) = instance { (tx, trapped) = ModuleHost::materialize_views(tx, instance, view_collector, sender, Workload::Subscribe)?; @@ -2778,6 +2793,68 @@ mod tests { Ok(()) } + /// Test that a failed v2 subscription is answered with an error message, + /// and that its query set id is left free to re-use. + #[tokio::test] + async fn subscribe_v2_error() -> anyhow::Result<()> { + let db = relational_db()?; + + let client_id = client_id_from_u8(1); + let (sender, mut rx) = v2_client_connection(client_id, &db); + + let auth = AuthCtx::new(db.owner_identity(), client_id.identity); + let subs = ModuleSubscriptions::for_test_enclosing_runtime(db.clone()); + + db.create_table_for_test("t", &[("x", AlgebraicType::U8)], &[])?; + + // Subscribe to an invalid query (r is not in scope). + subs.add_v2_subscription_inner::( + None, + sender.clone(), + auth.clone(), + ws_v2::Subscribe { + request_id: 1, + query_set_id: ws_v2::QuerySetId::new(1), + query_strings: ["select r.* from t".into()].into(), + }, + Instant::now(), + None, + )?; + + match rx.recv().await { + Some(OutboundMessage::V2(ws_v2::ServerMessage::SubscriptionError(msg))) => { + assert_eq!(msg.request_id, Some(1)); + assert_eq!(msg.query_set_id, ws_v2::QuerySetId::new(1)); + } + other => panic!("Expected v2 SubscriptionError, got: {other:?}"), + } + + // The failed subscription was not registered, + // so the same query set id can be used again. + subs.add_v2_subscription_inner::( + None, + sender.clone(), + auth, + ws_v2::Subscribe { + request_id: 2, + query_set_id: ws_v2::QuerySetId::new(1), + query_strings: ["select * from t".into()].into(), + }, + Instant::now(), + None, + )?; + + match rx.recv().await { + Some(OutboundMessage::V2(ws_v2::ServerMessage::SubscribeApplied(msg))) => { + assert_eq!(msg.request_id, 2); + assert_eq!(msg.query_set_id, ws_v2::QuerySetId::new(1)); + } + other => panic!("Expected v2 SubscribeApplied, got: {other:?}"), + } + + Ok(()) + } + /// Test that a batch subscription answers all sets in one message, /// applies and registers the valid sets, /// and reports an invalid set's error without failing the batch. From c6f1780228e0a4ba98e9fc171c8779dcea49fb8d Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Tue, 1 Sep 2026 13:28:20 +0200 Subject: [PATCH 04/14] Update doc comment wording --- crates/client-api-messages/src/websocket/v2.rs | 4 ++-- .../src/subscription/module_subscription_actor.rs | 15 +++++---------- 2 files changed, 7 insertions(+), 12 deletions(-) diff --git a/crates/client-api-messages/src/websocket/v2.rs b/crates/client-api-messages/src/websocket/v2.rs index c203d4de702..e56abc1a36f 100644 --- a/crates/client-api-messages/src/websocket/v2.rs +++ b/crates/client-api-messages/src/websocket/v2.rs @@ -97,7 +97,7 @@ pub enum UnsubscribeFlags { /// Sent by client to register multiple subscriptions in one atomic step. /// /// The server registers every subscription set under a single subscription-manager -/// lock and evaluates all of them at a single transaction snapshot, +/// lock and evaluates all of them at a single transaction offset, /// then responds with one [`SubscribeBatchApplied`] message carrying a result per set. /// No [`TransactionUpdate`] is delivered between the registration of the first set /// and the [`SubscribeBatchApplied`] response, @@ -335,7 +335,7 @@ pub struct SubscriptionError { /// This message's `request_id` matches the one the client provided in the [`SubscribeBatch`] message, /// and `results` contains exactly one entry per received [`SubscribeSet`], in the same order. /// -/// Every applied set's rows are evaluated at a single transaction snapshot. +/// Every applied set's rows are evaluated at the same transaction offset. #[derive(SpacetimeType, Debug)] #[sats(crate = spacetimedb_lib)] pub struct SubscribeBatchApplied { diff --git a/crates/core/src/subscription/module_subscription_actor.rs b/crates/core/src/subscription/module_subscription_actor.rs index 3df709378ef..f5ba88774b2 100644 --- a/crates/core/src/subscription/module_subscription_actor.rs +++ b/crates/core/src/subscription/module_subscription_actor.rs @@ -1379,9 +1379,9 @@ impl ModuleSubscriptions { /// [`ws_v2::SubscribeBatch`] message. /// /// Every set is registered under a single subscription-manager lock and - /// evaluated at a single transaction snapshot, so no transaction update + /// evaluated at a single transaction offset, so no transaction update /// for any of the new sets can precede the [`ws_v2::SubscribeBatchApplied`] - /// response, and updates resume after it, all relative to the same snapshot. + /// response, and updates resume after it, all relative to that same offset. /// /// A set which fails to compile or evaluate reports a per-set error in the /// response while the remaining sets still apply. @@ -1614,8 +1614,7 @@ impl ModuleSubscriptions { // We always get the db lock before the subscription lock to avoid deadlocks. // // A single transaction spans the compilation, registration and evaluation - // of every set, so no schema change can invalidate a compiled set before - // it is registered, at the cost of holding the db lock across all compilations. + // of every set. let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe); // Compile every set. A set which fails to compile does not fail the others. @@ -1635,10 +1634,6 @@ impl ModuleSubscriptions { } } - // Register every compiled set under a single write lock, so that no - // transaction committed between two registrations can be observed by - // some sets but not others. - // // We minimize locking so that other clients can add subscriptions concurrently. // We are protected from race conditions with broadcasts, because we have the db lock, // and `commit_and_broadcast_event` grabs a read lock on `subscriptions` while it still @@ -1671,7 +1666,7 @@ impl ModuleSubscriptions { if registered.is_empty() { // No set was registered, so there is nothing to evaluate, - // and no snapshot to evaluate it at. + // and no transaction offset to evaluate it at. // No update can concern a set which is not registered, // so the caller's response needs no transaction to order it. // The mutable transaction is committed when `mut_tx` is dropped. @@ -1688,7 +1683,7 @@ impl ModuleSubscriptions { let (mut tx, tx_offset, trapped) = self.materialize_views_and_downgrade_tx(mut_tx, instance, ®istered_queries, auth.caller())?; - // Evaluate every registered set against the single snapshot above. + // Evaluate every registered set at the single transaction offset above. // A set which fails has its registration removed, // so that it never receives transaction updates. let mut total_metrics = ExecutionMetrics::default(); From 895f754ba33bc0ddc0954e556ad947a712815cef Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Thu, 3 Sep 2026 15:58:45 +0200 Subject: [PATCH 05/14] Serialize connection takeover --- crates/client-api/src/routes/subscribe.rs | 57 +- crates/core/src/client.rs | 2 +- .../core/src/client/client_session_index.rs | 488 +++++++++++++----- .../tests/cluster/connection_session.rs | 47 ++ 4 files changed, 448 insertions(+), 146 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 9370d2fab15..5f3d549d4eb 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -289,27 +289,27 @@ where log::debug!("websocket: New client connected from {client_log_string}"); // If this connection resumes a session which a live connection still - // holds, that connection is taken over. Stop its actor and run its - // module-side disconnect to completion, so the module observes - // `client_disconnected` for it strictly before `client_connected` for - // this one, and never two live connections for one session. - if let Some(session_id) = session_id - && let Some(superseded) = sessions.claim_session(client_id, session_id, None) - { - log::debug!( - "websocket: Connection {} supersedes {} for session {session_id}", - client_id.connection_id, - superseded.client_id.connection_id, - ); - if let Some(sender) = &superseded.sender { - sender.kick(ClientDisconnectCause::ConnectionSuperseded); + // holds, that connection is taken over: its actor is stopped and its + // module-side disconnect runs to completion before the claim returns. + // So the module observes `client_disconnected` for it strictly before + // `client_connected` for this one, and never two live connections for + // one session. + // + // The claim is held until this connection is established below, so a + // third connection resuming the same session must wait for this handover. + let session_claim = match session_id { + Some(session_id) => { + let module = module_rx.borrow().clone(); + Some( + sessions + .claim_session(db_identity, client_id, session_id, async |superseded| { + module.disconnect_client(superseded).await + }) + .await, + ) } - // Awaiting this is what orders the two lifecycle reducers: - // both run on the module's main instance, and this one is - // enqueued first. - let module = module_rx.borrow().clone(); - module.disconnect_client(superseded.client_id).await; - } + None => None, + }; let connected = match ClientConnection::call_client_connected_maybe_reject( &mut module_rx, @@ -340,11 +340,6 @@ where } }; record_client_rejection(db_identity, cause); - // The session claim is only meaningful for a connection which - // exists, so give it up again. - if let Some(session_id) = session_id { - sessions.release_session(client_id, session_id); - } return; } }; @@ -356,7 +351,9 @@ where // Release the session claim when the actor ends, including when it is aborted. let session_guard = session_id.map(|session_id| { let sessions = sessions.clone(); - scopeguard::guard((), move |()| sessions.release_session(client_id, session_id)) + scopeguard::guard((), move |()| { + sessions.release_session(db_identity, client_id, session_id) + }) }); let actor = |client, receiver| async move { let _session_guard = session_guard; @@ -374,10 +371,10 @@ where ) .await; - // Now that the actor exists, register its sender so that a later - // connection resuming this session can stop it. - if let Some(session_id) = session_id { - sessions.attach_sender(client_id, session_id, &client.sender()); + // Now that the actor exists, complete the handover by registering its + // sender, so that a later connection resuming this session can stop it. + if let Some(session_claim) = session_claim { + session_claim.attach_sender(&client.sender()); } // Send the client their identity token message as the first message diff --git a/crates/core/src/client.rs b/crates/core/src/client.rs index 5c90cba6175..8e14bf5324f 100644 --- a/crates/core/src/client.rs +++ b/crates/core/src/client.rs @@ -17,7 +17,7 @@ pub use client_connection::{ WsVersion, }; pub use client_connection_index::ClientActorIndex; -pub use client_session_index::{ClientSessionIndex, SessionId, SupersededConnection}; +pub use client_session_index::{ClientSessionIndex, SessionClaim, SessionId}; pub use message_handlers::MessageHandleError; pub use message_handlers_v1::MessageExecutionError; pub use messages::OutboundMessage; diff --git a/crates/core/src/client/client_session_index.rs b/crates/core/src/client/client_session_index.rs index ddc3eae3b7b..d303fdf2b02 100644 --- a/crates/core/src/client/client_session_index.rs +++ b/crates/core/src/client/client_session_index.rs @@ -13,11 +13,15 @@ //! `client_disconnected` could run after the new connection's //! `client_connected`. -use std::collections::hash_map::Entry; +use std::collections::hash_map::{Entry, OccupiedEntry}; use std::collections::HashMap; +use std::future::Future; use std::sync::{Arc, Mutex, Weak}; use spacetimedb_lib::Identity; +use tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard}; + +use crate::worker_metrics::ClientDisconnectCause; use super::{ClientActorId, ClientConnectionSender}; @@ -44,21 +48,36 @@ impl std::fmt::Display for SessionId { } } -/// A session is identified by the client's identity together with the -/// client-generated session id, so a session can only ever be replaced by the -/// client that owns it. -type SessionKey = (Identity, SessionId); +/// A session is identified by the database and the client's identity together +/// with the client-generated session id, so a session can only ever be replaced +/// by the same client on the same database. +#[derive(PartialEq, Eq, Clone, Copy, Hash, Debug)] +struct SessionKey { + database_identity: Identity, + client_identity: Identity, + session_id: SessionId, +} /// The connection currently serving a session. -struct SessionEntry { +struct SessionHolder { client_id: ClientActorId, /// Used to stop the connection's actor when it is superseded. /// /// Weak so that a connection whose actor has already ended can be dropped - /// normally rather than being kept alive by this map. + /// normally rather than being kept alive by this map. Empty until the + /// claiming connection is established, see [`SessionClaim::attach_sender`]. sender: Weak, } +/// One session's holder, behind the lock which serializes handovers of it. +/// +/// The lock is held for the whole of a handover: from the moment a connection +/// claims the session until that connection is established, or gives up. +/// A connection claiming a session whose handover is still in flight waits +/// here, so it never finds a holder whose connection does not exist yet and +/// therefore cannot be stopped. +type SessionSlot = AsyncMutex>; + /// The map of live sessions for one host. /// /// Maps each live session to the connection currently serving it. The entry is @@ -66,19 +85,7 @@ struct SessionEntry { /// connection. #[derive(Default)] pub struct ClientSessionIndex { - sessions: Mutex>, -} - -/// A connection which a newly arriving connection supersedes. -/// Returned by `ClientSessionIndex::claim_session`. -pub struct SupersededConnection { - /// The client actor matching a session id and identity. - /// Used by the caller to run the module side disconnect - /// lifecyle before allowing the new connection's `client_connected` to run. - pub client_id: ClientActorId, - /// The superseded connection's sender, if its actor is still alive. - /// Used to terminate the connection's websocket actor. - pub sender: Option>, + sessions: Mutex>>, } impl ClientSessionIndex { @@ -86,70 +93,111 @@ impl ClientSessionIndex { Self::default() } - /// Claim `session_id` for `client`, returning the connection it supersedes, - /// if any. + /// Claim a session for `client`, tearing down the connection it supersedes. /// - /// The caller must tear that connection down before allowing `client`'s - /// `client_connected` to run, so that the module never observes two live - /// connections for one session. The claim takes effect immediately, so a - /// third connection racing for the same session supersedes `client` rather - /// than the connection returned here. + /// Returns once that connection is fully closed: its actor has been + /// stopped, and `teardown`, which runs the module-side disconnect, has run + /// to completion. The caller may therefore run `client_connected` for + /// `client` as soon as this returns, and the module observes the two + /// connections' lifecycle events in order. /// - /// `sender` is registered so that a later connection can stop this - /// connection's actor. It is `None` before the connection's actor exists, - /// in which case the entry is registered without one. - pub fn claim_session( - &self, + /// The returned [`SessionClaim`] holds the session for the rest of the + /// handover. Another connection claiming the same session waits until the + /// claim is completed with [`SessionClaim::attach_sender`] or dropped, + /// so every claim finds a connection which it can actually stop. + pub async fn claim_session( + self: &Arc, + database_identity: Identity, client: ClientActorId, session_id: SessionId, - sender: Option<&Arc>, - ) -> Option { - let key = (client.identity, session_id); - let entry = SessionEntry { - client_id: client, - sender: sender.map(Arc::downgrade).unwrap_or_default(), + teardown: F, + ) -> SessionClaim + where + F: FnOnce(ClientActorId) -> Fut, + Fut: Future, + { + let key = SessionKey { + database_identity, + client_identity: client.identity, + session_id, }; - let mut sessions = self.sessions.lock().expect("session index poisoned"); - match sessions.entry(key) { - Entry::Occupied(mut occupied) => { - let superseded = occupied.insert(entry); - (superseded.client_id.connection_id != client.connection_id).then(|| SupersededConnection { - client_id: superseded.client_id, - sender: superseded.sender.upgrade(), - }) - } - Entry::Vacant(vacant) => { - vacant.insert(entry); - None + let slot = { + let mut sessions = self.sessions.lock().expect("session index poisoned"); + sessions.entry(key).or_default().clone() + }; + + // Wait out any handover of this session which is still in flight. + let mut holder = slot.lock_owned().await; + let superseded = holder.replace(SessionHolder { + client_id: client, + sender: Weak::new(), + }); + + // Connections are told apart by their name, the host's per-connection + // counter, rather than by their connection id, which a client may + // repeat across connections. + if let Some(superseded) = superseded.filter(|superseded| superseded.client_id.name != client.name) { + log::debug!( + "websocket: Connection {} supersedes {} for session {session_id}", + client.connection_id, + superseded.client_id.connection_id, + ); + if let Some(sender) = superseded.sender.upgrade() { + sender.kick(ClientDisconnectCause::ConnectionSuperseded); } + teardown(superseded.client_id).await; } - } - /// Record the sender for the connection currently holding `session_id`. - /// - /// Called once the connection's actor exists. Does nothing if the session - /// has already been claimed by a newer connection. - pub fn attach_sender(&self, client: ClientActorId, session_id: SessionId, sender: &Arc) { - let key = (client.identity, session_id); - let mut sessions = self.sessions.lock().expect("session index poisoned"); - if let Some(entry) = sessions.get_mut(&key) - && entry.client_id.connection_id == client.connection_id - { - entry.sender = Arc::downgrade(sender); + SessionClaim { + index: self.clone(), + key, + holder: Some(holder), } } - /// Release `session_id` if it is still held by `client`. + /// Release a session if it is still held by `client`. /// /// Called when a connection ends. A connection which has already been /// superseded no longer holds the session, so it leaves the entry alone: /// otherwise a slow teardown would evict its own replacement. - pub fn release_session(&self, client: ClientActorId, session_id: SessionId) { - let key = (client.identity, session_id); + pub fn release_session(&self, database_identity: Identity, client: ClientActorId, session_id: SessionId) { + let key = SessionKey { + database_identity, + client_identity: client.identity, + session_id, + }; let mut sessions = self.sessions.lock().expect("session index poisoned"); - if let Entry::Occupied(entry) = sessions.entry(key) - && entry.get().client_id.connection_id == client.connection_id - { + let Entry::Occupied(entry) = sessions.entry(key) else { + return; + }; + // This runs while a connection is being dropped, so it must not block. + // Failing to take the lock means a handover of this session is in + // flight: either another connection's, in which case this connection + // no longer holds the session and there is nothing to release, or this + // connection's own, whose actor ended before it was established. The + // latter leaves the entry behind, holding a sender which can no longer + // be upgraded, until the next claim of the session replaces it. + let Ok(mut holder) = entry.get().clone().try_lock_owned() else { + return; + }; + if holder.as_ref().is_some_and(|held| held.client_id.name == client.name) { + *holder = None; + } + let is_vacant = holder.is_none(); + drop(holder); + if is_vacant { + Self::prune(entry); + } + } + + /// Remove a session which no connection holds and none is claiming. + /// + /// A connection waiting to claim the session holds a reference to the slot, + /// so a strong count of one means the map holds the only reference and the + /// entry can go. Anything else would leave a claimant waiting on a slot no + /// longer reachable from the map, which a later claimant would not find. + fn prune(entry: OccupiedEntry<'_, SessionKey, Arc>) { + if Arc::strong_count(entry.get()) == 1 { entry.remove(); } } @@ -164,20 +212,100 @@ impl ClientSessionIndex { } } +/// A session held for the duration of one handover. +/// +/// Returned by [`ClientSessionIndex::claim_session`]. Completed by +/// [`SessionClaim::attach_sender`] once the claiming connection exists; +/// dropping it without that gives the session up again, for a connection which +/// never came to be. +pub struct SessionClaim { + index: Arc, + key: SessionKey, + /// `Some` until the claim is completed or dropped. + holder: Option>>, +} + +impl SessionClaim { + /// Complete the handover, recording the now-established connection's sender + /// so that a later connection can stop it. + /// + /// Holding the claim is what proves the session is still this connection's, + /// so no ownership check is needed here. + pub fn attach_sender(mut self, sender: &Arc) { + let mut holder = self.holder.take().expect("a claim is completed at most once"); + if let Some(held) = holder.as_mut() { + held.sender = Arc::downgrade(sender); + } + } +} + +impl Drop for SessionClaim { + fn drop(&mut self) { + // Completed claims took the guard in `attach_sender`, leaving the + // session held by the connection which is now serving it. + let Some(mut holder) = self.holder.take() else { + return; + }; + *holder = None; + drop(holder); + + let mut sessions = self.index.sessions.lock().expect("session index poisoned"); + if let Entry::Occupied(entry) = sessions.entry(self.key) { + ClientSessionIndex::prune(entry); + } + } +} + #[cfg(test)] mod tests { + use super::super::client_connection::DurableOffsetSupply; use super::*; - use crate::client::ClientName; + use crate::client::{ClientConfig, ClientName}; + use crate::host::module_host::NoSuchModule; + use spacetimedb_durability::DurableOffset; use spacetimedb_lib::ConnectionId; + use std::sync::atomic::{AtomicBool, Ordering}; + + /// The dummy senders below never wait on durability. + struct NoDurability; + + impl DurableOffsetSupply for NoDurability { + fn durable_offset(&mut self) -> Result, NoSuchModule> { + Ok(None) + } + } + + fn index() -> Arc { + Arc::new(ClientSessionIndex::new()) + } + /// A client id whose `name`, the host's per-connection counter, matches its + /// connection id, so that tests naming distinct connections get distinct + /// names as the websocket handler would assign them. fn client(identity: Identity, connection_id: u128) -> ClientActorId { ClientActorId { identity, connection_id: ConnectionId::from_u128(connection_id), - name: ClientName(0), + name: ClientName(connection_id as u64), } } + fn sender(client: ClientActorId) -> Arc { + Arc::new(ClientConnectionSender::dummy( + client, + ClientConfig::for_test(), + NoDurability, + )) + } + + fn a_database() -> Identity { + Identity::from_byte_array([9; 32]) + } + + fn another_database() -> Identity { + Identity::from_byte_array([8; 32]) + } + fn an_identity() -> Identity { Identity::from_byte_array([1; 32]) } @@ -186,103 +314,233 @@ mod tests { Identity::from_byte_array([2; 32]) } - #[test] - fn first_connection_supersedes_nothing() { - let index = ClientSessionIndex::new(); + /// Claim a session, recording which connection was torn down, if any. + async fn claim( + index: &Arc, + database: Identity, + client: ClientActorId, + session: SessionId, + ) -> (SessionClaim, Option) { + let mut superseded = None; + let claim = index + .claim_session(database, client, session, async |old| superseded = Some(old)) + .await; + (claim, superseded) + } + + /// Claim a session and complete the handover, as an established connection + /// does. + async fn connect( + index: &Arc, + database: Identity, + client: ClientActorId, + session: SessionId, + ) -> (Arc, Option) { + let (claim, superseded) = claim(index, database, client, session).await; + let sender = sender(client); + claim.attach_sender(&sender); + (sender, superseded) + } + + #[tokio::test] + async fn first_connection_supersedes_nothing() { + let index = index(); let session = SessionId::from_u128(7); - assert!(index.claim_session(client(an_identity(), 1), session, None).is_none()); + + let (_sender, superseded) = connect(&index, a_database(), client(an_identity(), 1), session).await; + + assert!(superseded.is_none()); assert_eq!(index.len(), 1); } - #[test] - fn reconnect_supersedes_previous_connection() { - let index = ClientSessionIndex::new(); + #[tokio::test] + async fn reconnect_supersedes_previous_connection() { + let index = index(); let session = SessionId::from_u128(7); - index.claim_session(client(an_identity(), 1), session, None); + let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; - let superseded = index.claim_session(client(an_identity(), 2), session, None); + let (_second, superseded) = connect(&index, a_database(), client(an_identity(), 2), session).await; assert_eq!( - superseded.map(|s| s.client_id.connection_id), + superseded.map(|old| old.connection_id), Some(ConnectionId::from_u128(1)) ); + // The superseded connection's actor is stopped. + assert!(first.is_cancelled()); // The session is now held by the new connection, not the old one. assert_eq!(index.len(), 1); } - #[test] - fn different_identity_does_not_supersede() { - let index = ClientSessionIndex::new(); + /// A client may repeat a connection id across connections, so the session + /// is handed over on the connection's name rather than on that id. + #[tokio::test] + async fn reconnect_reusing_connection_id_supersedes() { + let index = index(); let session = SessionId::from_u128(7); - index.claim_session(client(an_identity(), 1), session, None); + let first = ClientActorId { + name: ClientName(1), + ..client(an_identity(), 1) + }; + let second = ClientActorId { + name: ClientName(2), + ..client(an_identity(), 1) + }; + let (first_sender, _) = connect(&index, a_database(), first, session).await; - let superseded = index.claim_session(client(another_identity(), 2), session, None); + let (_second_sender, superseded) = connect(&index, a_database(), second, session).await; + + assert_eq!(superseded.map(|old| old.name), Some(ClientName(1))); + assert!(first_sender.is_cancelled()); + assert_eq!(index.len(), 1); + } + + #[tokio::test] + async fn different_identity_does_not_supersede() { + let index = index(); + let session = SessionId::from_u128(7); + let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; + + let (_second, superseded) = connect(&index, a_database(), client(another_identity(), 2), session).await; assert!(superseded.is_none()); assert_eq!(index.len(), 2); } - #[test] - fn different_session_does_not_supersede() { - let index = ClientSessionIndex::new(); - index.claim_session(client(an_identity(), 1), SessionId::from_u128(7), None); + #[tokio::test] + async fn different_session_does_not_supersede() { + let index = index(); + let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), SessionId::from_u128(7)).await; - let superseded = index.claim_session(client(an_identity(), 2), SessionId::from_u128(8), None); + let (_second, superseded) = + connect(&index, a_database(), client(an_identity(), 2), SessionId::from_u128(8)).await; assert!(superseded.is_none()); assert_eq!(index.len(), 2); } - #[test] - fn release_removes_the_session() { - let index = ClientSessionIndex::new(); + /// A session belongs to one database, so a connection to another database + /// never tears down this one, whose module knows nothing about it. + #[tokio::test] + async fn different_database_does_not_supersede() { + let index = index(); + let session = SessionId::from_u128(7); + let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; + + let (_second, superseded) = connect(&index, another_database(), client(an_identity(), 2), session).await; + + assert!(superseded.is_none()); + assert!(!first.is_cancelled()); + assert_eq!(index.len(), 2); + } + + #[tokio::test] + async fn release_removes_the_session() { + let index = index(); let session = SessionId::from_u128(7); let connection = client(an_identity(), 1); - index.claim_session(connection, session, None); + let (_sender, _) = connect(&index, a_database(), connection, session).await; - index.release_session(connection, session); + index.release_session(a_database(), connection, session); assert!(index.is_empty()); } - #[test] - fn superseded_connection_release_does_not_evict_its_replacement() { - let index = ClientSessionIndex::new(); + #[tokio::test] + async fn superseded_connection_release_does_not_evict_its_replacement() { + let index = index(); let session = SessionId::from_u128(7); let old = client(an_identity(), 1); let new = client(an_identity(), 2); - index.claim_session(old, session, None); - index.claim_session(new, session, None); + let (_old_sender, _) = connect(&index, a_database(), old, session).await; + let (_new_sender, _) = connect(&index, a_database(), new, session).await; // The old connection tears down after being superseded. - index.release_session(old, session); + index.release_session(a_database(), old, session); // The replacement still holds the session. assert_eq!(index.len(), 1); - assert_eq!( - index - .claim_session(client(an_identity(), 3), session, None) - .map(|s| s.client_id.connection_id), - Some(new.connection_id) - ); + let (_third, superseded) = connect(&index, a_database(), client(an_identity(), 3), session).await; + assert_eq!(superseded.map(|old| old.connection_id), Some(new.connection_id)); } - #[test] - fn three_way_race_supersedes_the_most_recent_connection() { - let index = ClientSessionIndex::new(); + /// A connection which never came to be, because `client_connected` + /// rejected it, gives the session up again. + #[tokio::test] + async fn dropped_claim_releases_the_session() { + let index = index(); let session = SessionId::from_u128(7); - index.claim_session(client(an_identity(), 1), session, None); - let second = index.claim_session(client(an_identity(), 2), session, None); - let third = index.claim_session(client(an_identity(), 3), session, None); + let (claim, _) = claim(&index, a_database(), client(an_identity(), 1), session).await; + drop(claim); + + assert!(index.is_empty()); + } + + #[tokio::test] + async fn three_way_race_supersedes_the_most_recent_connection() { + let index = index(); + let session = SessionId::from_u128(7); + let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; + + let (_second, superseded_by_second) = connect(&index, a_database(), client(an_identity(), 2), session).await; + let (_third, superseded_by_third) = connect(&index, a_database(), client(an_identity(), 3), session).await; assert_eq!( - second.map(|s| s.client_id.connection_id), + superseded_by_second.map(|old| old.connection_id), Some(ConnectionId::from_u128(1)) ); assert_eq!( - third.map(|s| s.client_id.connection_id), + superseded_by_third.map(|old| old.connection_id), Some(ConnectionId::from_u128(2)) ); } + + /// A claim waits for the handover in flight, so it never supersedes a + /// connection which does not exist yet and so cannot be stopped. + #[tokio::test] + async fn handover_serializes_concurrent_claims() { + let index = index(); + let session = SessionId::from_u128(7); + let first = client(an_identity(), 1); + let second = client(an_identity(), 2); + + // The first connection claims the session but is not established yet, + // as it would not be while its `client_connected` runs. + let (claim, _) = claim(&index, a_database(), first, session).await; + + let torn_down = Arc::new(Mutex::new(None)); + let started = Arc::new(AtomicBool::new(false)); + let claimed = Arc::new(AtomicBool::new(false)); + let racing = tokio::spawn({ + let (index, torn_down) = (index.clone(), torn_down.clone()); + let (started, claimed) = (started.clone(), claimed.clone()); + async move { + started.store(true, Ordering::Release); + let claim = index + .claim_session(a_database(), second, session, async |old| { + *torn_down.lock().unwrap() = Some(old); + }) + .await; + claimed.store(true, Ordering::Release); + claim.attach_sender(&sender(second)); + } + }); + + // The second connection cannot claim the session while the first + // connection's handover is still in flight. + tokio::task::yield_now().await; + assert!(started.load(Ordering::Acquire), "the racing claim never ran"); + assert!(!claimed.load(Ordering::Acquire), "the racing claim should be waiting"); + assert!(torn_down.lock().unwrap().is_none()); + + // Once the first connection is established, the second supersedes it, + // and finds a connection which it can stop. + let first_sender = sender(first); + claim.attach_sender(&first_sender); + racing.await.unwrap(); + + assert_eq!(torn_down.lock().unwrap().map(|old| old.name), Some(first.name)); + assert!(first_sender.is_cancelled()); + assert_eq!(index.len(), 1); + } } diff --git a/crates/smoketests/tests/cluster/connection_session.rs b/crates/smoketests/tests/cluster/connection_session.rs index d6048b2868d..1f87311f7de 100644 --- a/crates/smoketests/tests/cluster/connection_session.rs +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -315,6 +315,53 @@ fn test_repeated_reconnects_leave_one_live_connection() { }); } +/// A reconnect which repeats its predecessor's connection id still supersedes +/// it. Connections are told apart by the server, not by the id a client sends, +/// which a client is free to repeat. +#[test] +fn test_reconnect_reusing_connection_id_replaces_connection() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + runtime().block_on(async { + let mut first = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("first connection failed"); + wait_for_log(&test, "connected", CONNECTION_A); + + // Reconnect under the same connection id as well as the same session. + let mut second = TestConnection::open(&test, CONNECTION_A, Some(SESSION)) + .await + .expect("second connection failed"); + + assert!( + !first.is_still_served().await, + "the superseded connection should no longer be served" + ); + assert!( + second.is_still_served().await, + "the newest connection should still be served" + ); + + // Both connections log the same id, so count the events rather than + // ordering them: the first connection was torn down exactly once. + let lines = wait_for_log(&test, "disconnected", CONNECTION_A); + let disconnects = lines + .iter() + .filter(|line| line.contains(&format!("disconnected {CONNECTION_A}"))) + .count(); + assert_eq!(disconnects, 1, "expected exactly one teardown: {lines:?}"); + + // Exactly one websocket client row remains for the session. The SQL + // query itself opens a short-lived connection, so allow for one extra. + let sql_out = test.sql("SELECT * FROM st_client").unwrap(); + let row_count = sql_out.lines().filter(|line| line.contains("0x")).count(); + assert!( + row_count <= 2, + "expected at most 2 st_client rows (the live connection and the SQL query's own), got {row_count}: {sql_out}" + ); + }); +} + /// A batch subscribe registers every query set atomically and answers with one /// result per set, in request order. #[test] From 3f9d3285a3e677d928a068785e6d53dc1f7a53d1 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Mon, 7 Sep 2026 14:08:10 +0200 Subject: [PATCH 06/14] Refuse connection of existing session and start teardown instead --- .../src/websocket/common.rs | 8 + crates/client-api/src/routes/subscribe.rs | 96 ++-- crates/core/src/client.rs | 2 +- crates/core/src/client/client_connection.rs | 10 +- .../core/src/client/client_session_index.rs | 434 +++++------------- crates/core/src/worker_metrics/mod.rs | 5 + .../tests/cluster/connection_session.rs | 99 +++- 7 files changed, 263 insertions(+), 391 deletions(-) diff --git a/crates/client-api-messages/src/websocket/common.rs b/crates/client-api-messages/src/websocket/common.rs index 45e6b4b381f..688baa9eab5 100644 --- a/crates/client-api-messages/src/websocket/common.rs +++ b/crates/client-api-messages/src/websocket/common.rs @@ -53,6 +53,14 @@ pub const SERVER_MSG_COMPRESSION_TAG_BROTLI: u8 = 1; /// The tag recognized by the host and SDKs to mean gzip compression of a `ServerMessage`. pub const SERVER_MSG_COMPRESSION_TAG_GZIP: u8 = 2; +/// Websocket close code sent when a connection supplies a `session_id` +/// which another live connection of the same identity still holds. +/// +/// The server stops that connection and tears it down. The refused client +/// should retry after a short delay, without counting this as a failed +/// connection attempt for backoff purposes. +pub const SESSION_BUSY_CLOSE_CODE: u16 = 4000; + pub type RowSize = u16; pub type RowOffset = u64; diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 5f3d549d4eb..e1293a3997a 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -28,7 +28,8 @@ use spacetimedb::client::messages::{ }; use spacetimedb::client::{ ClientActorId, ClientConfig, ClientConnection, ClientConnectionReceiver, DataMessage, MessageExecutionError, - MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, SessionId, WsVersion, + MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, SessionBusy, SessionId, + SessionReservation, WsVersion, }; use spacetimedb::host::module_host::ClientConnectedError; use spacetimedb::host::NoSuchModule; @@ -38,6 +39,7 @@ use spacetimedb::worker_metrics::{ record_client_rejection, ClientDisconnectCause, ClientDisconnectRecorder, ClientRejectCause, WORKER_METRICS, }; use spacetimedb::Identity; +use spacetimedb_client_api_messages::websocket::common::SESSION_BUSY_CLOSE_CODE; use spacetimedb_client_api_messages::websocket::v1 as ws_v1; use spacetimedb_client_api_messages::websocket::v2 as ws_v2; use spacetimedb_client_api_messages::websocket::v3 as ws_v3; @@ -89,10 +91,11 @@ pub struct SubscribeQueryParams { /// A client-generated identifier for a logical client session, /// stable across the reconnects of one client connection object. /// - /// When a connection supplies a session id already held by a live - /// connection of the same identity, the old connection is torn down before - /// this one runs `client_connected`, so the module never observes two live - /// connections for one session. + /// When a connection supplies a session id still held by a live + /// connection of the same identity, this connection is refused with + /// [`SESSION_BUSY_CLOSE_CODE`] and the old one is torn down. A retry + /// succeeds once the old connection's `client_disconnected` has run, so the + /// module never observes two live connections for one session. /// See [`spacetimedb::client::ClientSessionIndex`]. /// /// Connections which do not supply one behave exactly as before. @@ -269,7 +272,7 @@ where let ws_opts = ctx.websocket_options(); tokio::spawn(async move { - let ws = match ws_upgrade.upgrade(ws_config).await { + let mut ws = match ws_upgrade.upgrade(ws_config).await { Ok(ws) => ws, Err(err) => { record_client_rejection(db_identity, ClientRejectCause::WebsocketUpgradeError); @@ -288,26 +291,25 @@ where log::debug!("websocket: New client connected from {client_log_string}"); - // If this connection resumes a session which a live connection still - // holds, that connection is taken over: its actor is stopped and its - // module-side disconnect runs to completion before the claim returns. - // So the module observes `client_disconnected` for it strictly before - // `client_connected` for this one, and never two live connections for - // one session. - // - // The claim is held until this connection is established below, so a - // third connection resuming the same session must wait for this handover. - let session_claim = match session_id { - Some(session_id) => { - let module = module_rx.borrow().clone(); - Some( - sessions - .claim_session(db_identity, client_id, session_id, async |superseded| { - module.disconnect_client(superseded).await - }) - .await, - ) - } + // Reserved before `client_connected` so that no two connections of one + // session run it. Released by the actor's teardown after the + // module-side disconnect, so a retry finds `client_disconnected` run. + let session = match session_id { + Some(session_id) => match sessions.try_reserve(db_identity, client_id, session_id) { + Ok(reservation) => Some(reservation), + Err(SessionBusy) => { + WORKER_METRICS.ws_clients_session_busy.with_label_values(&db_identity).inc(); + log::debug!("websocket: Refusing connection for {client_log_string}: session {session_id} is busy"); + let close = CloseFrame { + code: CloseCode::from(SESSION_BUSY_CLOSE_CODE), + reason: "session busy".into(), + }; + if let Err(e) = ws.close(Some(close)).await { + log::debug!("websocket: Error refusing connection for {client_log_string}: {e}"); + } + return; + } + }, None => None, }; @@ -348,16 +350,11 @@ where "websocket: Database accepted connection from {client_log_string}; spawning ws_client_actor and ClientConnection" ); - // Release the session claim when the actor ends, including when it is aborted. - let session_guard = session_id.map(|session_id| { - let sessions = sessions.clone(); - scopeguard::guard((), move |()| { - sessions.release_session(db_identity, client_id, session_id) - }) - }); - let actor = |client, receiver| async move { - let _session_guard = session_guard; - ws_client_actor(ws_opts, client, ws, receiver).await; + let actor = |client: ClientConnection, receiver| { + if let Some(session) = &session { + session.establish(&client.sender()); + } + ws_client_actor(ws_opts, client, ws, receiver, session) }; let client = ClientConnection::spawn( client_id, @@ -371,12 +368,6 @@ where ) .await; - // Now that the actor exists, complete the handover by registering its - // sender, so that a later connection resuming this session can stop it. - if let Some(session_claim) = session_claim { - session_claim.attach_sender(&client.sender()); - } - // Send the client their identity token message as the first message // NOTE: We're adding this to the protocol because some client libraries are // unable to access the http response headers. @@ -581,15 +572,26 @@ async fn ws_client_actor( client: ClientConnection, ws: WebSocketStream, sendrx: ClientConnectionReceiver, + session: Option, ) { - // ensure that even if this task gets cancelled, we always cleanup the connection - let mut client = scopeguard::guard(client, |client| { - tokio::spawn(client.disconnect()); + // Runs the module-side disconnect even if this task gets cancelled. + let mut client = scopeguard::guard((client, session), |(client, session)| { + tokio::spawn(ws_client_teardown(client, session)); }); - ws_client_actor_inner(&mut client, options, ws, sendrx).await; + ws_client_actor_inner(&mut client.0, options, ws, sendrx).await; - ScopeGuard::into_inner(client).disconnect().await; + let (client, session) = ScopeGuard::into_inner(client); + ws_client_teardown(client, session).await; +} + +/// Run the module-side disconnect, then free the connection's session. +/// +/// The session is released only after `client_disconnected` has run, so that +/// a connection retrying for the same session observes it in order. +async fn ws_client_teardown(client: ClientConnection, session: Option) { + client.disconnect().await; + drop(session); } async fn ws_client_actor_inner( diff --git a/crates/core/src/client.rs b/crates/core/src/client.rs index 8e14bf5324f..7281f72c0f9 100644 --- a/crates/core/src/client.rs +++ b/crates/core/src/client.rs @@ -17,7 +17,7 @@ pub use client_connection::{ WsVersion, }; pub use client_connection_index::ClientActorIndex; -pub use client_session_index::{ClientSessionIndex, SessionClaim, SessionId}; +pub use client_session_index::{ClientSessionIndex, SessionBusy, SessionId, SessionReservation}; pub use message_handlers::MessageHandleError; pub use message_handlers_v1::MessageExecutionError; pub use messages::OutboundMessage; diff --git a/crates/core/src/client/client_connection.rs b/crates/core/src/client/client_connection.rs index e22ba3616d1..821d303e601 100644 --- a/crates/core/src/client/client_connection.rs +++ b/crates/core/src/client/client_connection.rs @@ -410,14 +410,10 @@ impl ClientConnectionSender { /// Stop this connection's websocket actor. /// - /// Used when a newer connection supersedes this one + /// Used when a newer connection arrives for this connection's session /// (see [`super::ClientSessionIndex`]), and when a client exceeds its - /// outgoing queue capacity. - /// - /// This only stops the actor. The module-side disconnect - /// ([`crate::host::ModuleHost::disconnect_client`]) is run separately by - /// the actor's teardown, or by the caller when it needs that teardown to - /// complete before some other work. + /// outgoing queue capacity. The actor's teardown runs the module-side + /// disconnect. pub fn kick(&self, cause: ClientDisconnectCause) { if let Some(metrics) = &self.metrics { metrics.disconnect_recorder.record(cause); diff --git a/crates/core/src/client/client_session_index.rs b/crates/core/src/client/client_session_index.rs index d303fdf2b02..6d968d774f2 100644 --- a/crates/core/src/client/client_session_index.rs +++ b/crates/core/src/client/client_session_index.rs @@ -1,25 +1,25 @@ -//! Tracking of client sessions, used to replace pre-existing connections. +//! Tracking of live client sessions, so that at most one connection serves a +//! session at a time. //! //! A client which reconnects automatically sends the same client-generated //! session id on every connection attempt. Each connection still receives its //! own [`ConnectionId`] and its own `client_connected` / `client_disconnected` //! events. The session id only identifies which earlier connection a new one -//! supersedes. +//! replaces. //! -//! A client frequently notices a dropped connection before the server does -//! as the server needs up to its idle timeout to notice an idle peer. -//! Without this index the module would briefly observe two live -//! connections for the same client, and the old connection's -//! `client_disconnected` could run after the new connection's -//! `client_connected`. - -use std::collections::hash_map::{Entry, OccupiedEntry}; +//! A client frequently notices a dropped connection before the server does, +//! as the server needs up to its idle timeout to notice an idle peer. When a +//! connection arrives for a session which a live connection still holds, the +//! new connection is refused and the old one is stopped. The session frees up +//! once the old connection is fully closed, including its module-side +//! disconnect, so a retry finds `client_disconnected` already run. +//! +//! [`ConnectionId`]: spacetimedb_lib::ConnectionId + use std::collections::HashMap; -use std::future::Future; use std::sync::{Arc, Mutex, Weak}; use spacetimedb_lib::Identity; -use tokio::sync::{Mutex as AsyncMutex, OwnedMutexGuard}; use crate::worker_metrics::ClientDisconnectCause; @@ -48,9 +48,8 @@ impl std::fmt::Display for SessionId { } } -/// A session is identified by the database and the client's identity together -/// with the client-generated session id, so a session can only ever be replaced -/// by the same client on the same database. +/// A session belongs to one client on one database, so it can only ever be +/// replaced by the same client on the same database. #[derive(PartialEq, Eq, Clone, Copy, Hash, Debug)] struct SessionKey { database_identity: Identity, @@ -58,34 +57,24 @@ struct SessionKey { session_id: SessionId, } -/// The connection currently serving a session. +/// The connection holding a session. struct SessionHolder { - client_id: ClientActorId, - /// Used to stop the connection's actor when it is superseded. + client: ClientActorId, + /// Stops the connection's actor when a newer connection arrives. /// - /// Weak so that a connection whose actor has already ended can be dropped - /// normally rather than being kept alive by this map. Empty until the - /// claiming connection is established, see [`SessionClaim::attach_sender`]. - sender: Weak, + /// Empty while the connection is still running `client_connected` and has + /// no actor yet. Weak so that the map never keeps a sender alive. + sender: Option>, } -/// One session's holder, behind the lock which serializes handovers of it. -/// -/// The lock is held for the whole of a handover: from the moment a connection -/// claims the session until that connection is established, or gives up. -/// A connection claiming a session whose handover is still in flight waits -/// here, so it never finds a holder whose connection does not exist yet and -/// therefore cannot be stopped. -type SessionSlot = AsyncMutex>; - -/// The map of live sessions for one host. -/// -/// Maps each live session to the connection currently serving it. The entry is -/// removed when that connection ends, so a session never outlives its -/// connection. +/// The session is held by another connection. +#[derive(Debug, PartialEq, Eq)] +pub struct SessionBusy; + +/// The live sessions of one host, each mapped to the connection holding it. #[derive(Default)] pub struct ClientSessionIndex { - sessions: Mutex>>, + sessions: Mutex>, } impl ClientSessionIndex { @@ -93,112 +82,55 @@ impl ClientSessionIndex { Self::default() } - /// Claim a session for `client`, tearing down the connection it supersedes. + /// Reserve a session for `client`. /// - /// Returns once that connection is fully closed: its actor has been - /// stopped, and `teardown`, which runs the module-side disconnect, has run - /// to completion. The caller may therefore run `client_connected` for - /// `client` as soon as this returns, and the module observes the two - /// connections' lifecycle events in order. + /// If another connection holds the session, its actor is stopped and + /// [`SessionBusy`] is returned. The session frees up once that connection + /// has been torn down, so the caller should refuse `client` and let it + /// retry. /// - /// The returned [`SessionClaim`] holds the session for the rest of the - /// handover. Another connection claiming the same session waits until the - /// claim is completed with [`SessionClaim::attach_sender`] or dropped, - /// so every claim finds a connection which it can actually stop. - pub async fn claim_session( + /// The returned reservation releases the session when dropped. The caller + /// should complete it with [`SessionReservation::establish`] once the + /// connection has an actor, and drop it only after the connection is fully + /// closed. + pub fn try_reserve( self: &Arc, database_identity: Identity, client: ClientActorId, session_id: SessionId, - teardown: F, - ) -> SessionClaim - where - F: FnOnce(ClientActorId) -> Fut, - Fut: Future, - { + ) -> Result { let key = SessionKey { database_identity, client_identity: client.identity, session_id, }; - let slot = { - let mut sessions = self.sessions.lock().expect("session index poisoned"); - sessions.entry(key).or_default().clone() - }; - - // Wait out any handover of this session which is still in flight. - let mut holder = slot.lock_owned().await; - let superseded = holder.replace(SessionHolder { - client_id: client, - sender: Weak::new(), - }); - - // Connections are told apart by their name, the host's per-connection - // counter, rather than by their connection id, which a client may - // repeat across connections. - if let Some(superseded) = superseded.filter(|superseded| superseded.client_id.name != client.name) { + let mut sessions = self.sessions.lock().expect("session index poisoned"); + if let Some(holder) = sessions.get(&key) { log::debug!( - "websocket: Connection {} supersedes {} for session {session_id}", + "websocket: Connection {} refused, session {session_id} still held by {}", client.connection_id, - superseded.client_id.connection_id, + holder.client.connection_id, ); - if let Some(sender) = superseded.sender.upgrade() { + if let Some(sender) = holder.sender.as_ref().and_then(Weak::upgrade) { sender.kick(ClientDisconnectCause::ConnectionSuperseded); } - teardown(superseded.client_id).await; + return Err(SessionBusy); } - - SessionClaim { + sessions.insert(key, SessionHolder { client, sender: None }); + Ok(SessionReservation { index: self.clone(), key, - holder: Some(holder), - } + client, + }) } - /// Release a session if it is still held by `client`. - /// - /// Called when a connection ends. A connection which has already been - /// superseded no longer holds the session, so it leaves the entry alone: - /// otherwise a slow teardown would evict its own replacement. - pub fn release_session(&self, database_identity: Identity, client: ClientActorId, session_id: SessionId) { - let key = SessionKey { - database_identity, - client_identity: client.identity, - session_id, - }; + /// Remove the session if `client` still holds it. + fn release(&self, key: SessionKey, client: ClientActorId) { let mut sessions = self.sessions.lock().expect("session index poisoned"); - let Entry::Occupied(entry) = sessions.entry(key) else { - return; - }; - // This runs while a connection is being dropped, so it must not block. - // Failing to take the lock means a handover of this session is in - // flight: either another connection's, in which case this connection - // no longer holds the session and there is nothing to release, or this - // connection's own, whose actor ended before it was established. The - // latter leaves the entry behind, holding a sender which can no longer - // be upgraded, until the next claim of the session replaces it. - let Ok(mut holder) = entry.get().clone().try_lock_owned() else { - return; - }; - if holder.as_ref().is_some_and(|held| held.client_id.name == client.name) { - *holder = None; - } - let is_vacant = holder.is_none(); - drop(holder); - if is_vacant { - Self::prune(entry); - } - } - - /// Remove a session which no connection holds and none is claiming. - /// - /// A connection waiting to claim the session holds a reference to the slot, - /// so a strong count of one means the map holds the only reference and the - /// entry can go. Anything else would leave a claimant waiting on a slot no - /// longer reachable from the map, which a later claimant would not find. - fn prune(entry: OccupiedEntry<'_, SessionKey, Arc>) { - if Arc::strong_count(entry.get()) == 1 { - entry.remove(); + // Connections are told apart by their name, the host's per-connection + // counter, rather than by their connection id, which a client may repeat. + if sessions.get(&key).is_some_and(|holder| holder.client.name == client.name) { + sessions.remove(&key); } } @@ -212,47 +144,30 @@ impl ClientSessionIndex { } } -/// A session held for the duration of one handover. +/// A session held by one connection for its whole lifetime. /// -/// Returned by [`ClientSessionIndex::claim_session`]. Completed by -/// [`SessionClaim::attach_sender`] once the claiming connection exists; -/// dropping it without that gives the session up again, for a connection which -/// never came to be. -pub struct SessionClaim { +/// Dropping it releases the session, so it must be dropped only after the +/// connection is fully closed, including its module-side disconnect. +pub struct SessionReservation { index: Arc, key: SessionKey, - /// `Some` until the claim is completed or dropped. - holder: Option>>, + client: ClientActorId, } -impl SessionClaim { - /// Complete the handover, recording the now-established connection's sender - /// so that a later connection can stop it. - /// - /// Holding the claim is what proves the session is still this connection's, - /// so no ownership check is needed here. - pub fn attach_sender(mut self, sender: &Arc) { - let mut holder = self.holder.take().expect("a claim is completed at most once"); - if let Some(held) = holder.as_mut() { - held.sender = Arc::downgrade(sender); +impl SessionReservation { + /// Record the connection's sender, so that a later connection for the + /// same session can stop it. + pub fn establish(&self, sender: &Arc) { + let mut sessions = self.index.sessions.lock().expect("session index poisoned"); + if let Some(holder) = sessions.get_mut(&self.key) { + holder.sender = Some(Arc::downgrade(sender)); } } } -impl Drop for SessionClaim { +impl Drop for SessionReservation { fn drop(&mut self) { - // Completed claims took the guard in `attach_sender`, leaving the - // session held by the connection which is now serving it. - let Some(mut holder) = self.holder.take() else { - return; - }; - *holder = None; - drop(holder); - - let mut sessions = self.index.sessions.lock().expect("session index poisoned"); - if let Entry::Occupied(entry) = sessions.entry(self.key) { - ClientSessionIndex::prune(entry); - } + self.index.release(self.key, self.client); } } @@ -264,7 +179,6 @@ mod tests { use crate::host::module_host::NoSuchModule; use spacetimedb_durability::DurableOffset; use spacetimedb_lib::ConnectionId; - use std::sync::atomic::{AtomicBool, Ordering}; /// The dummy senders below never wait on durability. struct NoDurability; @@ -279,9 +193,8 @@ mod tests { Arc::new(ClientSessionIndex::new()) } - /// A client id whose `name`, the host's per-connection counter, matches its - /// connection id, so that tests naming distinct connections get distinct - /// names as the websocket handler would assign them. + /// A client id whose `name` matches its connection id, as the websocket + /// handler would assign distinct names to distinct connections. fn client(identity: Identity, connection_id: u128) -> ClientActorId { ClientActorId { identity, @@ -314,69 +227,49 @@ mod tests { Identity::from_byte_array([2; 32]) } - /// Claim a session, recording which connection was torn down, if any. - async fn claim( - index: &Arc, - database: Identity, - client: ClientActorId, - session: SessionId, - ) -> (SessionClaim, Option) { - let mut superseded = None; - let claim = index - .claim_session(database, client, session, async |old| superseded = Some(old)) - .await; - (claim, superseded) + fn session() -> SessionId { + SessionId::from_u128(7) } - /// Claim a session and complete the handover, as an established connection - /// does. - async fn connect( + /// Reserve a session and establish it, as a connected client does. + fn connect( index: &Arc, database: Identity, client: ClientActorId, session: SessionId, - ) -> (Arc, Option) { - let (claim, superseded) = claim(index, database, client, session).await; + ) -> (SessionReservation, Arc) { + let reservation = index.try_reserve(database, client, session).expect("session should be free"); let sender = sender(client); - claim.attach_sender(&sender); - (sender, superseded) + reservation.establish(&sender); + (reservation, sender) } #[tokio::test] - async fn first_connection_supersedes_nothing() { + async fn first_connection_reserves_the_session() { let index = index(); - let session = SessionId::from_u128(7); - let (_sender, superseded) = connect(&index, a_database(), client(an_identity(), 1), session).await; + let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session()); - assert!(superseded.is_none()); assert_eq!(index.len(), 1); } #[tokio::test] - async fn reconnect_supersedes_previous_connection() { + async fn reconnect_is_refused_and_stops_the_holder() { let index = index(); - let session = SessionId::from_u128(7); - let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; - - let (_second, superseded) = connect(&index, a_database(), client(an_identity(), 2), session).await; - - assert_eq!( - superseded.map(|old| old.connection_id), - Some(ConnectionId::from_u128(1)) - ); - // The superseded connection's actor is stopped. - assert!(first.is_cancelled()); - // The session is now held by the new connection, not the old one. + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); + + let refused = index.try_reserve(a_database(), client(an_identity(), 2), session()); + + assert!(refused.is_err()); + assert!(first_sender.is_cancelled()); assert_eq!(index.len(), 1); } - /// A client may repeat a connection id across connections, so the session - /// is handed over on the connection's name rather than on that id. + /// A client may repeat a connection id across connections, so the holder + /// is told apart by its name rather than by that id. #[tokio::test] - async fn reconnect_reusing_connection_id_supersedes() { + async fn reconnect_reusing_connection_id_is_refused() { let index = index(); - let session = SessionId::from_u128(7); let first = ClientActorId { name: ClientName(1), ..client(an_identity(), 1) @@ -385,162 +278,81 @@ mod tests { name: ClientName(2), ..client(an_identity(), 1) }; - let (first_sender, _) = connect(&index, a_database(), first, session).await; + let (_first, first_sender) = connect(&index, a_database(), first, session()); - let (_second_sender, superseded) = connect(&index, a_database(), second, session).await; - - assert_eq!(superseded.map(|old| old.name), Some(ClientName(1))); + assert!(index.try_reserve(a_database(), second, session()).is_err()); assert!(first_sender.is_cancelled()); - assert_eq!(index.len(), 1); } #[tokio::test] - async fn different_identity_does_not_supersede() { + async fn holder_without_a_sender_still_refuses() { let index = index(); - let session = SessionId::from_u128(7); - let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; + let _first = index + .try_reserve(a_database(), client(an_identity(), 1), session()) + .unwrap(); - let (_second, superseded) = connect(&index, a_database(), client(another_identity(), 2), session).await; - - assert!(superseded.is_none()); - assert_eq!(index.len(), 2); + assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); + assert_eq!(index.len(), 1); } #[tokio::test] - async fn different_session_does_not_supersede() { + async fn different_identity_does_not_conflict() { let index = index(); - let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), SessionId::from_u128(7)).await; + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); - let (_second, superseded) = - connect(&index, a_database(), client(an_identity(), 2), SessionId::from_u128(8)).await; + let (_second, _) = connect(&index, a_database(), client(another_identity(), 2), session()); - assert!(superseded.is_none()); + assert!(!first_sender.is_cancelled()); assert_eq!(index.len(), 2); } - /// A session belongs to one database, so a connection to another database - /// never tears down this one, whose module knows nothing about it. #[tokio::test] - async fn different_database_does_not_supersede() { + async fn different_session_does_not_conflict() { let index = index(); - let session = SessionId::from_u128(7); - let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); - let (_second, superseded) = connect(&index, another_database(), client(an_identity(), 2), session).await; + let (_second, _) = connect(&index, a_database(), client(an_identity(), 2), SessionId::from_u128(8)); - assert!(superseded.is_none()); - assert!(!first.is_cancelled()); + assert!(!first_sender.is_cancelled()); assert_eq!(index.len(), 2); } #[tokio::test] - async fn release_removes_the_session() { + async fn different_database_does_not_conflict() { let index = index(); - let session = SessionId::from_u128(7); - let connection = client(an_identity(), 1); - let (_sender, _) = connect(&index, a_database(), connection, session).await; + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); - index.release_session(a_database(), connection, session); + let (_second, _) = connect(&index, another_database(), client(an_identity(), 2), session()); - assert!(index.is_empty()); + assert!(!first_sender.is_cancelled()); + assert_eq!(index.len(), 2); } #[tokio::test] - async fn superseded_connection_release_does_not_evict_its_replacement() { + async fn dropping_the_reservation_frees_the_session() { let index = index(); - let session = SessionId::from_u128(7); - let old = client(an_identity(), 1); - let new = client(an_identity(), 2); - let (_old_sender, _) = connect(&index, a_database(), old, session).await; - let (_new_sender, _) = connect(&index, a_database(), new, session).await; + let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session()); + assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); - // The old connection tears down after being superseded. - index.release_session(a_database(), old, session); + drop(first); - // The replacement still holds the session. + assert!(index.is_empty()); + let retry = index.try_reserve(a_database(), client(an_identity(), 2), session()); + assert!(retry.is_ok()); assert_eq!(index.len(), 1); - let (_third, superseded) = connect(&index, a_database(), client(an_identity(), 3), session).await; - assert_eq!(superseded.map(|old| old.connection_id), Some(new.connection_id)); } /// A connection which never came to be, because `client_connected` - /// rejected it, gives the session up again. + /// rejected it, frees the session without ever establishing it. #[tokio::test] - async fn dropped_claim_releases_the_session() { + async fn dropping_an_unestablished_reservation_frees_the_session() { let index = index(); - let session = SessionId::from_u128(7); + let reservation = index + .try_reserve(a_database(), client(an_identity(), 1), session()) + .unwrap(); - let (claim, _) = claim(&index, a_database(), client(an_identity(), 1), session).await; - drop(claim); + drop(reservation); assert!(index.is_empty()); } - - #[tokio::test] - async fn three_way_race_supersedes_the_most_recent_connection() { - let index = index(); - let session = SessionId::from_u128(7); - let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session).await; - - let (_second, superseded_by_second) = connect(&index, a_database(), client(an_identity(), 2), session).await; - let (_third, superseded_by_third) = connect(&index, a_database(), client(an_identity(), 3), session).await; - - assert_eq!( - superseded_by_second.map(|old| old.connection_id), - Some(ConnectionId::from_u128(1)) - ); - assert_eq!( - superseded_by_third.map(|old| old.connection_id), - Some(ConnectionId::from_u128(2)) - ); - } - - /// A claim waits for the handover in flight, so it never supersedes a - /// connection which does not exist yet and so cannot be stopped. - #[tokio::test] - async fn handover_serializes_concurrent_claims() { - let index = index(); - let session = SessionId::from_u128(7); - let first = client(an_identity(), 1); - let second = client(an_identity(), 2); - - // The first connection claims the session but is not established yet, - // as it would not be while its `client_connected` runs. - let (claim, _) = claim(&index, a_database(), first, session).await; - - let torn_down = Arc::new(Mutex::new(None)); - let started = Arc::new(AtomicBool::new(false)); - let claimed = Arc::new(AtomicBool::new(false)); - let racing = tokio::spawn({ - let (index, torn_down) = (index.clone(), torn_down.clone()); - let (started, claimed) = (started.clone(), claimed.clone()); - async move { - started.store(true, Ordering::Release); - let claim = index - .claim_session(a_database(), second, session, async |old| { - *torn_down.lock().unwrap() = Some(old); - }) - .await; - claimed.store(true, Ordering::Release); - claim.attach_sender(&sender(second)); - } - }); - - // The second connection cannot claim the session while the first - // connection's handover is still in flight. - tokio::task::yield_now().await; - assert!(started.load(Ordering::Acquire), "the racing claim never ran"); - assert!(!claimed.load(Ordering::Acquire), "the racing claim should be waiting"); - assert!(torn_down.lock().unwrap().is_none()); - - // Once the first connection is established, the second supersedes it, - // and finds a connection which it can stop. - let first_sender = sender(first); - claim.attach_sender(&first_sender); - racing.await.unwrap(); - - assert_eq!(torn_down.lock().unwrap().map(|old| old.name), Some(first.name)); - assert!(first_sender.is_cancelled()); - assert_eq!(index.len(), 1); - } } diff --git a/crates/core/src/worker_metrics/mod.rs b/crates/core/src/worker_metrics/mod.rs index 95a074afcd8..a0d5f45f717 100644 --- a/crates/core/src/worker_metrics/mod.rs +++ b/crates/core/src/worker_metrics/mod.rs @@ -280,6 +280,11 @@ metrics_group!( #[labels(database_identity: Identity)] pub ws_clients_idle_timed_out: IntCounterVec, + #[name = spacetime_worker_ws_clients_session_busy_total] + #[help = "The cumulative number of ws connections refused because their session was still held by a connection being torn down"] + #[labels(database_identity: Identity)] + pub ws_clients_session_busy: IntCounterVec, + // Compatibility counters above continue to be emitted for existing dashboards. // Accepted-client disconnection `cause` label values are: // client_close, idle_timeout, incoming_queue_full, outgoing_queue_full, diff --git a/crates/smoketests/tests/cluster/connection_session.rs b/crates/smoketests/tests/cluster/connection_session.rs index 1f87311f7de..be6ca0f48fe 100644 --- a/crates/smoketests/tests/cluster/connection_session.rs +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -1,9 +1,12 @@ //! Tests for connection replacement, the server side of SDK auto-reconnect. //! //! A reconnecting client supplies a stable `session_id`. When it reconnects -//! before the server has noticed the old socket died, the new connection -//! supersedes the old one. The old connection is torn down through the normal -//! disconnect sequence before the new connection's `client_connected` runs. +//! before the server has noticed the old socket died, the new connection is +//! refused with `SESSION_BUSY_CLOSE_CODE` and the old one is torn down through +//! the normal disconnect sequence. A retry succeeds once the old connection's +//! `client_disconnected` has run. + +use std::time::{Duration, Instant}; use anyhow::{bail, Context, Result}; use futures::{SinkExt, StreamExt}; @@ -27,7 +30,25 @@ struct TestConnection { impl TestConnection { /// Open a connection, optionally supplying a `session_id`, and wait for the /// server's `InitialConnection` message. + /// + /// Retries while the server refuses the connection because its session + /// is still held by a connection being torn down. async fn open(test: &Smoketest, connection_id: &str, session_id: Option<&str>) -> Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if let Some(connection) = Self::open_once(test, connection_id, session_id).await? { + return Ok(connection); + } + if Instant::now() > deadline { + bail!("timed out retrying a connection refused as session busy"); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// Open a connection once. Returns `None` if the server refused it with + /// `SESSION_BUSY_CLOSE_CODE`. + async fn open_once(test: &Smoketest, connection_id: &str, session_id: Option<&str>) -> Result> { let token = test.read_token()?; let host = test.server_host(); let database = test @@ -64,8 +85,8 @@ impl TestConnection { socket, connection_id: connection_id.to_string(), }; - match connection.next_message().await? { - ws_v2::ServerMessage::InitialConnection(initial) => { + match connection.next_message().await { + Ok(ws_v2::ServerMessage::InitialConnection(initial)) => { let established = initial.connection_id.to_hex().to_string(); if established != connection.connection_id { bail!( @@ -73,10 +94,12 @@ impl TestConnection { connection.connection_id ); } + Ok(Some(connection)) } - other => bail!("expected InitialConnection, got {other:?}"), + Ok(other) => bail!("expected InitialConnection, got {other:?}"), + Err(err) if err.downcast_ref::().is_some_and(|closed| closed.is_session_busy()) => Ok(None), + Err(err) => Err(err), } - Ok(connection) } /// Read the next server message, decoding the v3 framing, which packs one @@ -99,7 +122,7 @@ impl TestConnection { return Ok(bsatn::from_reader(&mut body)?); } Message::Ping(_) | Message::Pong(_) => continue, - Message::Close(frame) => bail!("websocket closed: {frame:?}"), + Message::Close(frame) => return Err(Closed(frame.map(|frame| frame.code.into())).into()), other => bail!("unexpected websocket message: {other:?}"), } } @@ -113,9 +136,9 @@ impl TestConnection { /// Whether the server still serves this connection. /// - /// A superseded connection's actor is stopped, so a request on it is never - /// answered. Note the server does not send a close frame. The peer's - /// socket stays half-open until it writes, which is what this does. + /// A stopped connection's actor never answers a request. The server does + /// not send a close frame, so the peer's socket stays half-open until it + /// writes, which is what this does. async fn is_still_served(&mut self) -> bool { if self .send(ws_v2::ClientMessage::Subscribe(ws_v2::Subscribe { @@ -135,6 +158,24 @@ impl TestConnection { } } +/// The server closed the websocket, with the close code it sent, if any. +#[derive(Debug)] +struct Closed(Option); + +impl Closed { + fn is_session_busy(&self) -> bool { + self.0 == Some(ws_common::SESSION_BUSY_CLOSE_CODE) + } +} + +impl std::fmt::Display for Closed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "websocket closed with code {:?}", self.0) + } +} + +impl std::error::Error for Closed {} + fn runtime() -> tokio::runtime::Runtime { tokio::runtime::Builder::new_current_thread() .enable_all() @@ -179,9 +220,9 @@ const CONNECTION_C: &str = "00000000000000000000000000000c33"; const SESSION: &str = "0000000000000000000000000000dead"; const OTHER_SESSION: &str = "0000000000000000000000000000beef"; -/// A second connection with the same session id supersedes the first: the old -/// connection is disconnected, and its `client_disconnected` runs strictly -/// before the new connection's `client_connected`. +/// A second connection with the same session id is refused while the first is +/// still live, and the first is stopped. Its `client_disconnected` runs +/// strictly before the retried connection's `client_connected`. #[test] fn test_reconnect_with_same_session_replaces_connection() { let test = Smoketest::builder().precompiled_module("connection-session").build(); @@ -193,15 +234,23 @@ fn test_reconnect_with_same_session_replaces_connection() { wait_for_log(&test, "connected", CONNECTION_A); // Reconnect with the same session before the server notices the drop. - let _second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + let refused = TestConnection::open_once(&test, CONNECTION_B, Some(SESSION)) .await .expect("second connection failed"); + assert!( + refused.is_none(), + "the reconnect should be refused while the first connection is live" + ); assert!( !first.is_still_served().await, - "the superseded connection should no longer be served" + "the first connection should no longer be served" ); + let _second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("retried connection failed"); + let lines = wait_for_log(&test, "connected", CONNECTION_B); let connected_a = position_of(&lines, "connected", CONNECTION_A).expect("A never connected"); let disconnected_a = position_of(&lines, "disconnected", CONNECTION_A).expect("A never disconnected"); @@ -218,7 +267,7 @@ fn test_reconnect_with_same_session_replaces_connection() { }); } -/// A connection with a different session id does not supersede: both stay live. +/// A connection with a different session id is not refused: both stay live. #[test] fn test_different_session_does_not_replace_connection() { let test = Smoketest::builder().precompiled_module("connection-session").build(); @@ -242,7 +291,7 @@ fn test_different_session_does_not_replace_connection() { } /// A connection which supplies no session id behaves exactly as before: it -/// neither supersedes nor is superseded. +/// neither refuses nor is stopped by a connection with a session id. #[test] fn test_connection_without_session_is_not_replaced() { let test = Smoketest::builder().precompiled_module("connection-session").build(); @@ -260,12 +309,12 @@ fn test_connection_without_session_is_not_replaced() { assert!( position_of(&lines, "disconnected", CONNECTION_A).is_none(), - "a connection without a session id should not be superseded: {lines:?}" + "a connection without a session id should not be stopped: {lines:?}" ); }); } -/// Repeated reconnects each supersede only the connection immediately before +/// Repeated reconnects each replace only the connection immediately before /// them, leaving exactly one live connection for the session. #[test] fn test_repeated_reconnects_leave_one_live_connection() { @@ -281,18 +330,18 @@ fn test_repeated_reconnects_leave_one_live_connection() { .await .expect("second connection failed"); wait_for_log(&test, "connected", CONNECTION_B); - assert!(!first.is_still_served().await, "A should have been superseded"); + assert!(!first.is_still_served().await, "A should have been stopped"); let mut third = TestConnection::open(&test, CONNECTION_C, Some(SESSION)) .await .expect("third connection failed"); wait_for_log(&test, "connected", CONNECTION_C); - assert!(!second.is_still_served().await, "B should have been superseded"); + assert!(!second.is_still_served().await, "B should have been stopped"); let lines = lifecycle_log(&test); assert!( position_of(&lines, "disconnected", CONNECTION_B).is_some(), - "B should have been superseded by C: {lines:?}" + "B should have been replaced by C: {lines:?}" ); assert!( position_of(&lines, "disconnected", CONNECTION_C).is_none(), @@ -315,7 +364,7 @@ fn test_repeated_reconnects_leave_one_live_connection() { }); } -/// A reconnect which repeats its predecessor's connection id still supersedes +/// A reconnect which repeats its predecessor's connection id still replaces /// it. Connections are told apart by the server, not by the id a client sends, /// which a client is free to repeat. #[test] @@ -335,7 +384,7 @@ fn test_reconnect_reusing_connection_id_replaces_connection() { assert!( !first.is_still_served().await, - "the superseded connection should no longer be served" + "the replaced connection should no longer be served" ); assert!( second.is_still_served().await, From 7ebdc8cf005a9514ee3c879b0f1a9ccdc667d517 Mon Sep 17 00:00:00 2001 From: Jeffrey Dallatezza Date: Thu, 10 Sep 2026 04:09:16 -0700 Subject: [PATCH 07/14] Add a drop safe async version of scope guard (#5910) --- crates/client-api/src/routes/subscribe.rs | 18 +- crates/client-api/src/util.rs | 185 ++++++++++++++++++ .../core/src/client/client_session_index.rs | 17 +- .../tests/cluster/connection_session.rs | 8 +- 4 files changed, 216 insertions(+), 12 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index e1293a3997a..462bace2d2f 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -20,7 +20,7 @@ use derive_more::From; use futures::{pin_mut, Sink, SinkExt, Stream, StreamExt}; use http::{HeaderValue, StatusCode}; use prometheus::{Histogram, IntGauge}; -use scopeguard::{defer, ScopeGuard}; +use scopeguard::defer; use serde::Deserialize; use spacetimedb::client::messages::{ serialize, serialize_v3, IdentityTokenMessage, InUseSerializeBuffer, SerializeBuffer, SwitchedServerMessage, @@ -58,7 +58,7 @@ use crate::util::serde::humantime_duration; use crate::util::websocket::{ CloseCode, CloseFrame, Message as WsMessage, WebSocketConfig, WebSocketStream, WebSocketUpgrade, WsError, }; -use crate::util::{NameOrIdentity, XForwardedFor}; +use crate::util::{async_cleanup_guard, NameOrIdentity, XForwardedFor}; use crate::{log_and_500, Authorization, ControlStateDelegate, NodeDelegate}; #[allow(clippy::declare_interior_mutable_const)] @@ -298,7 +298,10 @@ where Some(session_id) => match sessions.try_reserve(db_identity, client_id, session_id) { Ok(reservation) => Some(reservation), Err(SessionBusy) => { - WORKER_METRICS.ws_clients_session_busy.with_label_values(&db_identity).inc(); + WORKER_METRICS + .ws_clients_session_busy + .with_label_values(&db_identity) + .inc(); log::debug!("websocket: Refusing connection for {client_log_string}: session {session_id} is busy"); let close = CloseFrame { code: CloseCode::from(SESSION_BUSY_CLOSE_CODE), @@ -575,14 +578,15 @@ async fn ws_client_actor( session: Option, ) { // Runs the module-side disconnect even if this task gets cancelled. - let mut client = scopeguard::guard((client, session), |(client, session)| { - tokio::spawn(ws_client_teardown(client, session)); + let mut client = async_cleanup_guard((client, session), |(client, session)| { + ws_client_teardown(client, session) }); ws_client_actor_inner(&mut client.0, options, ws, sendrx).await; - let (client, session) = ScopeGuard::into_inner(client); - ws_client_teardown(client, session).await; + if let Err(e) = client.cleanup().await { + log::error!("websocket client teardown task failed: {e}"); + } } /// Run the module-side disconnect, then free the connection's session. diff --git a/crates/client-api/src/util.rs b/crates/client-api/src/util.rs index 0cc87a82bfd..66e5fc91158 100644 --- a/crates/client-api/src/util.rs +++ b/crates/client-api/src/util.rs @@ -3,7 +3,10 @@ pub(crate) mod serde; pub mod websocket; use core::fmt; +use std::future::Future; +use std::marker::PhantomData; use std::net::IpAddr; +use std::ops::{Deref, DerefMut}; use axum::body::Bytes; use axum::extract::{FromRequest, Request}; @@ -15,10 +18,117 @@ use http::{HeaderName, HeaderValue, StatusCode}; use hyper::body::Body; use spacetimedb::Identity; use spacetimedb_client_api_messages::name::DatabaseName; +use tokio::task::{JoinError, JoinHandle}; use crate::routes::identity::IdentityForUrl; use crate::{log_and_500, ControlStateReadAccess}; +/// Returns a guard that runs async cleanup for `value` when dropped. +/// +/// This is cancel-safe with respect to cancellation of the task holding the +/// guard: dropping the guard spawns the cleanup future in its own task instead +/// of trying to run async cleanup from `Drop`. +/// +/// This does not guarantee that cleanup survives shutdown of the Tokio runtime, +/// process exit, or explicit abortion of the spawned cleanup task. +/// +/// Dropping this guard calls [`tokio::spawn`], so it must be dropped from +/// within a Tokio runtime. +pub(crate) fn async_cleanup_guard(value: T, cleanup: F) -> AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + AsyncCleanupGuard { + value: Some(value), + cleanup: Some(cleanup), + future: PhantomData, + } +} + +/// Scope guard for values that require async cleanup. +/// +/// Dropping the guard is cancel-safe for the guarded task: it moves the guarded +/// value into a newly spawned cleanup task. Drop does not wait for cleanup to +/// complete. +/// +/// Call [`Self::cleanup`] on the normal path when the current task should wait +/// for cleanup. That method starts cleanup in a spawned task before awaiting it, +/// so cancelling the waiter does not cancel the cleanup task. +pub(crate) struct AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + value: Option, + cleanup: Option, + future: PhantomData Fut>, +} + +impl AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + fn spawn_cleanup(&mut self) -> JoinHandle<()> { + let value = self.value.take().expect("cleanup value already taken"); + let cleanup = self.cleanup.take().expect("cleanup function already taken"); + tokio::spawn(cleanup(value)) + } + + /// Starts cleanup and waits for the cleanup task to finish. + /// + /// This is cancel-safe with respect to cancellation of the caller: cleanup + /// is spawned before this method awaits, so dropping this future after its + /// first poll drops only the wait for completion, not the cleanup itself. + /// + /// This is not cancel-safe against explicit abortion of the returned + /// cleanup task by the runtime or against runtime shutdown. + pub(crate) async fn cleanup(mut self) -> Result<(), JoinError> { + self.spawn_cleanup().await + } +} + +impl Deref for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + type Target = T; + + fn deref(&self) -> &Self::Target { + self.value.as_ref().expect("cleanup value already taken") + } +} + +impl DerefMut for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + fn deref_mut(&mut self) -> &mut Self::Target { + self.value.as_mut().expect("cleanup value already taken") + } +} + +impl Drop for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + fn drop(&mut self) { + if let (Some(value), Some(cleanup)) = (self.value.take(), self.cleanup.take()) { + tokio::spawn(cleanup(value)); + } + } +} + pub struct ByteStringBody(pub ByteString); #[async_trait::async_trait] @@ -192,6 +302,9 @@ impl FromRequest for EmptyBody { mod tests { use super::*; use headers::Header; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + use tokio::sync::{oneshot, Notify}; fn decode_one(raw: &str) -> Result { let val = HeaderValue::from_str(raw).unwrap(); @@ -224,4 +337,76 @@ mod tests { assert!(decode_one("not-an-ip").is_err()); assert!(decode_one("not-an-ip, 10.0.0.1").is_err()); } + + #[tokio::test] + async fn async_cleanup_guard_runs_cleanup_when_dropped() { + let (tx, rx) = oneshot::channel(); + + drop(async_cleanup_guard((), move |()| async move { + tx.send(()).unwrap(); + })); + + rx.await.expect("cleanup should run"); + } + + #[tokio::test] + async fn async_cleanup_guard_cleanup_waits_for_cleanup() { + let cleaned_up = Arc::new(AtomicBool::new(false)); + let cleanup_started = Arc::new(Notify::new()); + let finish_cleanup = Arc::new(Notify::new()); + let cleaned_up_for_guard = Arc::clone(&cleaned_up); + let cleanup_started_for_guard = Arc::clone(&cleanup_started); + let finish_cleanup_for_guard = Arc::clone(&finish_cleanup); + + let cleanup = tokio::spawn( + async_cleanup_guard((), move |()| async move { + cleanup_started_for_guard.notify_one(); + finish_cleanup_for_guard.notified().await; + cleaned_up_for_guard.store(true, Ordering::Release); + }) + .cleanup(), + ); + + cleanup_started.notified().await; + assert!(!cleaned_up.load(Ordering::Acquire)); + + finish_cleanup.notify_one(); + cleanup + .await + .expect("cleanup join task should not panic") + .expect("cleanup task should not panic"); + assert!(cleaned_up.load(Ordering::Acquire)); + } + + #[tokio::test] + async fn async_cleanup_guard_cleanup_continues_if_waiter_is_aborted() { + let cleaned_up = Arc::new(AtomicBool::new(false)); + let cleanup_started = Arc::new(Notify::new()); + let finish_cleanup = Arc::new(Notify::new()); + let cleaned_up_for_guard = Arc::clone(&cleaned_up); + let cleanup_started_for_guard = Arc::clone(&cleanup_started); + let finish_cleanup_for_guard = Arc::clone(&finish_cleanup); + + let cleanup = tokio::spawn( + async_cleanup_guard((), move |()| async move { + cleanup_started_for_guard.notify_one(); + finish_cleanup_for_guard.notified().await; + cleaned_up_for_guard.store(true, Ordering::Release); + }) + .cleanup(), + ); + + cleanup_started.notified().await; + cleanup.abort(); + assert!(cleanup.await.unwrap_err().is_cancelled()); + + finish_cleanup.notify_one(); + tokio::time::timeout(std::time::Duration::from_secs(1), async { + while !cleaned_up.load(Ordering::Acquire) { + tokio::task::yield_now().await; + } + }) + .await + .expect("cleanup should continue after waiter abort"); + } } diff --git a/crates/core/src/client/client_session_index.rs b/crates/core/src/client/client_session_index.rs index 6d968d774f2..f958ebb9a34 100644 --- a/crates/core/src/client/client_session_index.rs +++ b/crates/core/src/client/client_session_index.rs @@ -129,7 +129,10 @@ impl ClientSessionIndex { let mut sessions = self.sessions.lock().expect("session index poisoned"); // Connections are told apart by their name, the host's per-connection // counter, rather than by their connection id, which a client may repeat. - if sessions.get(&key).is_some_and(|holder| holder.client.name == client.name) { + if sessions + .get(&key) + .is_some_and(|holder| holder.client.name == client.name) + { sessions.remove(&key); } } @@ -238,7 +241,9 @@ mod tests { client: ClientActorId, session: SessionId, ) -> (SessionReservation, Arc) { - let reservation = index.try_reserve(database, client, session).expect("session should be free"); + let reservation = index + .try_reserve(database, client, session) + .expect("session should be free"); let sender = sender(client); reservation.establish(&sender); (reservation, sender) @@ -291,7 +296,9 @@ mod tests { .try_reserve(a_database(), client(an_identity(), 1), session()) .unwrap(); - assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); + assert!(index + .try_reserve(a_database(), client(an_identity(), 2), session()) + .is_err()); assert_eq!(index.len(), 1); } @@ -332,7 +339,9 @@ mod tests { async fn dropping_the_reservation_frees_the_session() { let index = index(); let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session()); - assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); + assert!(index + .try_reserve(a_database(), client(an_identity(), 2), session()) + .is_err()); drop(first); diff --git a/crates/smoketests/tests/cluster/connection_session.rs b/crates/smoketests/tests/cluster/connection_session.rs index be6ca0f48fe..b7ea3758a1d 100644 --- a/crates/smoketests/tests/cluster/connection_session.rs +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -97,7 +97,13 @@ impl TestConnection { Ok(Some(connection)) } Ok(other) => bail!("expected InitialConnection, got {other:?}"), - Err(err) if err.downcast_ref::().is_some_and(|closed| closed.is_session_busy()) => Ok(None), + Err(err) + if err + .downcast_ref::() + .is_some_and(|closed| closed.is_session_busy()) => + { + Ok(None) + } Err(err) => Err(err), } } From 0ad26528714094ad750e13a24df15a5e2a547784 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Fri, 11 Sep 2026 15:39:56 +0200 Subject: [PATCH 08/14] Typescript SDK: auto-reconnect --- crates/bindings-typescript/README.md | 30 +- crates/bindings-typescript/src/lib/errors.ts | 34 + .../src/sdk/client_api/types.ts | 48 + .../src/sdk/connection_manager.ts | 180 +-- .../src/sdk/db_connection_builder.ts | 51 +- .../src/sdk/db_connection_impl.ts | 872 +++++++++++- .../src/sdk/subscription_builder_impl.ts | 24 +- .../src/sdk/table_cache.ts | 34 +- .../src/sdk/websocket_test_adapter.ts | 84 +- crates/bindings-typescript/src/sdk/ws.ts | 26 +- .../tests/connection_manager_liveness.test.ts | 319 ----- .../connection_manager_reconnect.test.ts | 261 +++- .../tests/db_connection_liveness.test.ts | 267 ++++ .../tests/db_connection_reconnect.test.ts | 1201 +++++++++++++++++ .../tests/table_cache_reconnect.test.ts | 62 + 15 files changed, 2905 insertions(+), 588 deletions(-) delete mode 100644 crates/bindings-typescript/tests/connection_manager_liveness.test.ts create mode 100644 crates/bindings-typescript/tests/db_connection_liveness.test.ts create mode 100644 crates/bindings-typescript/tests/db_connection_reconnect.test.ts create mode 100644 crates/bindings-typescript/tests/table_cache_reconnect.test.ts diff --git a/crates/bindings-typescript/README.md b/crates/bindings-typescript/README.md index 48e7cfd1535..6ebfd1d0e36 100644 --- a/crates/bindings-typescript/README.md +++ b/crates/bindings-typescript/README.md @@ -26,22 +26,28 @@ import { DbConnection, tables } from './module_bindings'; const connection = DbConnection.builder() .withUri('ws://localhost:3000') .withDatabaseName('MODULE_NAME') - .onDisconnect(() => { - console.log('disconnected'); + .withAutomaticReconnect() + .onConnect((_connection, identity) => { + console.log('Connected:', identity.toHexString()); }) - .onConnectError(() => { - console.log('client_error'); - }) - .onConnect((connection, identity, _token) => { + .onDisconnect((_ctx, error, attempt, delayMs) => { console.log( - 'Connected to SpacetimeDB with identity:', - identity.toHexString() + attempt === undefined + ? 'Disconnected' + : `Retry ${attempt} in ${delayMs} ms`, + error + ); + }) + .onConnectError((_ctx, error, attempt) => { + console.error( + attempt === undefined ? 'Connection failed' : 'Retry failed', + error ); - - connection.subscriptionBuilder().subscribe(tables.player); }) .withToken('TOKEN') .build(); + +connection.subscriptionBuilder().subscribe(tables.player); ``` If you need to disconnect the client: @@ -50,6 +56,10 @@ If you need to disconnect the client: connection.disconnect(); ``` +Automatic reconnection preserves the connection, cache, handles, and callbacks. Register subscriptions and row callbacks once, outside `onConnect`, which runs again after every reconnect. Cache reads remain available during outages. Initial connection failures are not retried by the core SDK. + +For expiring credentials, pass the initial token with `withToken` and add `withTokenProvider(() => auth.getAccessToken())`. The SDK asks for a fresh token before reconnecting when the retained token is near expiry. The provider must return a token for the same identity. + Typically, you will use the SDK with types generated from SpacetimeDB module. For example, given a table named `Player` you can subscribe to player updates like this: ```ts diff --git a/crates/bindings-typescript/src/lib/errors.ts b/crates/bindings-typescript/src/lib/errors.ts index c8ec99133c8..93890421b4d 100644 --- a/crates/bindings-typescript/src/lib/errors.ts +++ b/crates/bindings-typescript/src/lib/errors.ts @@ -24,3 +24,37 @@ export class InternalError extends Error { return 'InternalError'; } } + +/** The call was not sent because the connection was not established. */ +export class DisconnectedError extends Error { + constructor(message: string = 'Not connected to SpacetimeDB') { + super(message); + } + get name(): string { + return 'DisconnectedError'; + } +} + +/** The connection dropped before acknowledgement; the call may have run. */ +export class UnknownCallResultError extends Error { + constructor( + message: string = 'Connection lost before the call was acknowledged; it may or may not have run' + ) { + super(message); + } + get name(): string { + return 'UnknownCallResultError'; + } +} + +/** The reconnect returned a different identity, ending automatic reconnection. */ +export class IdentityChangedError extends Error { + constructor( + message: string = 'Reconnected with a different identity; the token was revoked or replaced' + ) { + super(message); + } + get name(): string { + return 'IdentityChangedError'; + } +} diff --git a/crates/bindings-typescript/src/sdk/client_api/types.ts b/crates/bindings-typescript/src/sdk/client_api/types.ts index 709a114da28..ce4f4ed0680 100644 --- a/crates/bindings-typescript/src/sdk/client_api/types.ts +++ b/crates/bindings-typescript/src/sdk/client_api/types.ts @@ -51,6 +51,9 @@ export const ClientMessage = __t.enum('ClientMessage', { get CallProcedure() { return CallProcedure; }, + get SubscribeBatch() { + return SubscribeBatch; + }, }); export type ClientMessage = __Infer; @@ -192,6 +195,9 @@ export const ServerMessage = __t.enum('ServerMessage', { get ProcedureResult() { return ProcedureResult; }, + get SubscribeBatchApplied() { + return SubscribeBatchApplied; + }, }); export type ServerMessage = __Infer; @@ -223,6 +229,48 @@ export const SubscribeApplied = __t.object('SubscribeApplied', { }); export type SubscribeApplied = __Infer; +export const SubscribeBatch = __t.object('SubscribeBatch', { + requestId: __t.u32(), + get sets() { + return __t.array(SubscribeSet); + }, +}); +export type SubscribeBatch = __Infer; + +export const SubscribeBatchApplied = __t.object('SubscribeBatchApplied', { + requestId: __t.u32(), + get results() { + return __t.array(SubscribeSetResult); + }, +}); +export type SubscribeBatchApplied = __Infer; + +export const SubscribeSet = __t.object('SubscribeSet', { + get querySetId() { + return QuerySetId; + }, + queryStrings: __t.array(__t.string()), +}); +export type SubscribeSet = __Infer; + +export const SubscribeSetOutcome = __t.enum('SubscribeSetOutcome', { + get Applied() { + return QueryRows; + }, + Error: __t.string(), +}); +export type SubscribeSetOutcome = __Infer; + +export const SubscribeSetResult = __t.object('SubscribeSetResult', { + get querySetId() { + return QuerySetId; + }, + get outcome() { + return SubscribeSetOutcome; + }, +}); +export type SubscribeSetResult = __Infer; + export const SubscriptionError = __t.object('SubscriptionError', { requestId: __t.option(__t.u32()), get querySetId() { diff --git a/crates/bindings-typescript/src/sdk/connection_manager.ts b/crates/bindings-typescript/src/sdk/connection_manager.ts index 2cf34ea2c36..052fa73e392 100644 --- a/crates/bindings-typescript/src/sdk/connection_manager.ts +++ b/crates/bindings-typescript/src/sdk/connection_manager.ts @@ -25,6 +25,16 @@ * Result: Single WebSocket survives ✓ * ``` * + * ## Reconnection: + * + * The manager forces {@link DbConnectionBuilder.withAutomaticReconnect} on + * every builder it builds from, so a connection lost mid-session reconnects + * *inside* the `DbConnection`: the object, its table handles and its callbacks + * all survive, and the manager merely mirrors the lifecycle events into its + * state snapshots. The manager itself rebuilds a connection only when the SDK + * reports it will not retry (a failed initial connection, or another terminal + * failure), preserving the frameworks' historical keep-trying behavior. + * * @module connection_manager */ import type { @@ -50,9 +60,13 @@ export const CONNECTION_MANAGER_RECONNECT_BASE_DELAY_MS = 1000; export const CONNECTION_MANAGER_RECONNECT_MAX_DELAY_MS = 30_000; /** - * Computes the reconnect delay for the given attempt (0-based) using + * Computes the rebuild delay for the given attempt (0-based) using * exponential backoff: the base delay doubles with each consecutive failed * attempt, capped at the maximum delay. + * + * This paces only the manager's own rebuild loop for failures the SDK will + * not retry; reconnects after a mid-session drop are paced by the SDK's own + * policy (see `computeReconnectDelayMs` in `db_connection_impl`). */ export function connectionManagerReconnectDelayMs(attempt: number): number { return Math.min( @@ -71,8 +85,18 @@ type ManagedConnection = { reconnectTimer: ReturnType | null; reconnectAttempt: number; onConnect?: (conn: DbConnectionImpl) => void; - onDisconnect?: (ctx: ErrorContextInterface, error?: Error) => void; - onConnectError?: (ctx: ErrorContextInterface, error: Error) => void; + onDisconnect?: ( + ctx: ErrorContextInterface, + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void; + onConnectError?: ( + ctx: ErrorContextInterface, + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void; }; function defaultState(): ConnectionState { @@ -92,98 +116,6 @@ function defaultState(): ConnectionState { class ConnectionManagerImpl { #connections = new Map(); - constructor() { - // Auto-reconnect otherwise relies entirely on the browser firing - // `onclose` plus a `setTimeout` backoff. Both are unreliable across a - // backgrounded/frozen tab: the close event may never be delivered (the - // socket dies while the event loop is suspended), and background timers - // are heavily throttled or paused, so a scheduled reconnect can stall - // indefinitely and never resume when the window is refocused. - // - // These listeners make the manager proactively re-check liveness when the - // page comes back to the foreground / the network returns, bringing any - // stalled reconnect forward and rebuilding sockets that died silently. - if ( - typeof document !== 'undefined' && - typeof document.addEventListener === 'function' - ) { - document.addEventListener('visibilitychange', () => { - if (document.visibilityState === 'visible') { - this.#handleResume(); - } - }); - } - if ( - typeof window !== 'undefined' && - typeof window.addEventListener === 'function' - ) { - window.addEventListener('focus', this.#handleResume); - window.addEventListener('online', this.#handleResume); - // `pageshow` fires on bfcache restores, where `visibilitychange` may not. - window.addEventListener('pageshow', this.#handleResume); - } - } - - /** - * Called when the page is likely resuming from a background/frozen state: - * the tab became visible, the window regained focus, the network came back, - * or a bfcache page was restored. For each retained connection this brings a - * stalled reconnect forward immediately (resetting backoff) and rebuilds any - * socket that died silently while we were hidden. - */ - #handleResume = (): void => { - for (const managed of this.#connections.values()) { - if (managed.refCount <= 0 || managed.pendingRelease) { - continue; - } - - // A reconnect was scheduled but its timer is stuck behind background - // timer throttling / page freezing. Fire it now and reset backoff so we - // reconnect promptly instead of waiting out a (capped 30s, possibly - // paused) delay. - if (managed.reconnectTimer && !managed.connection) { - clearTimeout(managed.reconnectTimer); - managed.reconnectTimer = null; - managed.reconnectAttempt = 0; - if (managed.builder) { - this.#buildManagedConnection(managed, managed.builder); - } - continue; - } - - // We believe we're connected, but the socket may have died silently. - this.#reviveIfZombie(managed); - } - }; - - /** - * If `managed` holds a connection whose socket has entered CLOSING/CLOSED - * without a clean `onclose` (see {@link DbConnectionImpl.isSocketClosed}), - * for example because it was torn down while the tab was frozen, tear it down - * and build a fresh one immediately, resetting backoff. - */ - #reviveIfZombie(managed: ManagedConnection): void { - const connection = managed.connection; - if ( - !connection || - connection.isDisconnectRequested || - !connection.isSocketClosed - ) { - return; - } - - this.#detachCallbacks(managed, connection); - managed.connection = undefined; - // Close the dead socket in case it is only CLOSING; callbacks are already - // detached, so this won't trigger a duplicate reconnect. - connection.disconnect(); - this.#updateState(managed, { isActive: false }); - managed.reconnectAttempt = 0; - if (managed.builder) { - this.#buildManagedConnection(managed, managed.builder); - } - } - /** Generates a unique key for a connection based on URI and module name. */ static getKey(uri: string, moduleName: string): string { return `${uri}::${moduleName}`; @@ -246,7 +178,12 @@ class ConnectionManagerImpl { }); }; - managed.onDisconnect = (ctx, error) => { + // With automatic reconnection forced on, a non-undefined + // `nextReconnectAttempt` means the SDK is retrying inside the same + // connection object: mirror the event into state and leave the connection + // alone. Only when the SDK reports it will not retry does the manager + // rebuild a replacement (see #scheduleRebuild). + managed.onDisconnect = (ctx, error, nextReconnectAttempt) => { if (ctx !== managed.connection) { return; } @@ -254,10 +191,12 @@ class ConnectionManagerImpl { isActive: false, connectionError: error ?? undefined, }); - this.#scheduleReconnect(managed); + if (nextReconnectAttempt === undefined) { + this.#scheduleRebuild(managed); + } }; - managed.onConnectError = (ctx, error) => { + managed.onConnectError = (ctx, error, nextReconnectAttempt) => { if (ctx !== managed.connection) { return; } @@ -265,7 +204,9 @@ class ConnectionManagerImpl { isActive: false, connectionError: error, }); - this.#scheduleReconnect(managed); + if (nextReconnectAttempt === undefined) { + this.#scheduleRebuild(managed); + } }; } @@ -284,44 +225,24 @@ class ConnectionManagerImpl { connection: DbConnectionImpl ): void { if (managed.onConnect) { - connection.removeOnConnect(managed.onConnect as any); + connection.removeOnConnect(managed.onConnect); } if (managed.onDisconnect) { - connection.removeOnDisconnect(managed.onDisconnect as any); + connection.removeOnDisconnect(managed.onDisconnect); } if (managed.onConnectError) { - connection.removeOnConnectError(managed.onConnectError as any); + connection.removeOnConnectError(managed.onConnectError); } } - /** - * Builds a connection for `managed` from `builder`, adopting it as the - * entry's retained builder. - * - * `resumeSession` (the default) re-applies the session's current token to the - * builder first. This matters because the builder is a *long-lived template*: - * the application hands it over once, and every automatic rebuild — scheduled - * reconnect, resume-from-background, zombie-socket revival — reuses that same - * object. Its token, though, is a snapshot taken when the application built - * it, typically read out of storage at module load, before any session - * existed. Rebuilding from it verbatim would reconnect *anonymously* for any - * user whose token was issued during this page's lifetime, and the server - * would answer by minting a brand-new identity: a silent account switch, with - * no error raised on either side, curable only by a page reload. - * - * `state.token` is the token of the most recent connection (set below at - * build time, and again by `onConnect` when the server issues one), so - * re-applying it keeps every automatic rebuild on the same principal. - * - * Pass `resumeSession: false` when the caller is deliberately changing - * identity — see {@link rebuild} — so the builder's own token wins. - */ + /** Reuse the latest session token when rebuilding from the retained builder. */ #buildManagedConnection>( managed: ManagedConnection, builder: DbConnectionBuilder, { resumeSession = true }: { resumeSession?: boolean } = {} ): T { managed.builder = builder; + builder.withAutomaticReconnect(); if (resumeSession && managed.state.token) { builder.withToken(managed.state.token); } @@ -340,7 +261,8 @@ class ConnectionManagerImpl { return connection as T; } - #scheduleReconnect(managed: ManagedConnection): void { + /** Preserve framework retries for failures the core connection will not retry. */ + #scheduleRebuild(managed: ManagedConnection): void { if ( managed.refCount <= 0 || managed.pendingRelease || @@ -425,8 +347,8 @@ class ConnectionManagerImpl { * …) re-bind to the new connection automatically. * * The old connection's callbacks are detached before it is closed, so its - * disconnect event never leaks into pool state, and any pending auto-reconnect - * is cancelled (the caller is driving the reconnect explicitly). Returns the + * disconnect event never leaks into pool state, and any pending rebuild is + * cancelled (the caller is driving the replacement explicitly). Returns the * newly-built connection, or `null` if the key has no retained entry. * * @param key - Unique identifier for the connection (use getKey to generate) @@ -442,8 +364,8 @@ class ConnectionManagerImpl { } // The caller is taking over the connection lifecycle explicitly; cancel a - // deferred release or a pending auto-reconnect so neither races the fresh - // connection, and reset the backoff so the next unexpected drop starts over. + // deferred release or a pending rebuild so neither races the fresh + // connection, and reset the backoff so the next terminal failure starts over. if (managed.pendingRelease) { clearTimeout(managed.pendingRelease); managed.pendingRelease = null; diff --git a/crates/bindings-typescript/src/sdk/db_connection_builder.ts b/crates/bindings-typescript/src/sdk/db_connection_builder.ts index 282cab02d37..9af4fad1020 100644 --- a/crates/bindings-typescript/src/sdk/db_connection_builder.ts +++ b/crates/bindings-typescript/src/sdk/db_connection_builder.ts @@ -1,4 +1,8 @@ -import { DbConnectionImpl, type ConnectionEvent } from './db_connection_impl'; +import { + DbConnectionImpl, + type ConnectionEvent, + type TokenProvider, +} from './db_connection_impl'; import { EventEmitter } from './event_emitter'; import type { DbConnectionConfig, @@ -27,6 +31,8 @@ export class DbConnectionBuilder> { #compression: 'gzip' | 'brotli' | 'none' = 'gzip'; #lightMode: boolean = false; #confirmedReads?: boolean; + #automaticReconnect: boolean = false; + #tokenProvider?: TokenProvider; #createWSFn: WebSocketFactory; /** @@ -145,6 +151,26 @@ export class DbConnectionBuilder> { return this; } + /** + * Reconnect after an established connection drops, preserving handles and callbacks. + * Retries use exponential backoff until disconnect() or a terminal failure. + * Initial connection failures are not retried. Lifecycle callbacks report + * the next attempt and delay, or undefined when no retry is scheduled. + */ + withAutomaticReconnect(): this { + this.#automaticReconnect = true; + return this; + } + + /** + * Refresh the retained token before reconnecting when it is near expiry or + * rejected. The provider must return a token for the same identity. + */ + withTokenProvider(provider: TokenProvider): this { + this.#tokenProvider = provider; + return this; + } + /** * Register a callback to be invoked upon authentication with the database. * @@ -190,11 +216,19 @@ export class DbConnectionBuilder> { * console.log("Error connecting to SpacetimeDB:", error); * }); * ``` + * + * With {@link DbConnectionBuilder.withAutomaticReconnect} enabled, this + * callback also reports each failed reconnect attempt: + * `nextReconnectAttempt` is the number of the upcoming attempt and + * `nextReconnectDelayMs` the wait before it. Both are `undefined` when the + * SDK will not retry, as for a failed initial connection. */ onConnectError( callback: ( ctx: ErrorContextInterface>, - error: Error + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number ) => void ): this { this.#emitter.on('connectError', callback); @@ -225,13 +259,22 @@ export class DbConnectionBuilder> { * This is a concession to ergonomics; there's no clean place to return a `CallbackId` from this method * or from `build`. * + * With {@link DbConnectionBuilder.withAutomaticReconnect} enabled, this + * callback also reports connections lost mid-session, and the SDK keeps + * reconnecting afterwards: `nextReconnectAttempt` is the number of the + * upcoming attempt and `nextReconnectDelayMs` the wait before it. Both are + * `undefined` when the SDK will not retry, which is always the case without + * automatic reconnection. + * * @param {function(error?: Error): void} callback - The callback to invoke upon disconnection. * @throws {Error} Throws an error if called multiple times on the same `DbConnectionBuilder`. */ onDisconnect( callback: ( ctx: ErrorContextInterface>, - error?: Error | undefined + error?: Error | undefined, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number ) => void ): this { this.#emitter.on('disconnect', callback); @@ -285,6 +328,8 @@ export class DbConnectionBuilder> { confirmedReads: this.#confirmedReads, createWSFn: this.#createWSFn, remoteModule: this.remoteModule, + automaticReconnect: this.#automaticReconnect, + tokenProvider: this.#tokenProvider, }); } } diff --git a/crates/bindings-typescript/src/sdk/db_connection_impl.ts b/crates/bindings-typescript/src/sdk/db_connection_impl.ts index cf73bc270bb..efbbea89df2 100644 --- a/crates/bindings-typescript/src/sdk/db_connection_impl.ts +++ b/crates/bindings-typescript/src/sdk/db_connection_impl.ts @@ -10,6 +10,8 @@ import { ServerMessage, TableUpdateRows, UnsubscribeFlags, + type SubscribeBatchApplied, + type SubscribeBatch, } from './client_api/types'; import { ClientCache } from './client_cache.ts'; import { DbConnectionBuilder } from './db_connection_builder.ts'; @@ -61,8 +63,19 @@ import type { UntypedSchemaDef } from '../lib/schema'; import type { ProceduresView } from './procedures.ts'; import type { Values } from '../lib/type_util.ts'; import type { TransactionUpdate } from './client_api/types.ts'; -import { InternalError, SenderError } from '../lib/errors.ts'; -import type { WebSocketAdapter, WebSocketFactory } from './ws.ts'; +import type { SubscriptionEntry } from './subscription_builder_impl'; +import { + DisconnectedError, + IdentityChangedError, + InternalError, + SenderError, + UnknownCallResultError, +} from '../lib/errors.ts'; +import { + WebSocketTokenError, + type WebSocketAdapter, + type WebSocketFactory, +} from './ws.ts'; import { normalizeWsProtocol, PREFERRED_WS_PROTOCOLS, @@ -96,7 +109,21 @@ export type { ReducerEvent, }; -export type ConnectionEvent = 'connect' | 'disconnect' | 'connectError'; +export type ConnectionEventArgs = { + connect: [identity: Identity, token: string]; + disconnect: [ + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number, + ]; + connectError: [ + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number, + ]; +}; + +export type ConnectionEvent = keyof ConnectionEventArgs; export type DbConnectionConfig = { uri: URL; @@ -109,14 +136,152 @@ export type DbConnectionConfig = { lightMode: boolean; confirmedReads?: boolean; remoteModule: RemoteModule; + /** + * Whether the connection reconnects on its own after losing its socket. + * Set by {@link DbConnectionBuilder.withAutomaticReconnect}. + */ + automaticReconnect?: boolean; + /** + * Supplies a fresh token for a reconnect attempt. + * Set by {@link DbConnectionBuilder.withTokenProvider}. + */ + tokenProvider?: TokenProvider; }; +/** + * Supplies an authentication token, called before a reconnect attempt whose + * retained token is close to expiring. + */ +export type TokenProvider = () => Promise; + +/** The delay before the first reconnect attempt. */ +export const RECONNECT_INITIAL_DELAY_MS = 1_000; +/** The upper bound on the delay between reconnect attempts. */ +export const RECONNECT_MAX_DELAY_MS = 30_000; +/** The random spread applied to each reconnect delay. */ +export const RECONNECT_JITTER = 0.5; +/** + * A token is refreshed when its remaining validity falls below this fraction + * of its lifetime, or below {@link TOKEN_REFRESH_MIN_MARGIN_MS}, whichever is + * larger. + */ +const TOKEN_REFRESH_MARGIN_FRACTION = 0.05; +const TOKEN_REFRESH_MIN_MARGIN_MS = 30_000; + +/** Exponential backoff with jitter; attempt numbers start at one. */ +export function computeReconnectDelayMs( + attempt: number, + random: () => number = Math.random +): number { + const base = Math.min( + RECONNECT_INITIAL_DELAY_MS * Math.pow(2, Math.max(0, attempt - 1)), + RECONNECT_MAX_DELAY_MS + ); + const jittered = base * (1 + RECONNECT_JITTER * (2 * random() - 1)); + return Math.max(0, Math.min(jittered, RECONNECT_MAX_DELAY_MS)); +} + +/** Refresh unreadable tokens, or tokens within 5% of expiry (at least 30 seconds). */ +export function tokenNeedsRefresh( + token: string | undefined, + nowMs: number = Date.now() +): boolean { + if (!token) { + return true; + } + const claims = decodeJwtClaims(token); + if (claims === undefined || claims.exp === undefined) { + return true; + } + const expiryMs = claims.exp * 1000; + const lifetimeMs = + claims.iat !== undefined ? expiryMs - claims.iat * 1000 : undefined; + const margin = Math.max( + TOKEN_REFRESH_MIN_MARGIN_MS, + lifetimeMs !== undefined ? lifetimeMs * TOKEN_REFRESH_MARGIN_FRACTION : 0 + ); + return expiryMs - nowMs <= margin; +} + +/** Normalize an unknown thrown value or error event into an `Error`. */ +function errorFromEvent(value: unknown, fallbackMessage: string): Error { + if (value instanceof Error) { + return value; + } + if (typeof value === 'object' && value !== null) { + const message = 'message' in value ? value.message : undefined; + const error = 'error' in value ? value.error : undefined; + if (error instanceof Error) { + return error; + } + if (typeof message === 'string' && message.length > 0) { + return new Error(message); + } + } + return new Error(fallbackMessage); +} + +/** + * Whether a failure to connect is one that retrying cannot fix, so that the + * SDK stops rather than looping. + */ +function isTerminalConnectError(error: Error): boolean { + return ( + error instanceof IdentityChangedError || + error instanceof WebSocketProtocolError + ); +} + +/** Whether a failure looks like the server rejecting our credentials. */ +function isAuthError(error: Error): boolean { + return ( + error instanceof WebSocketTokenError && + (error.status === 401 || error.status === 403) + ); +} + +class WebSocketProtocolError extends Error {} + +const SESSION_BUSY_CLOSE_CODE = 4000; + +type JwtClaims = { exp?: number; iat?: number }; + +function decodeJwtClaims(token: string): JwtClaims | undefined { + const parts = token.split('.'); + if (parts.length < 2) { + return undefined; + } + try { + const payload = parts[1].replace(/-/g, '+').replace(/_/g, '/'); + const padded = payload.padEnd( + payload.length + ((4 - (payload.length % 4)) % 4), + '=' + ); + const json = + typeof atob === 'function' + ? atob(padded) + : Buffer.from(padded, 'base64').toString('binary'); + const claims: unknown = JSON.parse(json); + if (typeof claims !== 'object' || claims === null) { + return undefined; + } + const exp = 'exp' in claims ? claims.exp : undefined; + const iat = 'iat' in claims ? claims.iat : undefined; + return { + exp: typeof exp === 'number' && Number.isFinite(exp) ? exp : undefined, + iat: typeof iat === 'number' && Number.isFinite(iat) ? iat : undefined, + }; + } catch { + return undefined; + } +} + type ProcedureCallback = (result: ProcedureResultMessage['result']) => void; type Deferred = { promise: Promise; resolve: (value: T | PromiseLike) => void; - reject: (reason?: unknown) => void; + reject: (reason: Error | string) => void; }; const TEXT_ENCODER = new TextEncoder(); @@ -177,8 +342,9 @@ export class DbConnectionImpl * Whether the underlying websocket has entered `CLOSING` (2) or `CLOSED` * (3). This becomes true even when the browser never delivered an * `onclose` event, for example if the socket was torn down while the tab was - * frozen or the machine was asleep. The `ConnectionManager` uses this to detect - * such "zombie" connections when the page resumes and to force a reconnect. + * frozen or the machine was asleep. The liveness listeners (see + * {@link DbConnectionImpl.#installLivenessListeners}) use this to detect + * such "zombie" sockets when the page resumes and to force a reconnect. * * Returns false while the socket is still `CONNECTING`/`OPEN`, or before * the socket has been created. @@ -231,6 +397,29 @@ export class DbConnectionImpl connectionId: ConnectionId = ConnectionId.random(); #connectionIdHex = this.connectionId.toHexString(); + // Stable across sockets; each attempt still gets a fresh ConnectionId. + #sessionIdHex = ConnectionId.random().toHexString(); + #automaticReconnect: boolean; + #tokenProvider?: TokenProvider; + #wsBaseUrl: URL; + #nameOrAddress: string; + #createWSFn: WebSocketFactory; + #compression: 'gzip' | 'brotli' | 'none'; + #lightMode: boolean; + #confirmedReads?: boolean; + // Invalidates events and asynchronous work from discarded sockets. + #socketGeneration = 0; + #hasEverConnected = false; + #connectionEnded = false; + #reconnectAttempt = 0; + #reconnectTimer?: ReturnType; + #forceTokenRefresh = false; + #usedFreshToken = false; + #livenessCleanup?: () => void; + #pendingReplay?: { requestId: number; querySetIds: Set }; + #preparingReplay = false; + #socketEstablished = false; + // These fields are meant to be strictly private. #queryId = 0; #requestId = 0; @@ -250,6 +439,10 @@ export class DbConnectionImpl >(); #reducerCallInfo = new Map(); #procedureCallbacks = new Map(); + /** + * Reject pending reducer and procedure calls when the socket drops. + */ + #pendingCallRejecters = new Map void>(); #rowDeserializers: Record>; #rowIdMetadata: Record< string, @@ -291,6 +484,8 @@ export class DbConnectionImpl compression, lightMode, confirmedReads, + automaticReconnect, + tokenProvider, }: DbConnectionConfig) { stdbLogger('info', 'Connecting to SpacetimeDB WS...'); @@ -357,43 +552,323 @@ export class DbConnectionImpl ); } - url.searchParams.set('connection_id', this.#connectionIdHex); - this.clientCache = new ClientCache(); this.db = this.#makeDbView(); this.reducers = this.#makeReducers(remoteModule); this.procedures = this.#makeProcedures(remoteModule); - this.wsPromise = createWSFn({ - url, - nameOrAddress, - wsProtocol: [...PREFERRED_WS_PROTOCOLS], - authToken: token, - compression: compression, - lightMode: lightMode, - confirmedReads: confirmedReads, - }) - .then(v => { - this.ws = v; - - this.ws.onclose = () => { - this.isActive = false; - this.#emitter.emit('disconnect', this); - }; - this.ws.onerror = (e: ErrorEvent) => { - this.isActive = false; - this.#emitter.emit('connectError', this, e); - }; - this.ws.onopen = this.#handleOnOpen.bind(this); - this.ws.onmessage = this.#handleOnMessage.bind(this); - return v; - }) - .catch(e => { - stdbLogger('error', 'Error connecting to SpacetimeDB WS'); - this.#emitter.emit('connectError', this, e); + this.#automaticReconnect = automaticReconnect ?? false; + this.#preparingReplay = this.#automaticReconnect; + this.#tokenProvider = tokenProvider; + this.#wsBaseUrl = url; + this.#nameOrAddress = nameOrAddress; + this.#createWSFn = createWSFn; + this.#compression = compression; + this.#lightMode = lightMode; + this.#confirmedReads = confirmedReads; + + this.wsPromise = this.#openSocket(); + } + + async #openSocket(): Promise { + const generation = ++this.#socketGeneration; + if (this.#hasEverConnected) { + this.#setConnectionId(ConnectionId.random()); + } + this.#usedFreshToken = false; + const url = new URL(this.#wsBaseUrl.toString()); + url.searchParams.set('connection_id', this.#connectionIdHex); + try { + const authToken = await this.#tokenForAttempt(); + if (generation !== this.#socketGeneration || this.#connectionEnded) { return undefined; + } + this.token = authToken; + const ws = await this.#createWSFn({ + url, + nameOrAddress: this.#nameOrAddress, + wsProtocol: [...PREFERRED_WS_PROTOCOLS], + authToken, + compression: this.#compression, + lightMode: this.#lightMode, + confirmedReads: this.#confirmedReads, + connectionId: this.#connectionIdHex, + sessionId: this.#automaticReconnect ? this.#sessionIdHex : undefined, }); + + if (generation !== this.#socketGeneration || this.#connectionEnded) { + // A newer attempt superseded this socket while it was opening, or the + // application disconnected in the meantime. + ws.close(); + return undefined; + } + + this.ws = ws; + this.#socketEstablished = false; + const isCurrent = (): boolean => + generation === this.#socketGeneration && !this.#connectionEnded; + const handleLoss = (error: Error, isErrorEvent: boolean): void => { + if (!isCurrent()) return; + this.isActive = false; + if (!this.#automaticReconnect) { + this.#emitter.emit( + isErrorEvent ? 'connectError' : 'disconnect', + this, + isErrorEvent ? error : undefined + ); + } else if (this.#socketEstablished) { + this.#handleConnectionLoss(error); + } else { + this.#handleAttemptFailure(error); + } + }; + + ws.onclose = event => { + if (!isCurrent()) return; + if ( + this.#automaticReconnect && + this.#hasEverConnected && + !this.#socketEstablished && + event.code === SESSION_BUSY_CLOSE_CODE + ) { + // The server is tearing down the previous session holder. + this.#discardSocket(); + const delayMs = computeReconnectDelayMs(1); + this.#emitter.emit( + 'connectError', + this, + new Error('Session busy'), + this.#reconnectAttempt, + delayMs + ); + this.#scheduleReconnect(this.#reconnectAttempt, delayMs); + return; + } + const message = `WebSocket closed (code ${event.code}${event.reason ? `: ${event.reason}` : ''})`; + const error = [1002, 1003, 1007, 1008].includes(event.code) + ? new WebSocketProtocolError(message) + : new Error(message); + handleLoss(error, false); + }; + ws.onerror = event => + handleLoss(errorFromEvent(event, 'WebSocket error'), true); + ws.onopen = () => { + if (isCurrent()) this.#handleOnOpen(); + }; + ws.onmessage = message => { + if (isCurrent()) this.#handleOnMessage(message); + }; + return ws; + } catch (e) { + if (generation !== this.#socketGeneration || this.#connectionEnded) { + return undefined; + } + stdbLogger('error', 'Error connecting to SpacetimeDB WS'); + this.isActive = false; + this.#handleAttemptFailure(errorFromEvent(e, 'Failed to connect')); + return undefined; + } + } + + async #tokenForAttempt(): Promise { + // The initial connection uses the token the application configured. + if (!this.#tokenProvider || !this.#hasEverConnected) { + return this.token; + } + if (!this.#forceTokenRefresh && !tokenNeedsRefresh(this.token)) { + return this.token; + } + const token = await this.#tokenProvider(); + this.#forceTokenRefresh = false; + this.#usedFreshToken = true; + return token; + } + + #handleConnectionLoss(error: Error): void { + if (this.#connectionEnded) { + return; + } + this.#discardSocket(); + this.#failInFlightCalls(new UnknownCallResultError()); + + const willReconnect = + this.#automaticReconnect && + !this.isDisconnectRequested && + !(error instanceof WebSocketProtocolError); + + if (!willReconnect) { + this.#endConnection(error); + return; + } + + const attempt = this.#reconnectAttempt + 1; + const delayMs = computeReconnectDelayMs(attempt); + this.#emitter.emit('disconnect', this, error, attempt, delayMs); + this.#scheduleReconnect(attempt, delayMs); + } + + #handleAttemptFailure(error: Error): void { + if (this.#connectionEnded) { + return; + } + this.#discardSocket(); + this.#failInFlightCalls(new UnknownCallResultError()); + + // A failed *initial* connection is not retried: the cause is usually a + // misconfigured URI or database name that no retry will fix. + const willReconnect = + this.#automaticReconnect && + this.#hasEverConnected && + !this.isDisconnectRequested && + !isTerminalConnectError(error) && + !(isAuthError(error) && (!this.#tokenProvider || this.#usedFreshToken)); + + if (!willReconnect) { + this.#endConnection(undefined, { alreadyReported: true }); + this.#emitter.emit('connectError', this, error); + return; + } + + // A rejected token is worth one forced refresh: the retained token may + // have been revoked, or the clock may be skewed. + if (this.#tokenProvider && isAuthError(error)) { + this.#forceTokenRefresh = true; + } + + const attempt = this.#reconnectAttempt + 1; + const delayMs = computeReconnectDelayMs(attempt); + this.#emitter.emit('connectError', this, error, attempt, delayMs); + this.#scheduleReconnect(attempt, delayMs); + } + + #scheduleReconnect(attempt: number, delayMs: number): void { + if (this.#connectionEnded || this.isDisconnectRequested) return; + this.#reconnectAttempt = attempt; + this.#clearReconnectTimer(); + this.#reconnectTimer = setTimeout(() => { + this.#reconnectTimer = undefined; + if (this.#connectionEnded || this.isDisconnectRequested) { + return; + } + this.wsPromise = this.#openSocket(); + }, delayMs); + } + + #clearReconnectTimer(): void { + if (this.#reconnectTimer !== undefined) { + clearTimeout(this.#reconnectTimer); + this.#reconnectTimer = undefined; + } + } + + #endConnection( + disconnectError: Error | undefined, + options?: { alreadyReported?: boolean } + ): void { + if (this.#connectionEnded) { + return; + } + this.#connectionEnded = true; + this.isActive = false; + this.#clearReconnectTimer(); + this.#discardSocket(); + this.#removeLivenessListeners(); + this.#failInFlightCalls(new UnknownCallResultError()); + if (!options?.alreadyReported) { + this.#emitter.emit('disconnect', this, disconnectError); + } + } + + #failInFlightCalls(error: Error): void { + const rejecters = [...this.#pendingCallRejecters.values()]; + this.#pendingCallRejecters.clear(); + this.#reducerCallbacks.clear(); + this.#reducerCallInfo.clear(); + this.#procedureCallbacks.clear(); + + for (const reject of rejecters) { + reject(error); + } + } + + /** True while the SDK is between a lost connection and a completed reconnect. */ + get isReconnecting(): boolean { + return ( + this.#automaticReconnect && + this.#hasEverConnected && + !this.isActive && + !this.#connectionEnded + ); + } + + #discardSocket(): void { + this.#socketGeneration += 1; + this.isActive = false; + this.#socketEstablished = false; + this.#pendingReplay = undefined; + this.#preparingReplay = this.#hasEverConnected; + this.#outboundQueue.length = 0; + this.#inboundQueue.length = 0; + const ws = this.ws; + this.ws = undefined; + ws?.close(); + for (const [id, entry] of this.#subscriptionManager.subscriptions) { + if (entry.unsubscribeRequested) this.#endSubscription(id); + } + } + + // Resume events recover sockets that closed silently while the page was frozen. + #installLivenessListeners(): void { + if (!this.#automaticReconnect || this.#livenessCleanup) { + return; + } + const doc = typeof document !== 'undefined' ? document : undefined; + const win = typeof window !== 'undefined' ? window : undefined; + if (!doc && !win) { + return; + } + + const onResume = (): void => this.#handleLivenessResume(); + const onVisibilityChange = (): void => { + if (doc?.visibilityState === 'visible') { + onResume(); + } + }; + + doc?.addEventListener('visibilitychange', onVisibilityChange); + win?.addEventListener('focus', onResume); + win?.addEventListener('online', onResume); + win?.addEventListener('pageshow', onResume); + + this.#livenessCleanup = () => { + doc?.removeEventListener('visibilitychange', onVisibilityChange); + win?.removeEventListener('focus', onResume); + win?.removeEventListener('online', onResume); + win?.removeEventListener('pageshow', onResume); + this.#livenessCleanup = undefined; + }; + } + + #removeLivenessListeners(): void { + this.#livenessCleanup?.(); + } + + #handleLivenessResume(): void { + if (this.#connectionEnded || this.isDisconnectRequested) { + return; + } + if (this.isSocketClosed) { + const error = new Error('WebSocket closed while suspended'); + if (this.#socketEstablished) this.#handleConnectionLoss(error); + else this.#handleAttemptFailure(error); + return; + } + if (this.#reconnectTimer !== undefined) { + // Retry now rather than waiting out a backoff computed before the pause. + this.#clearReconnectTimer(); + this.wsPromise = this.#openSocket(); + } } #getNextQueryId = () => { @@ -531,7 +1006,14 @@ export class DbConnectionImpl this.#subscriptionManager.subscriptions.set(querySetId, { handle, emitter: handleEmitter, + // Retained so the subscription can be replayed after a reconnect. + querySql: [...querySql], }); + if (!this.#preparingReplay) this.#sendSubscription(querySetId, querySql); + return querySetId; + } + + #sendSubscription(querySetId: number, querySql: string[]): void { const requestId = this.#getNextRequestId(); this.#sendMessage( ClientMessage.Subscribe({ @@ -540,10 +1022,145 @@ export class DbConnectionImpl requestId, }) ); - return querySetId; + } + + #replaySubscriptions(): void { + const entries = [...this.#subscriptionManager.subscriptions.entries()]; + const sets: SubscribeBatch['sets'] = []; + const replayed = new Map>(); + for (const [, entry] of entries) { + const querySetId = this.#getNextQueryId(); + entry.handle.rebindQuerySetId(querySetId); + replayed.set(querySetId, entry); + sets.push({ + querySetId: { id: querySetId }, + queryStrings: entry.querySql, + }); + } + this.#subscriptionManager.subscriptions = replayed; + + const requestId = this.#getNextRequestId(); + this.#pendingReplay = { + requestId, + querySetIds: new Set(replayed.keys()), + }; + this.#preparingReplay = false; + if (sets.length === 0) { + this.#applyReplayBatch({ requestId, results: [] }); + return; + } + this.#sendMessage( + ClientMessage.SubscribeBatch({ + requestId, + sets, + }) + ); + } + + #applyReplayBatch(applied: SubscribeBatchApplied): void { + const pending = this.#pendingReplay; + const resultIds = new Set( + applied.results.map(result => result.querySetId.id) + ); + if ( + !pending || + pending.requestId !== applied.requestId || + resultIds.size !== applied.results.length || + resultIds.size !== pending.querySetIds.size || + [...resultIds].some(id => !pending.querySetIds.has(id)) + ) { + this.#handleProtocolError( + new Error('Unexpected subscription replay response') + ); + return; + } + this.#pendingReplay = undefined; + + const event: Event = { + id: this.#nextEventId(), + tag: 'SubscribeApplied', + }; + const eventContext = this.#makeEventContext(event); + + // The removal half: every row the cache holds from the old connection. + // Only tables which have been populated exist in the cache. + const tableUpdates: CacheTableUpdate[] = []; + for (const [tableName, table] of this.clientCache.tables) { + const operations = table.snapshotDeleteOperations(); + if (operations.length > 0) { + tableUpdates.push({ tableName, operations }); + } + } + + // The addition half: the rows of every set which applied. + const failures: { + entry: SubscriptionEntry; + error: string; + }[] = []; + for (const result of applied.results) { + const entry = this.#subscriptionManager.subscriptions.get( + result.querySetId.id + ); + if (!entry) { + continue; + } + if (result.outcome.tag === 'Error') { + // The set is not registered, so drop it and report it below. + this.#subscriptionManager.subscriptions.delete(result.querySetId.id); + failures.push({ entry, error: result.outcome.value }); + continue; + } + tableUpdates.push( + ...this.#queryRowsToTableUpdates(result.outcome.value, 'insert') + ); + } + + const merged = this.#mergeTableUpdates(tableUpdates); + const callbacks = this.#applyTableUpdates(merged, eventContext, { + // A row which is unchanged across the outage appears as a + // delete/insert pair and must produce no callback. + skipIdenticalUpdates: true, + }); + const { event: _, ...subscriptionEventContext } = eventContext; + for (const [querySetId, entry] of this.#subscriptionManager.subscriptions) { + if (pending.querySetIds.has(querySetId)) { + entry.emitter.emit('applied', subscriptionEventContext); + } + } + for (const { entry, error: message } of failures) { + const error = Error(message); + const errorEventContext = this.#makeEventContext({ + id: this.#nextEventId(), + tag: 'Error', + value: error, + }); + entry.emitter.emit( + 'error', + { ...errorEventContext, event: error }, + error + ); + } + this.#dispatchPendingCallbacks(callbacks); + } + + #endSubscription(querySetId: number): void { + const entry = this.#subscriptionManager.subscriptions.get(querySetId); + this.#subscriptionManager.subscriptions.delete(querySetId); + const { event: _, ...ctx } = this.#makeEventContext({ + id: this.#nextEventId(), + tag: 'UnsubscribeApplied', + }); + entry?.emitter.emit('end', ctx); } unregisterSubscription(querySetId: number): void { + const entry = this.#subscriptionManager.subscriptions.get(querySetId); + if (!entry) return; + entry.unsubscribeRequested = true; + if (this.#automaticReconnect && (!this.isActive || this.#preparingReplay)) { + this.#endSubscription(querySetId); + return; + } const requestId = this.#getNextRequestId(); this.#sendMessage( ClientMessage.Unsubscribe({ @@ -772,6 +1389,12 @@ export class DbConnectionImpl } } + #rejectCallIfDisconnected(): DisconnectedError | undefined { + return this.#automaticReconnect && !this.isActive + ? new DisconnectedError() + : undefined; + } + #sendMessage(message: ClientMessage): void { const writer = this.#clientMessageEncoder; writer.clear(); @@ -840,15 +1463,17 @@ export class DbConnectionImpl if (this.ws) { this.#negotiatedWsProtocol = normalizeWsProtocol(this.ws.protocol); } - this.isActive = true; - if (this.ws) { + this.isActive = !this.#automaticReconnect; + this.#installLivenessListeners(); + if (this.ws && this.isActive) { this.#flushOutboundQueue(this.ws); } } #applyTableUpdates( tableUpdates: CacheTableUpdate[], - eventContext: EventContextInterface + eventContext: EventContextInterface, + options?: { skipIdenticalUpdates?: boolean } ): PendingCallback[] { const pendingCallbacks: PendingCallback[] = []; for (const tableUpdate of tableUpdates) { @@ -860,7 +1485,8 @@ export class DbConnectionImpl tableUpdate.operations as Operation< RowType> >[], - eventContext + eventContext, + options ); for (const callback of newCallbacks) { pendingCallbacks.push(callback); @@ -902,14 +1528,54 @@ export class DbConnectionImpl 'trace', () => `Processing server message: ${stringify(serverMessage)}` ); + if ( + this.#automaticReconnect && + (serverMessage.tag === 'InitialConnection') === this.#socketEstablished + ) { + this.#handleProtocolError( + new Error('Unexpected message during connection handshake') + ); + return; + } switch (serverMessage.tag) { case 'InitialConnection': { + const isReconnect = this.#hasEverConnected; + if ( + isReconnect && + this.identity && + !this.identity.isEqual(serverMessage.value.identity) + ) { + // Retrying cannot recover the old identity, so stop here rather + // than serving the application someone else's data. + const error = new IdentityChangedError(); + this.#endConnection(undefined, { alreadyReported: true }); + this.#emitter.emit('connectError', this, error); + break; + } + this.identity = serverMessage.value.identity; + // The server issues a token on the first connection; retain it so + // reconnects present the same identity. if (!this.token && serverMessage.value.token) { this.token = serverMessage.value.token; } this.#setConnectionId(serverMessage.value.connectionId); + this.isActive = true; + this.#hasEverConnected = true; + this.#socketEstablished = true; + // A connection was established, so the backoff schedule starts over. + this.#reconnectAttempt = 0; this.#emitter.emit('connect', this, this.identity, this.token); + if (this.#connectionEnded) break; + if (isReconnect) { + this.#replaySubscriptions(); + } else if (this.#preparingReplay) { + this.#preparingReplay = false; + for (const [id, entry] of this.#subscriptionManager.subscriptions) { + this.#sendSubscription(id, entry.querySql); + } + } + if (this.ws) this.#flushOutboundQueue(this.ws); break; } case 'SubscribeApplied': { @@ -1073,13 +1739,24 @@ export class DbConnectionImpl ); break; } + case 'SubscribeBatchApplied': { + this.#applyReplayBatch(serverMessage.value); + break; + } } } #processV2Message(data: Uint8Array): void { const reader = this.#messageReader; reader.reset(data); - this.#processServerMessage(ServerMessage.deserialize(reader)); + let message: ServerMessage; + try { + message = ServerMessage.deserialize(reader); + } catch (cause) { + this.#handleProtocolError(cause); + return; + } + this.#processServerMessage(message); } #processMessage(data: Uint8Array): void { @@ -1088,17 +1765,28 @@ export class DbConnectionImpl return; } - const messageCount = forEachServerMessageV3( - this.#messageReader, - data, - serverMessage => { + let dispatching = false; + const generation = this.#socketGeneration; + try { + forEachServerMessageV3(this.#messageReader, data, serverMessage => { + if (generation !== this.#socketGeneration) return; + dispatching = true; this.#processServerMessage(serverMessage); - } - ); - stdbLogger( - 'trace', - () => `Processing server v3 payload with ${messageCount} message(s)` - ); + dispatching = false; + }); + } catch (cause) { + if (dispatching) throw cause; + this.#handleProtocolError(cause); + } + } + + #handleProtocolError(cause: unknown): void { + if (!this.#automaticReconnect) throw cause; + const error = new WebSocketProtocolError('Invalid server message', { + cause, + }); + if (this.#socketEstablished) this.#handleConnectionLoss(error); + else this.#handleAttemptFailure(error); } /** @@ -1166,6 +1854,10 @@ export class DbConnectionImpl argsBuffer: Uint8Array, reducerArgs?: object ): Promise { + const rejected = this.#rejectCallIfDisconnected(); + if (rejected) { + return Promise.reject(rejected); + } const { promise, resolve, reject } = createDeferred(); const requestId = this.#getNextRequestId(); this.#sendCallReducerMessage(requestId, encodedReducerName, argsBuffer); @@ -1175,7 +1867,9 @@ export class DbConnectionImpl args: reducerArgs, }); } + this.#pendingCallRejecters.set(requestId, reject); this.#reducerCallbacks.set(requestId, result => { + this.#pendingCallRejecters.delete(requestId); if (result.tag === 'Ok' || result.tag === 'OkEmpty') { resolve(); } else { @@ -1201,6 +1895,10 @@ export class DbConnectionImpl argsBuffer: Uint8Array, reducerArgs?: object ): Promise { + const rejected = this.#rejectCallIfDisconnected(); + if (rejected) { + return Promise.reject(rejected); + } const { promise, resolve, reject } = createDeferred(); const requestId = this.#getNextRequestId(); const message = ClientMessage.CallReducer({ @@ -1216,7 +1914,9 @@ export class DbConnectionImpl args: reducerArgs, }); } + this.#pendingCallRejecters.set(requestId, reject); this.#reducerCallbacks.set(requestId, result => { + this.#pendingCallRejecters.delete(requestId); if (result.tag === 'Ok' || result.tag === 'OkEmpty') { resolve(); } else { @@ -1282,10 +1982,16 @@ export class DbConnectionImpl encodedProcedureName: Uint8Array, argsBuffer: Uint8Array ): Promise { + const rejected = this.#rejectCallIfDisconnected(); + if (rejected) { + return Promise.reject(rejected); + } const { promise, resolve, reject } = createDeferred(); const requestId = this.#getNextRequestId(); this.#sendCallProcedureMessage(requestId, encodedProcedureName, argsBuffer); + this.#pendingCallRejecters.set(requestId, reject); this.#procedureCallbacks.set(requestId, result => { + this.#pendingCallRejecters.delete(requestId); if (result.tag === 'Ok') { resolve(result.value); } else { @@ -1299,6 +2005,10 @@ export class DbConnectionImpl procedureName: string, argsBuffer: Uint8Array ): Promise { + const rejected = this.#rejectCallIfDisconnected(); + if (rejected) { + return Promise.reject(rejected); + } const { promise, resolve, reject } = createDeferred(); const requestId = this.#getNextRequestId(); const message = ClientMessage.CallProcedure({ @@ -1309,7 +2019,9 @@ export class DbConnectionImpl flags: 0, }); this.#sendMessage(message); + this.#pendingCallRejecters.set(requestId, reject); this.#procedureCallbacks.set(requestId, result => { + this.#pendingCallRejecters.delete(requestId); if (result.tag === 'Ok') { resolve(result.value); } else { @@ -1355,55 +2067,87 @@ export class DbConnectionImpl */ disconnect(): void { this.isDisconnectRequested = true; - this.wsPromise.then(ws => ws?.close()); + if (this.#automaticReconnect) { + if (this.#connectionEnded) { + this.#emitter.emit('disconnect', this); + } else { + this.#endConnection(undefined); + } + } else { + this.wsPromise.then(ws => ws?.close()); + } } - private on( - eventName: ConnectionEvent, - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + private on( + eventName: E, + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs[E] + ) => void ): void { this.#emitter.on(eventName, callback); } - private off( - eventName: ConnectionEvent, - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + private off( + eventName: E, + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs[E] + ) => void ): void { this.#emitter.off(eventName, callback); } private onConnect( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['connect'] + ) => void ): void { this.#emitter.on('connect', callback); } private onDisconnect( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['disconnect'] + ) => void ): void { this.#emitter.on('disconnect', callback); } private onConnectError( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['connectError'] + ) => void ): void { this.#emitter.on('connectError', callback); } removeOnConnect( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['connect'] + ) => void ): void { this.#emitter.off('connect', callback); } removeOnDisconnect( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['disconnect'] + ) => void ): void { this.#emitter.off('disconnect', callback); } removeOnConnectError( - callback: (ctx: DbConnectionImpl, ...args: any[]) => void + callback: ( + ctx: DbConnectionImpl, + ...args: ConnectionEventArgs['connectError'] + ) => void ): void { this.#emitter.off('connectError', callback); } diff --git a/crates/bindings-typescript/src/sdk/subscription_builder_impl.ts b/crates/bindings-typescript/src/sdk/subscription_builder_impl.ts index f7e271b4e85..cf72dad84b2 100644 --- a/crates/bindings-typescript/src/sdk/subscription_builder_impl.ts +++ b/crates/bindings-typescript/src/sdk/subscription_builder_impl.ts @@ -160,14 +160,15 @@ export class SubscriptionBuilderImpl { export type SubscribeEvent = 'applied' | 'error' | 'end'; +export type SubscriptionEntry = { + handle: SubscriptionHandleImpl; + emitter: EventEmitter; + querySql: string[]; + unsubscribeRequested?: boolean; +}; + export class SubscriptionManager { - subscriptions: Map< - number, - { - handle: SubscriptionHandleImpl; - emitter: EventEmitter; - } - > = new Map(); + subscriptions: Map> = new Map(); } export class SubscriptionHandleImpl { @@ -210,6 +211,11 @@ export class SubscriptionHandleImpl { ); } + /** @internal Rebind the retained handle to its replayed query set. */ + rebindQuerySetId(querySetId: number): void { + this.#querySetId = querySetId; + } + /** * Consumes self and issues an `Unsubscribe` message, * removing this query from the client's set of subscribed queries. @@ -220,7 +226,6 @@ export class SubscriptionHandleImpl { throw new Error('Unsubscribe has already been called'); } this.#unsubscribeCalled = true; - this.db.unregisterSubscription(this.#querySetId); this.#emitter.on( 'end', (_ctx: SubscriptionEventContextInterface) => { @@ -228,6 +233,7 @@ export class SubscriptionHandleImpl { this.#activeState = false; } ); + this.db.unregisterSubscription(this.#querySetId); } /** @@ -250,7 +256,6 @@ export class SubscriptionHandleImpl { throw new Error('Unsubscribe has already been called'); } this.#unsubscribeCalled = true; - this.db.unregisterSubscription(this.#querySetId); this.#emitter.on( 'end', (ctx: SubscriptionEventContextInterface) => { @@ -259,6 +264,7 @@ export class SubscriptionHandleImpl { onEnd(ctx); } ); + this.db.unregisterSubscription(this.#querySetId); } /** diff --git a/crates/bindings-typescript/src/sdk/table_cache.ts b/crates/bindings-typescript/src/sdk/table_cache.ts index 543d3fb78ed..32184e73426 100644 --- a/crates/bindings-typescript/src/sdk/table_cache.ts +++ b/crates/bindings-typescript/src/sdk/table_cache.ts @@ -262,11 +262,30 @@ export class TableCacheImpl< return this.iter(); } + /** Delete every cached reference, including overlaps between subscriptions. */ + snapshotDeleteOperations = (): Operation< + RowType> + >[] => { + const operations: Operation< + RowType> + >[] = []; + for (const [rowId, [row, refCount]] of this.rows) { + for (let i = 0; i < refCount; i++) { + operations.push({ type: 'delete', rowId, row }); + } + } + return operations; + }; + applyOperations = ( operations: Operation< RowType> >[], - ctx: EventContextInterface + ctx: EventContextInterface, + options?: { + /** Suppress unchanged rows during reconnect reconciliation. */ + skipIdenticalUpdates?: boolean; + } ): PendingCallback[] => { const pendingCallbacks: PendingCallback[] = []; @@ -286,7 +305,7 @@ export class TableCacheImpl< return pendingCallbacks; } - if (this.hasPrimaryKey) { + if (this.hasPrimaryKey || options?.skipIdenticalUpdates) { const insertMap = new Map< ComparablePrimitive, [ @@ -322,7 +341,8 @@ export class TableCacheImpl< ctx, primaryKey, insertOp.row, - refCountDelta + refCountDelta, + options?.skipIdenticalUpdates ); if (maybeCb) { pendingCallbacks.push(maybeCb); @@ -363,7 +383,8 @@ export class TableCacheImpl< ctx: EventContextInterface, rowId: ComparablePrimitive, newRow: RowType>, - refCountDelta: number = 0 + refCountDelta: number = 0, + skipIfIdentical: boolean = false ): PendingCallback | undefined => { const existingEntry = this.rows.get(rowId); if (!existingEntry) { @@ -384,6 +405,11 @@ export class TableCacheImpl< return undefined; } this.rows.set(rowId, [newRow, refCount]); + if (skipIfIdentical && deepEqual(oldRow, newRow)) { + // The row is unchanged; the reference count was adjusted above but no + // callback fires. + return undefined; + } // This indicates something is wrong, so we could arguably crash here. if (previousCount === 0) { stdbLogger( diff --git a/crates/bindings-typescript/src/sdk/websocket_test_adapter.ts b/crates/bindings-typescript/src/sdk/websocket_test_adapter.ts index 257f9cb2806..07ca43b3a89 100644 --- a/crates/bindings-typescript/src/sdk/websocket_test_adapter.ts +++ b/crates/bindings-typescript/src/sdk/websocket_test_adapter.ts @@ -1,7 +1,7 @@ import BinaryReader from '../lib/binary_reader.ts'; import BinaryWriter from '../lib/binary_writer.ts'; import { ClientMessage, ServerMessage } from './client_api/types'; -import type { WebSocketAdapter, WebSocketFactory } from './ws'; +import type { WebSocketAdapter, WebSocketArgs, WebSocketFactory } from './ws'; import { PREFERRED_WS_PROTOCOLS, V3_WS_PROTOCOL } from './websocket_protocols'; import { decodeClientMessagesV3, @@ -11,6 +11,9 @@ import { class WebsocketTestAdapter implements WebSocketAdapter { protocol: string = ''; + /** The arguments the connection passed to `openWebSocket`, for assertions. */ + connectArgs?: WebSocketArgs; + // WebSocket.CLOSED (3) / WebSocket.OPEN (1). Uses literals rather than the // `WebSocket` global, which is not defined when these tests run under Node. get readyState(): number { @@ -51,7 +54,15 @@ class WebsocketTestAdapter implements WebSocketAdapter { } error(error: Error): void { - this.#onerror(error as unknown as ErrorEvent); + this.#onerror( + Object.assign(new Event('error'), { + error, + message: error.message, + filename: '', + lineno: 0, + colno: 0, + }) + ); } send(message: Uint8Array): void { @@ -70,12 +81,31 @@ class WebsocketTestAdapter implements WebSocketAdapter { } close(): void { + this.serverClose(1000, 'normal closure', true); + } + + /** + * Simulate a close initiated by the server or the network, with an + * arbitrary close code (e.g. an abnormal closure or an + * application-specific code such as session-expired). + */ + serverClose( + code: number, + reason: string = '', + wasClean: boolean = false + ): void { + this.closed = true; + this.#onclose( + Object.assign(new Event('close'), { code, reason, wasClean }) + ); + } + + /** + * Mark the socket as closed without delivering any event, simulating a + * socket that died while the page was suspended (a "zombie" socket). + */ + dieSilently(): void { this.closed = true; - this.#onclose({ - code: 1000, - reason: 'normal closure', - wasClean: true, - } as CloseEvent); } acceptConnection(): void { @@ -98,16 +128,52 @@ class WebsocketTestAdapter implements WebSocketAdapter { this.#onmessage({ data: outboundData }); } - openWebSocket: WebSocketFactory = async ({ wsProtocol }) => { - const negotiatedProtocol = wsProtocol.find(protocol => + openWebSocket: WebSocketFactory = async args => { + const negotiatedProtocol = args.wsProtocol.find(protocol => this.supportedProtocols.includes(protocol) ); if (!negotiatedProtocol) { throw new Error('No compatible websocket protocol'); } this.protocol = negotiatedProtocol; + this.connectArgs = args; return this; }; } +/** + * A websocket factory that hands out a fresh {@link WebsocketTestAdapter} per + * connection attempt and records them all. Used to test automatic + * reconnection, where each attempt opens a new socket. + */ +export class WebsocketTestAdapterFactory { + /** Every adapter created so far, in creation order. */ + sockets: WebsocketTestAdapter[] = []; + /** + * When set, the next `openWebSocket` calls reject with this error instead + * of producing a socket (simulating an unreachable server or a failed + * token exchange). + */ + connectError?: Error; + + /** The most recently created adapter. */ + get current(): WebsocketTestAdapter { + const socket = this.sockets[this.sockets.length - 1]; + if (!socket) { + throw new Error('No websocket has been opened yet'); + } + return socket; + } + + openWebSocket: WebSocketFactory = async args => { + if (this.connectError) { + throw this.connectError; + } + const adapter = new WebsocketTestAdapter(); + await adapter.openWebSocket(args); + this.sockets.push(adapter); + return adapter; + }; +} + export default WebsocketTestAdapter; diff --git a/crates/bindings-typescript/src/sdk/ws.ts b/crates/bindings-typescript/src/sdk/ws.ts index 99d1b688ece..96eb7e5fe51 100644 --- a/crates/bindings-typescript/src/sdk/ws.ts +++ b/crates/bindings-typescript/src/sdk/ws.ts @@ -49,6 +49,16 @@ export interface WebSocketAdapter { set onerror(handler: (msg: ErrorEvent) => void); } +export class WebSocketTokenError extends Error { + constructor( + readonly status: number, + statusText: string + ) { + super(`Failed to verify token: ${status} ${statusText}`); + this.name = 'WebSocketTokenError'; + } +} + export interface WebSocketArgs { url: URL; wsProtocol: string[]; @@ -57,6 +67,10 @@ export interface WebSocketArgs { compression: 'gzip' | 'brotli' | 'none'; lightMode: boolean; confirmedReads?: boolean; + /** Hex-encoded id for this socket. */ + connectionId?: string; + /** Stable session id, sent only when automatic reconnection is enabled. */ + sessionId?: string; } export type WebSocketFactory = ( args: WebSocketArgs @@ -74,6 +88,8 @@ export async function openWebSocket({ compression, lightMode, confirmedReads, + connectionId, + sessionId, }: WebSocketArgs): Promise { const headers = new Headers(); @@ -92,7 +108,7 @@ export async function openWebSocket({ const { token } = await response.json(); temporaryAuthToken = token; } else { - throw new Error(`Failed to verify token: ${response.statusText}`); + throw new WebSocketTokenError(response.status, response.statusText); } } @@ -110,6 +126,14 @@ export async function openWebSocket({ if (confirmedReads !== undefined) { databaseUrl.searchParams.set('confirmed', confirmedReads.toString()); } + // Note that `url`'s own query parameters are not carried over by the `URL` + // constructor above, so these must be set here. + if (connectionId) { + databaseUrl.searchParams.set('connection_id', connectionId); + } + if (sessionId) { + databaseUrl.searchParams.set('session_id', sessionId); + } const ws = new WS(databaseUrl.toString(), wsProtocol); ws.binaryType = 'arraybuffer'; diff --git a/crates/bindings-typescript/tests/connection_manager_liveness.test.ts b/crates/bindings-typescript/tests/connection_manager_liveness.test.ts deleted file mode 100644 index 5c2e1e75c98..00000000000 --- a/crates/bindings-typescript/tests/connection_manager_liveness.test.ts +++ /dev/null @@ -1,319 +0,0 @@ -import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; -import { ConnectionId } from '../src'; -import { connectionManagerReconnectDelayMs } from '../src/sdk/connection_manager.ts'; - -// These tests exercise the page-resume + zombie-socket liveness recovery in the -// ConnectionManager. That logic wires itself to `document`/`window` events in -// its constructor, so each test installs minimal DOM stubs and re-imports the -// module to get a fresh singleton bound to those stubs. - -type ErrorContextInterface = { isActive: boolean }; - -class MockConnection { - isActive = false; - identity = undefined; - // A real DbConnectionImpl is constructed with the builder's token and keeps - // it in this field, so the mock takes it the same way. - token: string | undefined; - connectionId = ConnectionId.random(); - isDisconnectRequested = false; - disconnected = false; - // Controls the `isSocketClosed` signal the manager reads to detect a socket - // that died silently (CLOSING/CLOSED without a clean `onclose`). - socketClosed = false; - - #onConnect = new Set<(conn: MockConnection) => void>(); - #onDisconnect = new Set< - (ctx: ErrorContextInterface, error?: Error) => void - >(); - #onConnectError = new Set< - (ctx: ErrorContextInterface, error: Error) => void - >(); - - get isSocketClosed(): boolean { - return this.socketClosed; - } - - disconnect(): void { - this.isDisconnectRequested = true; - this.disconnected = true; - this.isActive = false; - } - - removeOnConnect(cb: (conn: MockConnection) => void): void { - this.#onConnect.delete(cb); - } - removeOnDisconnect( - cb: (ctx: ErrorContextInterface, error?: Error) => void - ): void { - this.#onDisconnect.delete(cb); - } - removeOnConnectError( - cb: (ctx: ErrorContextInterface, error: Error) => void - ): void { - this.#onConnectError.delete(cb); - } - - register( - type: 'connect' | 'disconnect' | 'connectError', - cb: (...args: any[]) => void - ): void { - if (type === 'connect') this.#onConnect.add(cb); - else if (type === 'disconnect') this.#onDisconnect.add(cb); - else this.#onConnectError.add(cb); - } - - /** - * @param issuedToken - the token the server hands back on connect, emulating - * a client being issued credentials it did not have when the builder was - * constructed. - */ - simulateConnect(issuedToken?: string): void { - this.isActive = true; - if (issuedToken !== undefined) this.token = issuedToken; - for (const cb of this.#onConnect) cb(this); - } - simulateDisconnect(error?: Error): void { - this.isActive = false; - for (const cb of this.#onDisconnect) - cb(this as unknown as ErrorContextInterface, error); - } -} - -class MockBuilder { - buildCount = 0; - connections: MockConnection[] = []; - /** The token each `build()` will stamp onto its connection. */ - token: string | undefined; - - #onConnect = new Set<(conn: MockConnection) => void>(); - #onDisconnect = new Set< - (ctx: ErrorContextInterface, error?: Error) => void - >(); - #onConnectError = new Set< - (ctx: ErrorContextInterface, error: Error) => void - >(); - - withToken(token?: string): MockBuilder { - this.token = token; - return this; - } - - build(): MockConnection { - const connection = new MockConnection(); - connection.token = this.token; - this.buildCount += 1; - this.connections.push(connection); - for (const cb of this.#onConnect) connection.register('connect', cb); - for (const cb of this.#onDisconnect) connection.register('disconnect', cb); - for (const cb of this.#onConnectError) - connection.register('connectError', cb); - return connection; - } - - onConnect(cb: (conn: MockConnection) => void): MockBuilder { - this.#onConnect.add(cb); - for (const c of this.connections) c.register('connect', cb); - return this; - } - onDisconnect( - cb: (ctx: ErrorContextInterface, error?: Error) => void - ): MockBuilder { - this.#onDisconnect.add(cb); - for (const c of this.connections) c.register('disconnect', cb); - return this; - } - onConnectError( - cb: (ctx: ErrorContextInterface, error: Error) => void - ): MockBuilder { - this.#onConnectError.add(cb); - for (const c of this.connections) c.register('connectError', cb); - return this; - } -} - -let keyCounter = 0; -function nextKey(): string { - keyCounter += 1; - return `connection-manager-liveness-${keyCounter}`; -} - -type DocStub = { - visibilityState: 'visible' | 'hidden'; - addEventListener: (ev: string, h: () => void) => void; -}; - -let ConnectionManager: typeof import('../src/sdk/connection_manager.ts').ConnectionManager; -let doc: DocStub; -let listeners: Record void>>; - -function retain(key: string, builder: MockBuilder): MockConnection { - return ConnectionManager.retain( - key, - builder as any - ) as unknown as MockConnection; -} - -function fire(name: string): void { - for (const h of listeners[name] ?? []) h(); -} - -async function loadManager(): Promise { - listeners = {}; - doc = { - visibilityState: 'visible', - addEventListener: (ev, h) => { - (listeners[`doc:${ev}`] ??= []).push(h); - }, - }; - const win = { - addEventListener: (ev: string, h: () => void) => { - (listeners[`win:${ev}`] ??= []).push(h); - }, - }; - (globalThis as any).document = doc; - (globalThis as any).window = win; - vi.resetModules(); - ({ ConnectionManager } = await import('../src/sdk/connection_manager.ts')); -} - -describe('ConnectionManager liveness recovery', () => { - beforeEach(async () => { - // Fake timers let us drive the reconnect backoff deterministically. - vi.useFakeTimers(); - await loadManager(); - }); - - afterEach(() => { - vi.runOnlyPendingTimers(); - vi.useRealTimers(); - delete (globalThis as any).document; - delete (globalThis as any).window; - }); - - test('registers resume + network listeners on construction', () => { - expect(listeners['doc:visibilitychange']?.length).toBe(1); - expect(listeners['win:focus']?.length).toBe(1); - expect(listeners['win:online']?.length).toBe(1); - expect(listeners['win:pageshow']?.length).toBe(1); - }); - - test('revives a silently-dead socket when the network returns', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateConnect(); - expect(ConnectionManager.getSnapshot(key)?.isActive).toBe(true); - - // Socket dies while backgrounded: no disconnect event ever fires, but the - // underlying readyState is now CLOSED. - first.socketClosed = true; - expect(ConnectionManager.getConnection(key)).toBe(first); - - fire('win:online'); - - expect(builder.buildCount).toBe(2); - expect(ConnectionManager.getConnection(key)).toBe(builder.connections[1]); - ConnectionManager.release(key); - }); - - test('does not rebuild a healthy connection on resume', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateConnect(); - - fire('win:focus'); - fire('doc:visibilitychange'); - - expect(builder.buildCount).toBe(1); - expect(ConnectionManager.getConnection(key)).toBe(first); - ConnectionManager.release(key); - }); - - test('does not revive a connection that was intentionally disconnected', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateConnect(); - first.isDisconnectRequested = true; - first.socketClosed = true; - - fire('win:online'); - - expect(builder.buildCount).toBe(1); - ConnectionManager.release(key); - }); - - test('brings a stalled reconnect forward on resume and resets backoff', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateDisconnect(); - - // The reconnect timer is scheduled but has not fired yet (simulating a - // background tab whose timers are throttled/frozen). - vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0) - 1); - expect(builder.buildCount).toBe(1); - - // Regaining focus rebuilds immediately instead of waiting out the delay. - fire('doc:visibilitychange'); - expect(builder.buildCount).toBe(2); - - // Backoff was reset: the next failure reconnects after the base delay. - builder.connections[1].simulateDisconnect(); - vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); - expect(builder.buildCount).toBe(3); - - ConnectionManager.release(key); - }); - - test('visibilitychange while still hidden does not reconnect', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateDisconnect(); - - doc.visibilityState = 'hidden'; - fire('doc:visibilitychange'); - - expect(builder.buildCount).toBe(1); - ConnectionManager.release(key); - }); - - // The resume paths rebuild from the retained builder, whose token is a - // snapshot from before the session existed. Reconnecting anonymously here - // makes the server mint a new identity, so a user who merely switched tabs - // comes back as a stranger. - test('reviving a dead socket on resume keeps the session identity', () => { - const key = nextKey(); - // A first-time visitor: no stored credentials when the builder was made. - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateConnect('session-token'); - - // Tab is backgrounded and the socket dies silently. - first.socketClosed = true; - fire('doc:visibilitychange'); - - expect(builder.buildCount).toBe(2); - expect(builder.connections[1].token).toBe('session-token'); - ConnectionManager.release(key); - }); - - test('a stalled reconnect brought forward on resume keeps the session identity', () => { - const key = nextKey(); - const builder = new MockBuilder(); - const first = retain(key, builder); - first.simulateConnect('session-token'); - first.simulateDisconnect(); - - // Timer still pending behind background throttling; focus fires it early. - vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0) - 1); - fire('win:focus'); - - expect(builder.buildCount).toBe(2); - expect(builder.connections[1].token).toBe('session-token'); - ConnectionManager.release(key); - }); -}); diff --git a/crates/bindings-typescript/tests/connection_manager_reconnect.test.ts b/crates/bindings-typescript/tests/connection_manager_reconnect.test.ts index ee981ccf2f9..2d57d1ef657 100644 --- a/crates/bindings-typescript/tests/connection_manager_reconnect.test.ts +++ b/crates/bindings-typescript/tests/connection_manager_reconnect.test.ts @@ -6,6 +6,13 @@ import { ConnectionManager, } from '../src/sdk/connection_manager.ts'; +// Reconnection after a mid-session drop lives in the SDK: the manager forces +// `withAutomaticReconnect()` on every builder, and a `nextReconnectAttempt` +// on a disconnect/connect-error report means the SDK is retrying inside the +// same connection object, so the manager must leave it alone. The manager +// rebuilds only when the SDK reports it will not retry (no attempt number): +// a failed initial connection, or another terminal failure. + type ErrorContextInterface = { isActive: boolean; }; @@ -26,10 +33,20 @@ class MockConnection { #onConnectCallbacks = new Set<(conn: MockConnection) => void>(); #onDisconnectCallbacks = new Set< - (ctx: ErrorContextInterface, error?: Error) => void + ( + ctx: ErrorContextInterface, + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void >(); #onConnectErrorCallbacks = new Set< - (ctx: ErrorContextInterface, error: Error) => void + ( + ctx: ErrorContextInterface, + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void >(); disconnect(): void { @@ -40,7 +57,7 @@ class MockConnection { this.disconnected = true; this.isActive = false; for (const cb of this.#onDisconnectCallbacks) { - cb(this as unknown as ErrorContextInterface); + cb(this); } } @@ -87,17 +104,29 @@ class MockConnection { } } - simulateDisconnect(error?: Error): void { + /** + * Passing `nextReconnectAttempt` emulates the SDK announcing it will retry + * internally; omitting it emulates a terminal report. + */ + simulateDisconnect( + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ): void { this.isActive = false; for (const cb of this.#onDisconnectCallbacks) { - cb(this as unknown as ErrorContextInterface, error); + cb(this, error, nextReconnectAttempt, nextReconnectDelayMs); } } - simulateConnectError(error: Error): void { + simulateConnectError( + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ): void { this.isActive = false; for (const cb of this.#onConnectErrorCallbacks) { - cb(this as unknown as ErrorContextInterface, error); + cb(this, error, nextReconnectAttempt, nextReconnectDelayMs); } } @@ -106,13 +135,23 @@ class MockConnection { } registerOnDisconnect( - cb: (ctx: ErrorContextInterface, error?: Error) => void + cb: ( + ctx: ErrorContextInterface, + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void ): void { this.#onDisconnectCallbacks.add(cb); } registerOnConnectError( - cb: (ctx: ErrorContextInterface, error: Error) => void + cb: ( + ctx: ErrorContextInterface, + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void ): void { this.#onConnectErrorCallbacks.add(cb); } @@ -125,6 +164,7 @@ class MockBuilder { token: string | undefined; /** Every token this builder was asked to carry, oldest first. */ tokenHistory: (string | undefined)[] = []; + automaticReconnect = false; constructor(token?: string) { this.token = token; @@ -144,6 +184,11 @@ class MockBuilder { return this; } + withAutomaticReconnect(): MockBuilder { + this.automaticReconnect = true; + return this; + } + build(): MockConnection { const connection = new MockConnection(this.token); this.buildCount += 1; @@ -205,7 +250,7 @@ function retainMock(key: string, builder: MockBuilder): MockConnection { ) as unknown as MockConnection; } -describe('ConnectionManager retained reconnect behavior', () => { +describe('ConnectionManager forces SDK automatic reconnection', () => { beforeEach(() => { vi.useFakeTimers(); }); @@ -215,7 +260,134 @@ describe('ConnectionManager retained reconnect behavior', () => { vi.useRealTimers(); }); - test('rebuilds a retained connection after disconnect', () => { + test('retain enables automatic reconnection on the builder', () => { + const key = nextKey(); + const builder = new MockBuilder(); + expect(builder.automaticReconnect).toBe(false); + + retainMock(key, builder); + + expect(builder.automaticReconnect).toBe(true); + + ConnectionManager.release(key); + }); + + test('rebuild enables automatic reconnection on the replacement builder', () => { + const key = nextKey(); + retainMock(key, new MockBuilder()); + + const replacement = new MockBuilder(); + ConnectionManager.rebuild(key, replacement as any); + + expect(replacement.automaticReconnect).toBe(true); + + ConnectionManager.release(key); + }); +}); + +describe('ConnectionManager during SDK-managed reconnection', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + test('a drop the SDK will retry does not rebuild the connection', () => { + const key = nextKey(); + const builder = new MockBuilder(); + const error = new Error('connection lost'); + + const first = retainMock(key, builder); + first.simulateConnect(); + + first.simulateDisconnect(error, 1, 1000); + + // The connection object stays managed and untouched: the SDK reconnects + // inside it. + expect(ConnectionManager.getConnection(key)).toBe(first); + expect(ConnectionManager.getSnapshot(key)?.isActive).toBe(false); + expect(ConnectionManager.getSnapshot(key)?.connectionError).toBe(error); + + vi.advanceTimersByTime(CONNECTION_MANAGER_RECONNECT_MAX_DELAY_MS); + expect(builder.buildCount).toBe(1); + + ConnectionManager.release(key); + }); + + test('failed attempts the SDK will retry do not rebuild the connection', () => { + const key = nextKey(); + const builder = new MockBuilder(); + + const first = retainMock(key, builder); + first.simulateConnect(); + first.simulateDisconnect(new Error('connection lost'), 1, 1000); + + const attemptError = new Error('still down'); + first.simulateConnectError(attemptError, 2, 2000); + + expect(ConnectionManager.getConnection(key)).toBe(first); + expect(ConnectionManager.getSnapshot(key)?.connectionError).toBe( + attemptError + ); + + vi.advanceTimersByTime(CONNECTION_MANAGER_RECONNECT_MAX_DELAY_MS); + expect(builder.buildCount).toBe(1); + + ConnectionManager.release(key); + }); + + test('manager callbacks stay attached while the SDK retries', () => { + const key = nextKey(); + const builder = new MockBuilder(); + + const first = retainMock(key, builder); + first.simulateConnect(); + first.simulateDisconnect(new Error('connection lost'), 1, 1000); + + expect(first.callbackCounts()).toEqual({ + connect: 1, + disconnect: 1, + connectError: 1, + }); + + ConnectionManager.release(key); + }); + + test('a successful SDK reconnect restores the state snapshot', () => { + const key = nextKey(); + const builder = new MockBuilder(); + + const first = retainMock(key, builder); + first.simulateConnect('session-token'); + first.simulateDisconnect(new Error('connection lost'), 1, 1000); + expect(ConnectionManager.getSnapshot(key)?.isActive).toBe(false); + + // The SDK reconnects inside the same object and fires onConnect again. + first.simulateConnect('session-token'); + + expect(ConnectionManager.getSnapshot(key)?.isActive).toBe(true); + expect(ConnectionManager.getSnapshot(key)?.connectionError).toBeUndefined(); + expect(ConnectionManager.getConnection(key)).toBe(first); + expect(builder.buildCount).toBe(1); + + ConnectionManager.release(key); + }); +}); + +describe('ConnectionManager rebuild on terminal failures', () => { + beforeEach(() => { + vi.useFakeTimers(); + }); + + afterEach(() => { + vi.runOnlyPendingTimers(); + vi.useRealTimers(); + }); + + test('rebuilds a retained connection after a terminal disconnect', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -242,9 +414,11 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('rebuilds a retained connection after connectError', () => { + test('rebuilds a retained connection after a terminal connectError', () => { const key = nextKey(); const builder = new MockBuilder(); + // A failed *initial* connection is the SDK's main terminal case: it + // reports it through onConnectError with no next attempt. const error = new Error('network unavailable'); const first = retainMock(key, builder); @@ -263,7 +437,7 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('same-key retain after disconnect returns a fresh connection immediately', () => { + test('same-key retain after a terminal failure returns a fresh connection immediately', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -283,7 +457,7 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('reconnect uses callbacks from a replacement same-key builder', () => { + test('rebuild uses callbacks from a replacement same-key builder', () => { const key = nextKey(); const firstBuilder = new MockBuilder(); const secondBuilder = new MockBuilder(); @@ -315,7 +489,7 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('disconnect removes manager callbacks from the old connection before pending reconnect', () => { + test('a terminal failure removes manager callbacks from the old connection', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -337,7 +511,7 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('release cancels a pending reconnect', () => { + test('release cancels a pending rebuild', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -351,7 +525,7 @@ describe('ConnectionManager retained reconnect behavior', () => { expect(ConnectionManager.getConnection(key)).toBeNull(); }); - test('manual disconnect does not trigger a reconnect', () => { + test('manual disconnect does not trigger a rebuild', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -387,14 +561,14 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('reconnect delay backs off exponentially across consecutive failures', () => { + test('rebuild delay backs off exponentially across consecutive failures', () => { const key = nextKey(); const builder = new MockBuilder(); const first = retainMock(key, builder); - first.simulateDisconnect(); + first.simulateConnectError(new Error('server unreachable')); - // First reconnect fires after the base delay. + // First rebuild fires after the base delay. vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); expect(builder.buildCount).toBe(2); @@ -415,12 +589,12 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('successful connect resets the reconnect backoff', () => { + test('successful connect resets the rebuild backoff', () => { const key = nextKey(); const builder = new MockBuilder(); const first = retainMock(key, builder); - first.simulateDisconnect(); + first.simulateConnectError(new Error('server unreachable')); vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); builder.connections[1].simulateConnectError(new Error('still down')); @@ -437,7 +611,7 @@ describe('ConnectionManager retained reconnect behavior', () => { ConnectionManager.release(key); }); - test('reconnect delay is capped at the maximum delay', () => { + test('rebuild delay is capped at the maximum delay', () => { expect(connectionManagerReconnectDelayMs(0)).toBeLessThan( CONNECTION_MANAGER_RECONNECT_MAX_DELAY_MS ); @@ -532,18 +706,19 @@ describe('ConnectionManager.rebuild', () => { ConnectionManager.release(key); }); - test('cancels a pending auto-reconnect and resets the backoff', () => { + test('cancels a pending terminal-failure rebuild and resets the backoff', () => { const key = nextKey(); const builder = new MockBuilder(); const first = retainMock(key, builder); - // Two consecutive failures so the backoff has advanced past the base delay. - first.simulateDisconnect(); + // Two consecutive terminal failures so the backoff has advanced past the + // base delay. + first.simulateConnectError(new Error('server unreachable')); vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); builder.connections[1].simulateConnectError(new Error('still down')); expect(builder.buildCount).toBe(2); - // rebuild() takes over: the scheduled reconnect must not also fire. + // rebuild() takes over: the scheduled rebuild must not also fire. const replacement = new MockBuilder(); ConnectionManager.rebuild(key, replacement as any); expect(replacement.buildCount).toBe(1); @@ -553,7 +728,8 @@ describe('ConnectionManager.rebuild', () => { expect(builder.buildCount).toBe(2); expect(replacement.buildCount).toBe(1); - // ...and the backoff was reset: a fresh drop reconnects after the base delay. + // ...and the backoff was reset: a fresh terminal failure rebuilds after + // the base delay. replacement.connections[0].simulateConnect(); replacement.connections[0].simulateDisconnect(); vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); @@ -576,6 +752,9 @@ describe('ConnectionManager.rebuild', () => { build() { throw buildError; }, + withAutomaticReconnect() { + return this; + }, onConnect() { return this; }, @@ -629,7 +808,7 @@ describe('ConnectionManager session continuity across rebuilds', () => { vi.useRealTimers(); }); - test('auto-reconnect reuses the token issued after the builder was built', () => { + test('a terminal-failure rebuild reuses the token issued after the builder was built', () => { const key = nextKey(); // A first-time visitor: nothing in storage, so the builder carries no token. const builder = new MockBuilder(); @@ -642,8 +821,9 @@ describe('ConnectionManager session continuity across rebuilds', () => { first.simulateConnect('session-token'); expect(ConnectionManager.getSnapshot(key)?.token).toBe('session-token'); - // The socket drops and the manager auto-reconnects from the retained - // builder — which still holds the empty token it was constructed with. + // The connection ends terminally and the manager rebuilds from the + // retained builder — which still holds the empty token it was constructed + // with. first.simulateDisconnect(); vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); @@ -654,7 +834,7 @@ describe('ConnectionManager session continuity across rebuilds', () => { expect(second.token).toBe('session-token'); }); - test('resumed session survives repeated reconnects', () => { + test('resumed session survives repeated rebuilds', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -708,7 +888,7 @@ describe('ConnectionManager session continuity across rebuilds', () => { ConnectionManager.release(key); }); - test('retain after a drop resumes the session rather than the stale builder', () => { + test('retain after a terminal failure resumes the session rather than the stale builder', () => { const key = nextKey(); const builder = new MockBuilder(); @@ -716,7 +896,7 @@ describe('ConnectionManager session continuity across rebuilds', () => { first.simulateConnect('session-token'); first.simulateDisconnect(); - // A provider remount rebuilds through retain(), not the reconnect timer. + // A provider remount rebuilds through retain(), not the rebuild timer. const second = retainMock(key, builder); expect(second.token).toBe('session-token'); @@ -745,15 +925,16 @@ describe('ConnectionManager session continuity across rebuilds', () => { ConnectionManager.release(key); }); - test('auto-reconnect with a replacement builder keeps the session identity', () => { + test('a terminal-failure rebuild with a replacement builder keeps the session identity', () => { const key = nextKey(); const anonymous = new MockBuilder(); const first = retainMock(key, anonymous); first.simulateConnect('anonymous-token'); - // Swap the builder while the connection is live, then drop: the reconnect - // uses the replacement's callbacks but must not adopt its token. + // Swap the builder while the connection is live, then end it terminally: + // the rebuild uses the replacement's callbacks but must not adopt its + // token. ConnectionManager.release(key); const signedIn = new MockBuilder('signed-in-token'); retainMock(key, signedIn); @@ -787,7 +968,7 @@ describe('ConnectionManager session continuity across rebuilds', () => { ConnectionManager.release(key); }); - test('auto-reconnect after rebuild() keeps the new identity', () => { + test('a terminal-failure rebuild after rebuild() keeps the new identity', () => { const key = nextKey(); const anonymous = new MockBuilder(); @@ -799,8 +980,8 @@ describe('ConnectionManager session continuity across rebuilds', () => { signedIn as any ) as unknown as MockConnection; - // Drop *before* the new connection completes its handshake: the manager - // must not fall back to the identity rebuild() just replaced. + // Fail terminally *before* the new connection completes its handshake: + // the manager must not fall back to the identity rebuild() just replaced. second.simulateDisconnect(); vi.advanceTimersByTime(connectionManagerReconnectDelayMs(0)); diff --git a/crates/bindings-typescript/tests/db_connection_liveness.test.ts b/crates/bindings-typescript/tests/db_connection_liveness.test.ts new file mode 100644 index 00000000000..41e36479f85 --- /dev/null +++ b/crates/bindings-typescript/tests/db_connection_liveness.test.ts @@ -0,0 +1,267 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; +import { Identity } from '../src'; +import { ServerMessage } from '../src/sdk/client_api/types'; +import { WebsocketTestAdapterFactory } from '../src/sdk/websocket_test_adapter'; +import { ConnectionId } from '../src'; +import { DbConnection } from '../test-app/src/module_bindings'; +import { anIdentity } from './utils'; + +// These tests exercise the page-resume liveness recovery in DbConnectionImpl: +// with automatic reconnection enabled, the connection listens for the page +// coming back to the foreground (visibilitychange/focus/online/pageshow) and +// uses the moment to notice sockets that died silently while the tab was +// frozen, and to bring a backoff-stalled reconnect forward. +// +// The listeners bind to `document`/`window` when the socket opens, so each +// test installs minimal DOM stubs first. + +type ReconnectReport = { + error?: Error; + nextReconnectAttempt?: number; + nextReconnectDelayMs?: number; +}; + +type Harness = { + connection: DbConnection; + factory: WebsocketTestAdapterFactory; + connects: { identity: Identity; token: string }[]; + disconnects: ReconnectReport[]; + connectErrors: ReconnectReport[]; +}; + +let listeners: Record void>>; +let visibilityState: 'visible' | 'hidden'; + +function installDomStubs(): void { + listeners = {}; + visibilityState = 'visible'; + const record = + (scope: string) => + (ev: string, h: () => void): void => { + (listeners[`${scope}:${ev}`] ??= []).push(h); + }; + const remove = + (scope: string) => + (ev: string, h: () => void): void => { + const bucket = listeners[`${scope}:${ev}`]; + if (bucket) { + const at = bucket.indexOf(h); + if (at >= 0) bucket.splice(at, 1); + } + }; + vi.stubGlobal('document', { + get visibilityState() { + return visibilityState; + }, + addEventListener: record('doc'), + removeEventListener: remove('doc'), + }); + vi.stubGlobal('window', { + addEventListener: record('win'), + removeEventListener: remove('win'), + }); +} + +function removeDomStubs(): void { + vi.unstubAllGlobals(); +} + +function fire(name: string): void { + for (const h of [...(listeners[name] ?? [])]) h(); +} + +function listenerCounts(): Record { + return Object.fromEntries( + Object.entries(listeners).map(([name, hs]) => [name, hs.length]) + ); +} + +function build(options?: { automaticReconnect?: boolean }): Harness { + const factory = new WebsocketTestAdapterFactory(); + const connects: { identity: Identity; token: string }[] = []; + const disconnects: ReconnectReport[] = []; + const connectErrors: ReconnectReport[] = []; + + let builder = DbConnection.builder() + .withUri('ws://127.0.0.1:1234') + .withDatabaseName('db') + .withWSFn(factory.openWebSocket) + .onConnect((_conn, identity, token) => connects.push({ identity, token })) + .onDisconnect((_ctx, error, nextReconnectAttempt, nextReconnectDelayMs) => + disconnects.push({ error, nextReconnectAttempt, nextReconnectDelayMs }) + ) + .onConnectError((_ctx, error, nextReconnectAttempt, nextReconnectDelayMs) => + connectErrors.push({ error, nextReconnectAttempt, nextReconnectDelayMs }) + ); + if (options?.automaticReconnect ?? true) { + builder = builder.withAutomaticReconnect(); + } + + return { + connection: builder.build(), + factory, + connects, + disconnects, + connectErrors, + }; +} + +/** Let the connection's pending socket promise settle. */ +async function settle(harness: Harness): Promise { + await harness.connection['wsPromise']; + await Promise.resolve(); +} + +/** Bring a connection up to an established state on its current socket. */ +async function establish(harness: Harness): Promise { + await settle(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient( + ServerMessage.InitialConnection({ + identity: anIdentity, + connectionId: ConnectionId.random(), + token: 'issued-token', + }) + ); + await Promise.resolve(); +} + +beforeEach(() => { + vi.useFakeTimers(); + installDomStubs(); +}); + +afterEach(() => { + vi.useRealTimers(); + removeDomStubs(); +}); + +describe('liveness listeners', () => { + test('are installed when the socket opens with automatic reconnect enabled', async () => { + const harness = build(); + await establish(harness); + + expect(listenerCounts()).toEqual({ + 'doc:visibilitychange': 1, + 'win:focus': 1, + 'win:online': 1, + 'win:pageshow': 1, + }); + }); + + test('are not installed without automatic reconnect', async () => { + const harness = build({ automaticReconnect: false }); + await establish(harness); + + expect(listeners).toEqual({}); + }); + + test('are removed when the connection ends', async () => { + const harness = build(); + await establish(harness); + + harness.connection.disconnect(); + await settle(harness); + + expect(listenerCounts()).toEqual({ + 'doc:visibilitychange': 0, + 'win:focus': 0, + 'win:online': 0, + 'win:pageshow': 0, + }); + }); +}); + +describe('liveness recovery on page resume', () => { + test('treats a silently-dead socket as a lost connection when the network returns', async () => { + const harness = build(); + await establish(harness); + const firstSocket = harness.factory.current; + + // Socket dies while backgrounded: no close event is ever delivered, but + // the underlying readyState is now CLOSED. + firstSocket.dieSilently(); + expect(harness.disconnects).toHaveLength(0); + + fire('win:online'); + + // The loss is reported like any mid-session drop, announcing a retry... + expect(harness.disconnects).toHaveLength(1); + expect(harness.disconnects[0].nextReconnectAttempt).toBe(1); + + // ...and the scheduled attempt builds a fresh socket the connection can + // re-establish on. + await vi.runOnlyPendingTimersAsync(); + await settle(harness); + expect(harness.factory.current).not.toBe(firstSocket); + await establish(harness); + expect(harness.connects).toHaveLength(2); + expect(harness.connection.isActive).toBe(true); + }); + + test('does not disturb a healthy connection on resume', async () => { + const harness = build(); + await establish(harness); + const firstSocket = harness.factory.current; + + fire('win:focus'); + fire('doc:visibilitychange'); + await settle(harness); + + expect(harness.disconnects).toHaveLength(0); + expect(harness.factory.current).toBe(firstSocket); + expect(harness.connection.isActive).toBe(true); + }); + + test('does not revive a connection after an explicit disconnect', async () => { + const harness = build(); + await establish(harness); + + harness.connection.disconnect(); + await settle(harness); + const disconnectsBefore = harness.disconnects.length; + + harness.factory.current.dieSilently(); + fire('win:online'); + await vi.runOnlyPendingTimersAsync(); + + expect(harness.disconnects).toHaveLength(disconnectsBefore); + expect(harness.connection.isActive).toBe(false); + }); + + test('brings a backoff-stalled reconnect forward on resume', async () => { + const harness = build(); + await establish(harness); + const firstSocket = harness.factory.current; + + // Drop the connection; a reconnect is now waiting out its backoff delay + // (simulating a background tab whose timers are throttled/frozen). + firstSocket.serverClose(1006); + expect(harness.disconnects).toHaveLength(1); + const delay = harness.disconnects[0].nextReconnectDelayMs!; + vi.advanceTimersByTime(delay - 1); + expect(harness.factory.current).toBe(firstSocket); + + // Regaining visibility retries immediately instead of waiting out the + // remaining delay. + fire('doc:visibilitychange'); + await settle(harness); + expect(harness.factory.current).not.toBe(firstSocket); + + await establish(harness); + expect(harness.connects).toHaveLength(2); + }); + + test('visibilitychange while still hidden does nothing', async () => { + const harness = build(); + await establish(harness); + const firstSocket = harness.factory.current; + + firstSocket.serverClose(1006); + visibilityState = 'hidden'; + fire('doc:visibilitychange'); + await settle(harness); + + expect(harness.factory.current).toBe(firstSocket); + }); +}); diff --git a/crates/bindings-typescript/tests/db_connection_reconnect.test.ts b/crates/bindings-typescript/tests/db_connection_reconnect.test.ts new file mode 100644 index 00000000000..f255a9098f5 --- /dev/null +++ b/crates/bindings-typescript/tests/db_connection_reconnect.test.ts @@ -0,0 +1,1201 @@ +import { afterEach, beforeEach, describe, expect, test, vi } from 'vitest'; +import { ConnectionId, Identity } from '../src'; +import { + DisconnectedError, + IdentityChangedError, + UnknownCallResultError, +} from '../src/lib/errors'; +import { + computeReconnectDelayMs, + tokenNeedsRefresh, + RECONNECT_INITIAL_DELAY_MS, + RECONNECT_MAX_DELAY_MS, +} from '../src/sdk/db_connection_impl'; +import { + ServerMessage, + type SubscribeBatch, +} from '../src/sdk/client_api/types'; +import { WebSocketTokenError } from '../src/sdk/ws'; +import { WebsocketTestAdapterFactory } from '../src/sdk/websocket_test_adapter'; +import { DbConnection } from '../test-app/src/module_bindings'; +import { anIdentity, bobIdentity, encodeUser } from './utils'; + +/** The disconnect/connect-error reports an application sees. */ +type ReconnectReport = { + error?: Error; + nextReconnectAttempt?: number; + nextReconnectDelayMs?: number; +}; + +type Harness = { + connection: DbConnection; + factory: WebsocketTestAdapterFactory; + connects: { identity: Identity; token: string }[]; + disconnects: ReconnectReport[]; + connectErrors: ReconnectReport[]; +}; + +const TOKEN = 'issued-token'; + +function build(options?: { + automaticReconnect?: boolean; + token?: string; + tokenProvider?: () => Promise; +}): Harness { + const factory = new WebsocketTestAdapterFactory(); + const connects: { identity: Identity; token: string }[] = []; + const disconnects: ReconnectReport[] = []; + const connectErrors: ReconnectReport[] = []; + + let builder = DbConnection.builder() + .withUri('ws://127.0.0.1:1234') + .withDatabaseName('db') + .withWSFn(factory.openWebSocket) + .onConnect((_conn, identity, token) => connects.push({ identity, token })) + .onDisconnect((_ctx, error, nextReconnectAttempt, nextReconnectDelayMs) => + disconnects.push({ error, nextReconnectAttempt, nextReconnectDelayMs }) + ) + .onConnectError((_ctx, error, nextReconnectAttempt, nextReconnectDelayMs) => + connectErrors.push({ error, nextReconnectAttempt, nextReconnectDelayMs }) + ); + if (options?.token) { + builder = builder.withToken(options.token); + } + if (options?.automaticReconnect ?? true) { + builder = builder.withAutomaticReconnect(); + } + if (options?.tokenProvider) { + builder = builder.withTokenProvider(options.tokenProvider); + } + + return { + connection: builder.build(), + factory, + connects, + disconnects, + connectErrors, + }; +} + +/** Let the connection's pending socket promise settle. */ +async function settle(harness: Harness): Promise { + await harness.connection['wsPromise']; + await Promise.resolve(); +} + +function initialConnection( + identity: Identity = anIdentity, + connectionId: ConnectionId = ConnectionId.random() +): ServerMessage { + return ServerMessage.InitialConnection({ + identity, + connectionId, + token: TOKEN, + }); +} + +/** Bring a connection up to an established state on its current socket. */ +async function establish( + harness: Harness, + identity: Identity = anIdentity +): Promise { + await settle(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection(identity)); + await Promise.resolve(); +} + +/** Run the scheduled reconnect timer and let its socket be created. */ +async function runReconnectTimer(harness: Harness): Promise { + await vi.runOnlyPendingTimersAsync(); + await settle(harness); +} + +beforeEach(() => { + vi.useFakeTimers(); +}); + +afterEach(() => { + vi.useRealTimers(); +}); + +describe('reconnect policy', () => { + test('delays grow exponentially from the initial delay', () => { + const noJitter = () => 0.5; + expect(computeReconnectDelayMs(1, noJitter)).toBe( + RECONNECT_INITIAL_DELAY_MS + ); + expect(computeReconnectDelayMs(2, noJitter)).toBe( + RECONNECT_INITIAL_DELAY_MS * 2 + ); + expect(computeReconnectDelayMs(3, noJitter)).toBe( + RECONNECT_INITIAL_DELAY_MS * 4 + ); + }); + + test('delays are capped', () => { + const noJitter = () => 0.5; + expect(computeReconnectDelayMs(20, noJitter)).toBe(RECONNECT_MAX_DELAY_MS); + }); + + test('jitter spreads the delay around the base but never exceeds the cap', () => { + expect(computeReconnectDelayMs(3, () => 0)).toBeLessThan( + computeReconnectDelayMs(3, () => 1) + ); + expect(computeReconnectDelayMs(30, () => 1)).toBeLessThanOrEqual( + RECONNECT_MAX_DELAY_MS + ); + expect(computeReconnectDelayMs(1, () => 0)).toBeGreaterThanOrEqual(0); + }); +}); + +describe('token refresh', () => { + const nowSeconds = 1_000_000; + const nowMs = nowSeconds * 1000; + + function jwt(claims: object): string { + const payload = Buffer.from(JSON.stringify(claims)).toString('base64url'); + return `header.${payload}.signature`; + } + + test('a token with plenty of life left is not refreshed', () => { + const token = jwt({ iat: nowSeconds - 60, exp: nowSeconds + 3600 }); + expect(tokenNeedsRefresh(token, nowMs)).toBe(false); + }); + + test('a token close to expiring is refreshed', () => { + const token = jwt({ iat: nowSeconds - 3590, exp: nowSeconds + 10 }); + expect(tokenNeedsRefresh(token, nowMs)).toBe(true); + }); + + test('a short-lived token is refreshed on the 30 second floor', () => { + // 5% of a 60 second lifetime is only 3 seconds, so the floor applies. + const token = jwt({ iat: nowSeconds - 40, exp: nowSeconds + 20 }); + expect(tokenNeedsRefresh(token, nowMs)).toBe(true); + }); + + test('an expired token is refreshed', () => { + const token = jwt({ iat: nowSeconds - 3600, exp: nowSeconds - 1 }); + expect(tokenNeedsRefresh(token, nowMs)).toBe(true); + }); + + test('a token whose expiry cannot be read is always refreshed', () => { + expect(tokenNeedsRefresh('not-a-jwt', nowMs)).toBe(true); + expect(tokenNeedsRefresh(jwt({ sub: 'no-exp' }), nowMs)).toBe(true); + expect(tokenNeedsRefresh(undefined, nowMs)).toBe(true); + }); +}); + +describe('losing an established connection', () => { + test('onDisconnect announces the first reconnect attempt', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await Promise.resolve(); + + expect(harness.disconnects).toHaveLength(1); + expect(harness.disconnects[0].nextReconnectAttempt).toBe(1); + expect(harness.disconnects[0].nextReconnectDelayMs).toBeGreaterThan(0); + expect(harness.disconnects[0].error).toBeInstanceOf(Error); + }); + + test('a reconnect opens a new socket and fires onConnect again', async () => { + const harness = build(); + await establish(harness); + expect(harness.factory.sockets).toHaveLength(1); + + harness.factory.current.close(); + await runReconnectTimer(harness); + expect(harness.factory.sockets).toHaveLength(2); + + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + expect(harness.connects).toHaveLength(2); + expect(harness.connection.isActive).toBe(true); + }); + + test('the connection object and its cache survive a reconnect', async () => { + const harness = build(); + await establish(harness); + const cacheBefore = harness.connection.db; + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + expect(harness.connection.db).toBe(cacheBefore); + }); + + test('the retained token is used to reconnect, keeping the identity stable', async () => { + // Connect anonymously; the server issues a token. + const harness = build(); + await establish(harness); + expect(harness.connects[0].token).toBe(TOKEN); + + harness.factory.current.close(); + await runReconnectTimer(harness); + + expect(harness.factory.current.connectArgs?.authToken).toBe(TOKEN); + }); + + test('a session id is sent so the server can supersede the old connection', async () => { + const harness = build(); + await establish(harness); + const sessionId = harness.factory.sockets[0].connectArgs?.sessionId; + expect(sessionId).toBeTruthy(); + + harness.factory.current.close(); + await runReconnectTimer(harness); + + // The same session id identifies both connections as one client session. + expect(harness.factory.current.connectArgs?.sessionId).toBe(sessionId); + // Each connection still has its own connection id. + expect(harness.factory.current.connectArgs?.connectionId).toBeTruthy(); + }); + + test('reconnection is off unless requested', async () => { + const harness = build({ automaticReconnect: false }); + await establish(harness); + + harness.factory.current.close(); + await vi.runOnlyPendingTimersAsync(); + + expect(harness.disconnects).toHaveLength(1); + expect(harness.disconnects[0].nextReconnectAttempt).toBeUndefined(); + expect(harness.factory.sockets).toHaveLength(1); + }); +}); + +describe('failed reconnect attempts', () => { + test('onConnectError announces the next attempt', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + // The attempt's socket dies before completing its handshake. + harness.factory.current.close(); + await Promise.resolve(); + + expect(harness.connectErrors).toHaveLength(1); + expect(harness.connectErrors[0].nextReconnectAttempt).toBe(2); + // A failed attempt is not a lost connection, so no second onDisconnect. + expect(harness.disconnects).toHaveLength(1); + }); + + test('the attempt number grows across consecutive failures', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + for (let expected = 2; expected <= 4; expected++) { + await runReconnectTimer(harness); + harness.factory.current.close(); + await Promise.resolve(); + expect( + harness.connectErrors[harness.connectErrors.length - 1] + .nextReconnectAttempt + ).toBe(expected); + } + }); + + test('the delay grows across consecutive failures', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + const delays: number[] = [harness.disconnects[0].nextReconnectDelayMs!]; + for (let i = 0; i < 3; i++) { + await runReconnectTimer(harness); + harness.factory.current.close(); + await Promise.resolve(); + delays.push( + harness.connectErrors[harness.connectErrors.length - 1] + .nextReconnectDelayMs! + ); + } + + // Jitter makes individual steps noisy, so compare the ends of the run. + expect(delays[delays.length - 1]).toBeGreaterThan(delays[0]); + }); + + test('a failure to open the socket at all counts as a failed attempt', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + harness.factory.connectError = new Error('server unreachable'); + await runReconnectTimer(harness); + + expect(harness.connectErrors).toHaveLength(1); + expect(harness.connectErrors[0].error?.message).toBe('server unreachable'); + expect(harness.connectErrors[0].nextReconnectAttempt).toBe(2); + }); + + test('the attempt counter resets once a connection is established', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.close(); + await Promise.resolve(); + expect( + harness.connectErrors[harness.connectErrors.length - 1] + .nextReconnectAttempt + ).toBe(2); + + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + // A later drop starts again at attempt 1. + harness.factory.current.close(); + await Promise.resolve(); + expect( + harness.disconnects[harness.disconnects.length - 1].nextReconnectAttempt + ).toBe(1); + }); + + test('an initial connection failure is not retried', async () => { + const harness = build(); + await settle(harness); + + // The socket dies before ever completing a handshake. + harness.factory.current.close(); + await vi.runOnlyPendingTimersAsync(); + + expect(harness.connectErrors).toHaveLength(1); + expect(harness.connectErrors[0].nextReconnectAttempt).toBeUndefined(); + expect(harness.factory.sockets).toHaveLength(1); + }); +}); + +describe('terminal failures', () => { + test('a reconnect under a different identity stops the SDK', async () => { + const harness = build(); + await establish(harness, anIdentity); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + // The server hands us a different identity: the token was replaced. + harness.factory.current.sendToClient(initialConnection(bobIdentity)); + await Promise.resolve(); + + const lastError = + harness.connectErrors[harness.connectErrors.length - 1].error; + expect(lastError).toBeInstanceOf(IdentityChangedError); + expect( + harness.connectErrors[harness.connectErrors.length - 1] + .nextReconnectAttempt + ).toBeUndefined(); + + // No further attempts are scheduled. + const socketsBefore = harness.factory.sockets.length; + await vi.runOnlyPendingTimersAsync(); + expect(harness.factory.sockets).toHaveLength(socketsBefore); + }); +}); + +describe('explicit disconnect()', () => { + test('fires onDisconnect and stops reconnecting', async () => { + const harness = build(); + await establish(harness); + + harness.connection.disconnect(); + harness.factory.current.close(); + await vi.runOnlyPendingTimersAsync(); + + expect(harness.disconnects).toHaveLength(1); + expect(harness.disconnects[0].nextReconnectAttempt).toBeUndefined(); + expect(harness.factory.sockets).toHaveLength(1); + }); + + test('fires onDisconnect when called while reconnecting', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await Promise.resolve(); + expect(harness.disconnects[0].nextReconnectAttempt).toBe(1); + + // Called between attempts, when there is no live socket whose close event + // would otherwise end the connection. + harness.connection.disconnect(); + await vi.runOnlyPendingTimersAsync(); + + expect(harness.disconnects).toHaveLength(2); + expect(harness.disconnects[1].nextReconnectAttempt).toBeUndefined(); + // The scheduled attempt was cancelled. + expect(harness.factory.sockets).toHaveLength(1); + }); + + test('cancels a scheduled attempt even after several failures', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.close(); + await Promise.resolve(); + + harness.connection.disconnect(); + const socketsBefore = harness.factory.sockets.length; + await vi.runOnlyPendingTimersAsync(); + + expect(harness.factory.sockets).toHaveLength(socketsBefore); + expect( + harness.disconnects[harness.disconnects.length - 1].nextReconnectAttempt + ).toBeUndefined(); + }); +}); + +describe('calls while reconnecting', () => { + test('a reducer call fails immediately rather than queueing', async () => { + const harness = build(); + await establish(harness); + harness.factory.current.close(); + await Promise.resolve(); + + await expect( + harness.connection.reducers.createPlayer({ + name: 'Alice', + location: { x: 1, y: 2 }, + }) + ).rejects.toBeInstanceOf(DisconnectedError); + }); + + test('an in-flight call settles with an unknown-result error', async () => { + const harness = build(); + await establish(harness); + + const pending = harness.connection.reducers.createPlayer({ + name: 'Alice', + location: { x: 1, y: 2 }, + }); + // The connection drops before the server acknowledges the call. + harness.factory.current.close(); + await Promise.resolve(); + + await expect(pending).rejects.toBeInstanceOf(UnknownCallResultError); + }); + + test('calls are not rejected without automatic reconnection', async () => { + const harness = build({ automaticReconnect: false }); + await establish(harness); + harness.factory.current.close(); + await Promise.resolve(); + + // Legacy behavior: the call queues on the dead socket rather than failing. + let settled = false; + void harness.connection.reducers + .createPlayer({ name: 'Alice', location: { x: 1, y: 2 } }) + .then( + () => (settled = true), + () => (settled = true) + ); + await Promise.resolve(); + expect(settled).toBe(false); + }); +}); + +describe('replaying subscriptions', () => { + /** The last batch-subscribe message the connection sent, if any. */ + function lastSubscribeBatch(harness: Harness): SubscribeBatch | undefined { + const messages = harness.factory.current.outgoingMessages; + for (let i = messages.length - 1; i >= 0; i--) { + const message = messages[i]; + if (message.tag === 'SubscribeBatch') { + return message.value; + } + } + return undefined; + } + + async function establishWithSubscription(): Promise { + const harness = build(); + await establish(harness); + harness.connection.subscriptionBuilder().subscribe(['SELECT * FROM user']); + await Promise.resolve(); + return harness; + } + + test('a reconnect replays live subscriptions in one batch', async () => { + const harness = await establishWithSubscription(); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + const batch = lastSubscribeBatch(harness); + expect(batch).toBeDefined(); + expect(batch!.sets).toHaveLength(1); + }); + + test('replayed sets are registered under fresh query set ids', async () => { + const harness = await establishWithSubscription(); + const originalSubscribe = harness.factory.current.outgoingMessages.find( + message => message.tag === 'Subscribe' + ); + const originalQuerySetId = originalSubscribe!.value.querySetId.id; + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + const batch = lastSubscribeBatch(harness); + expect(batch!.sets[0].querySetId.id).not.toBe(originalQuerySetId); + }); + + test('nothing is replayed when there are no subscriptions', async () => { + const harness = build(); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + expect(lastSubscribeBatch(harness)).toBeUndefined(); + }); + + test('rows unchanged across the outage produce no callbacks', async () => { + const harness = await establishWithSubscription(); + const querySetId = harness.factory.current.outgoingMessages.find( + message => message.tag === 'Subscribe' + )!.value.querySetId.id; + + // The server delivers one row for the subscription. + harness.factory.current.sendToClient( + ServerMessage.SubscribeApplied({ + requestId: 1, + querySetId: { id: querySetId }, + rows: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alice', + }), + }, + }, + ], + }, + }) + ); + await Promise.resolve(); + + const inserts: string[] = []; + const updates: string[] = []; + const deletes: string[] = []; + harness.connection.db.user.onInsert((_ctx, row) => + inserts.push(row.username) + ); + harness.connection.db.user.onUpdate((_ctx, _old, row) => + updates.push(row.username) + ); + harness.connection.db.user.onDelete((_ctx, row) => + deletes.push(row.username) + ); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + // The replay returns the same row it had before. + const batch = lastSubscribeBatch(harness)!; + harness.factory.current.sendToClient( + ServerMessage.SubscribeBatchApplied({ + requestId: batch.requestId, + results: [ + { + querySetId: batch.sets[0].querySetId, + outcome: { + tag: 'Applied', + value: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alice', + }), + }, + }, + ], + }, + }, + }, + ], + }) + ); + await Promise.resolve(); + + expect(inserts).toEqual([]); + expect(updates).toEqual([]); + expect(deletes).toEqual([]); + // The row is still readable from the cache. + expect(harness.connection.db.user.count()).toBe(1n); + }); + + test('a row which changed during the outage produces one update callback', async () => { + const harness = await establishWithSubscription(); + const querySetId = harness.factory.current.outgoingMessages.find( + message => message.tag === 'Subscribe' + )!.value.querySetId.id; + + harness.factory.current.sendToClient( + ServerMessage.SubscribeApplied({ + requestId: 1, + querySetId: { id: querySetId }, + rows: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alice', + }), + }, + }, + ], + }, + }) + ); + await Promise.resolve(); + + const updates: { from: string; to: string }[] = []; + harness.connection.db.user.onUpdate((_ctx, oldRow, newRow) => + updates.push({ from: oldRow.username, to: newRow.username }) + ); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + const batch = lastSubscribeBatch(harness)!; + harness.factory.current.sendToClient( + ServerMessage.SubscribeBatchApplied({ + requestId: batch.requestId, + results: [ + { + querySetId: batch.sets[0].querySetId, + outcome: { + tag: 'Applied', + value: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + // The same identity, renamed while we were away. + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alicia', + }), + }, + }, + ], + }, + }, + }, + ], + }) + ); + await Promise.resolve(); + + expect(updates).toEqual([{ from: 'Alice', to: 'Alicia' }]); + }); + + test('a row deleted during the outage produces a delete callback', async () => { + const harness = await establishWithSubscription(); + const querySetId = harness.factory.current.outgoingMessages.find( + message => message.tag === 'Subscribe' + )!.value.querySetId.id; + + harness.factory.current.sendToClient( + ServerMessage.SubscribeApplied({ + requestId: 1, + querySetId: { id: querySetId }, + rows: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alice', + }), + }, + }, + ], + }, + }) + ); + await Promise.resolve(); + + const deletes: string[] = []; + harness.connection.db.user.onDelete((_ctx, row) => + deletes.push(row.username) + ); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + // The replay returns no rows: the row is gone. + const batch = lastSubscribeBatch(harness)!; + harness.factory.current.sendToClient( + ServerMessage.SubscribeBatchApplied({ + requestId: batch.requestId, + results: [ + { + querySetId: batch.sets[0].querySetId, + outcome: { tag: 'Applied', value: { tables: [] } }, + }, + ], + }) + ); + await Promise.resolve(); + + expect(deletes).toEqual(['Alice']); + expect(harness.connection.db.user.count()).toBe(0n); + }); + + test('a rejected replayed query reports its error while the rest apply', async () => { + const harness = build(); + await establish(harness); + + const errors: string[] = []; + const applied: number[] = []; + harness.connection + .subscriptionBuilder() + .onApplied(() => applied.push(1)) + .onError(ctx => errors.push(ctx.event!.message)) + .subscribe(['SELECT * FROM user']); + harness.connection + .subscriptionBuilder() + .onApplied(() => applied.push(2)) + .onError(ctx => errors.push(ctx.event!.message)) + .subscribe(['SELECT * FROM no_such_table']); + await Promise.resolve(); + + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.sendToClient(initialConnection()); + await Promise.resolve(); + + const batch = lastSubscribeBatch(harness)!; + expect(batch.sets).toHaveLength(2); + harness.factory.current.sendToClient( + ServerMessage.SubscribeBatchApplied({ + requestId: batch.requestId, + results: [ + { + querySetId: batch.sets[0].querySetId, + outcome: { tag: 'Applied', value: { tables: [] } }, + }, + { + querySetId: batch.sets[1].querySetId, + outcome: { tag: 'Error', value: 'no such table: no_such_table' }, + }, + ], + }) + ); + await Promise.resolve(); + + expect(errors).toEqual(['no such table: no_such_table']); + // The healthy set applied, firing its onApplied again on the new connection. + expect(applied).toContain(1); + }); +}); + +describe('token provider', () => { + test('is not called while the retained token has life left', async () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const longLived = `header.${Buffer.from( + JSON.stringify({ iat: nowSeconds, exp: nowSeconds + 3600 }) + ).toString('base64url')}.sig`; + + const provider = vi.fn(async () => 'fresh-token'); + const harness = build({ token: longLived, tokenProvider: provider }); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + + expect(provider).not.toHaveBeenCalled(); + expect(harness.factory.current.connectArgs?.authToken).toBe(longLived); + }); + + test('supplies a fresh token when the retained one is close to expiring', async () => { + const nowSeconds = Math.floor(Date.now() / 1000); + const expiring = `header.${Buffer.from( + JSON.stringify({ iat: nowSeconds - 3595, exp: nowSeconds + 5 }) + ).toString('base64url')}.sig`; + + const provider = vi.fn(async () => 'fresh-token'); + const harness = build({ token: expiring, tokenProvider: provider }); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + + expect(provider).toHaveBeenCalled(); + expect(harness.factory.current.connectArgs?.authToken).toBe('fresh-token'); + }); + + test('a provider failure counts as a failed attempt, not a terminal error', async () => { + const provider = vi.fn(async () => { + throw new Error('token endpoint down'); + }); + const harness = build({ token: 'opaque', tokenProvider: provider }); + await establish(harness); + + harness.factory.current.close(); + await runReconnectTimer(harness); + + expect(harness.connectErrors).toHaveLength(1); + expect(harness.connectErrors[0].error?.message).toBe('token endpoint down'); + // The SDK keeps trying. + expect(harness.connectErrors[0].nextReconnectAttempt).toBe(2); + }); +}); + +describe('reconnect regressions', () => { + test('uses a fresh connection id for every attempt and retains the session id', async () => { + const harness = build(); + await establish(harness); + const first = harness.factory.current.connectArgs!; + const establishedId = harness.connection.connectionId.toHexString(); + harness.factory.current.close(); + await runReconnectTimer(harness); + const second = harness.factory.current.connectArgs!; + expect(second.connectionId).not.toBe(establishedId); + expect(second.connectionId).not.toBe(first.connectionId); + expect(second.sessionId).toBe(first.sessionId); + harness.factory.current.close(); + await runReconnectTimer(harness); + expect(harness.factory.current.connectArgs!.connectionId).not.toBe( + second.connectionId + ); + }); + + test('session-busy responses retry without advancing the backoff', async () => { + const harness = build(); + await establish(harness); + harness.factory.current.close(); + for (let i = 0; i < 3; i++) { + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + harness.factory.current.serverClose(4000, 'session busy'); + expect(harness.connectErrors.at(-1)?.nextReconnectAttempt).toBe(1); + expect( + harness.connectErrors.at(-1)?.nextReconnectDelayMs + ).toBeLessThanOrEqual(1500); + } + }); + + test('ignores all late events from a discarded socket', async () => { + const harness = build(); + await establish(harness); + const old = harness.factory.current; + old.error(new Error('lost')); + expect(old.closed).toBe(true); + old.acceptConnection(); + old.sendToClient(initialConnection(bobIdentity)); + old.close(); + expect(harness.connection.isActive).toBe(false); + expect(harness.disconnects).toHaveLength(1); + expect(harness.connects).toHaveLength(1); + await runReconnectTimer(harness); + await establish(harness); + old.sendToClient(initialConnection(bobIdentity)); + expect(harness.connection.identity).toEqual(anIdentity); + expect(harness.connects).toHaveLength(2); + }); + + test('waits for InitialConnection before allowing calls on a reconnect', async () => { + const harness = build(); + await establish(harness); + harness.factory.current.close(); + await runReconnectTimer(harness); + harness.factory.current.acceptConnection(); + expect(harness.connection.isActive).toBe(false); + expect(harness.connection.isReconnecting).toBe(true); + await expect( + harness.connection.callReducer('test', new Uint8Array()) + ).rejects.toBeInstanceOf(DisconnectedError); + expect(harness.factory.current.outgoingMessages).toEqual([]); + }); + + test('does not send a queued call on the next socket after rejecting it', async () => { + const harness = build(); + await establish(harness); + const pending = harness.connection.callReducer('test', new Uint8Array()); + const rejected = expect(pending).rejects.toBeInstanceOf( + UnknownCallResultError + ); + harness.factory.current.close(); + await rejected; + await runReconnectTimer(harness); + await establish(harness); + expect(harness.factory.current.outgoingMessages).toEqual([]); + }); + + test('explicit disconnect reports no error and settles in-flight procedures', async () => { + const harness = build(); + await establish(harness); + const rejected = expect( + harness.connection.callProcedure('test', new Uint8Array()) + ).rejects.toBeInstanceOf(UnknownCallResultError); + harness.connection.disconnect(); + await rejected; + expect(harness.disconnects).toEqual([ + { + error: undefined, + nextReconnectAttempt: undefined, + nextReconnectDelayMs: undefined, + }, + ]); + }); + + test('disconnect from onDisconnect cancels retries immediately', async () => { + const harness = build(); + await establish(harness); + harness.connection['onDisconnect'](() => { + if (!harness.connection.isDisconnectRequested) + harness.connection.disconnect(); + }); + harness.factory.current.close(); + expect(vi.getTimerCount()).toBe(0); + expect(harness.disconnects).toHaveLength(2); + }); + + test('disconnect during token refresh prevents opening another socket', async () => { + let resolveToken!: (token: string) => void; + const provider = vi.fn( + () => + new Promise(resolve => { + resolveToken = resolve; + }) + ); + const harness = build({ tokenProvider: provider }); + await establish(harness); + harness.factory.current.close(); + await vi.runOnlyPendingTimersAsync(); + expect(provider).toHaveBeenCalledOnce(); + harness.connection.disconnect(); + resolveToken('fresh-token'); + await settle(harness); + expect(harness.factory.sockets).toHaveLength(1); + expect(harness.connection.isReconnecting).toBe(false); + }); + + test('subscribes during the outage and onConnect are sent only in the replay batch', async () => { + const harness = build(); + await establish(harness); + harness.connection.subscriptionBuilder().subscribe('SELECT * FROM user'); + harness.factory.current.close(); + harness.connection.subscriptionBuilder().subscribe('SELECT * FROM player'); + harness.connection['onConnect'](conn => + conn.subscriptionBuilder().subscribe('SELECT * FROM user') + ); + await runReconnectTimer(harness); + await establish(harness); + const messages = harness.factory.current.outgoingMessages; + expect(messages.map(m => m.tag)).toEqual(['SubscribeBatch']); + const batch = messages[0]; + if (batch.tag !== 'SubscribeBatch') throw new Error('Expected replay'); + expect(batch.value.sets).toHaveLength(3); + }); + + test.each(['before drop', 'during outage'] as const)( + 'unsubscribe %s ends locally and removes stale rows after reconnect', + async timing => { + const harness = build(); + await establish(harness); + const handle = harness.connection + .subscriptionBuilder() + .subscribe('SELECT * FROM user'); + await Promise.resolve(); + const subscribe = harness.factory.current.outgoingMessages[0]; + if (subscribe.tag !== 'Subscribe') + throw new Error('Expected subscription'); + harness.factory.current.sendToClient( + ServerMessage.SubscribeApplied({ + ...subscribe.value, + rows: { + tables: [ + { + table: 'user', + rows: { + sizeHint: { tag: 'RowOffsets', value: [0n] }, + rowsData: encodeUser({ + identity: anIdentity, + username: 'Alice', + }), + }, + }, + ], + }, + }) + ); + const onEnd = vi.fn(); + if (timing === 'before drop') handle.unsubscribeThen(onEnd); + harness.factory.current.close(); + if (timing === 'during outage') handle.unsubscribeThen(onEnd); + expect(handle.isEnded()).toBe(true); + expect(handle.isActive()).toBe(false); + expect(onEnd).toHaveBeenCalledOnce(); + expect(harness.connection.db.user.count()).toBe(1n); + await runReconnectTimer(harness); + await establish(harness); + expect(harness.factory.current.outgoingMessages).toEqual([]); + expect(harness.connection.db.user.count()).toBe(0n); + } + ); + + test('legacy sockets still emit both error and close events', async () => { + const harness = build({ automaticReconnect: false }); + await establish(harness); + harness.factory.current.error(new Error('network error')); + harness.factory.current.close(); + expect(harness.connectErrors).toHaveLength(1); + expect(harness.disconnects).toHaveLength(1); + }); + + test.each([1002, 1003, 1007, 1008])( + 'protocol/policy close code %i is terminal', + async code => { + const harness = build(); + await establish(harness); + harness.factory.current.serverClose(code); + expect(harness.disconnects[0].nextReconnectAttempt).toBeUndefined(); + expect(harness.connection.isReconnecting).toBe(false); + expect(vi.getTimerCount()).toBe(0); + } + ); +}); + +describe('token rejection classification', () => { + test.each([401, 403])( + 'status %i without a provider is terminal', + async status => { + const harness = build(); + await establish(harness); + harness.factory.current.close(); + harness.factory.connectError = new WebSocketTokenError( + status, + 'Rejected' + ); + await runReconnectTimer(harness); + expect(harness.connectErrors[0].nextReconnectAttempt).toBeUndefined(); + expect(harness.connection.isReconnecting).toBe(false); + } + ); + + test('refreshes a rejected retained token once, then stops if the fresh token is rejected', async () => { + const now = Date.now() / 1000; + const token = `header.${Buffer.from(JSON.stringify({ iat: now, exp: now + 3600 })).toString('base64url')}.sig`; + const provider = vi.fn(async () => 'replacement'); + const harness = build({ token, tokenProvider: provider }); + await establish(harness); + harness.factory.current.close(); + harness.factory.connectError = new WebSocketTokenError(401, 'Unauthorized'); + await runReconnectTimer(harness); + expect(provider).not.toHaveBeenCalled(); + expect(harness.connectErrors[0].nextReconnectAttempt).toBe(2); + await runReconnectTimer(harness); + expect(provider).toHaveBeenCalledOnce(); + expect(harness.connectErrors[1].nextReconnectAttempt).toBeUndefined(); + }); + + test('a token exchange server error retries without classifying it as bad credentials', async () => { + const harness = build(); + await establish(harness); + harness.factory.current.close(); + harness.factory.connectError = new WebSocketTokenError(503, 'Unavailable'); + await runReconnectTimer(harness); + expect(harness.connectErrors[0].nextReconnectAttempt).toBe(2); + }); +}); + +describe('handshake and replay boundaries', () => { + test('subscriptions cancelled before the initial handshake are never sent', async () => { + const harness = build(); + const cancelled = harness.connection + .subscriptionBuilder() + .subscribe('SELECT * FROM user'); + cancelled.unsubscribe(); + harness.connection.subscriptionBuilder().subscribe('SELECT * FROM player'); + await establish(harness); + expect(cancelled.isEnded()).toBe(true); + const messages = harness.factory.current.outgoingMessages; + expect(messages).toHaveLength(1); + expect(messages[0]).toMatchObject({ + tag: 'Subscribe', + value: { queryStrings: ['SELECT * FROM player'] }, + }); + }); + + test('an incomplete replay response is terminal', async () => { + const harness = build(); + await establish(harness); + harness.connection.subscriptionBuilder().subscribe('SELECT * FROM user'); + harness.factory.current.close(); + await runReconnectTimer(harness); + await establish(harness); + const message = harness.factory.current.outgoingMessages[0]; + if (message.tag !== 'SubscribeBatch') throw new Error('Expected replay'); + harness.factory.current.sendToClient( + ServerMessage.SubscribeBatchApplied({ + requestId: message.value.requestId, + results: [], + }) + ); + expect(harness.disconnects.at(-1)?.nextReconnectAttempt).toBeUndefined(); + expect(harness.connection.isReconnecting).toBe(false); + expect(harness.factory.current.closed).toBe(true); + }); + + test('disconnect in onConnect prevents subscription replay', async () => { + const harness = build(); + await establish(harness); + harness.connection.subscriptionBuilder().subscribe('SELECT * FROM user'); + harness.factory.current.close(); + harness.connection['onConnect'](() => harness.connection.disconnect()); + await runReconnectTimer(harness); + await establish(harness); + expect(harness.factory.current.closed).toBe(true); + expect(harness.factory.current.outgoingMessages).toEqual([]); + expect(vi.getTimerCount()).toBe(0); + }); +}); diff --git a/crates/bindings-typescript/tests/table_cache_reconnect.test.ts b/crates/bindings-typescript/tests/table_cache_reconnect.test.ts new file mode 100644 index 00000000000..657e434ab67 --- /dev/null +++ b/crates/bindings-typescript/tests/table_cache_reconnect.test.ts @@ -0,0 +1,62 @@ +import { describe, expect, test } from 'vitest'; +import { ModuleContext, tablesToSchema } from '../src/lib/schema'; +import { table } from '../src/lib/table'; +import { t } from '../src/lib/type_builders'; +import { DbConnectionImpl } from '../src/sdk/db_connection_impl'; +import { EventEmitter } from '../src/sdk/event_emitter'; +import type { EventContextInterface } from '../src/sdk/event_context'; +import { TableCacheImpl, type Operation } from '../src/sdk/table_cache'; +import { WebsocketTestAdapterFactory } from '../src/sdk/websocket_test_adapter'; + +const schema = tablesToSchema(new ModuleContext(), { + item: table({ name: 'item' }, { value: t.string() }), +}); +const remoteModule = { + ...schema, + reducers: [], + procedures: [], + versionInfo: { cliVersion: '2.8.3' }, +}; + +describe('reconciliation without primary keys', () => { + test('preserves unchanged rows and adjusts overlapping subscription references', () => { + const connection = new DbConnectionImpl({ + uri: new URL('ws://localhost'), + nameOrAddress: 'test', + emitter: new EventEmitter(), + remoteModule, + createWSFn: new WebsocketTestAdapterFactory().openWebSocket, + compression: 'none', + lightMode: false, + }); + const cache = new TableCacheImpl( + schema.tables.item + ); + const ctx: EventContextInterface = { + db: connection.db, + reducers: connection.reducers, + isActive: true, + subscriptionBuilder: () => connection.subscriptionBuilder(), + disconnect: () => connection.disconnect(), + event: { id: 'test', tag: 'SubscribeApplied' }, + }; + const insert: Operation<{ value: string }> = { + type: 'insert', + rowId: 'row-bytes', + row: { value: 'same' }, + }; + cache.applyOperations([insert, insert], ctx); + const callbacks = cache.applyOperations( + [...cache.snapshotDeleteOperations(), insert], + ctx, + { skipIdenticalUpdates: true } + ); + expect(callbacks).toEqual([]); + expect(cache.count()).toBe(1n); + expect(cache.snapshotDeleteOperations()).toHaveLength(1); + const deletes = cache.applyOperations([{ ...insert, type: 'delete' }], ctx); + expect(deletes.map(callback => callback.type)).toEqual(['delete']); + expect(cache.count()).toBe(0n); + connection.disconnect(); + }); +}); From 812ee9b067634806a7cbdc3d9eadc3882bc4fa1e Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Fri, 11 Sep 2026 15:01:19 +0200 Subject: [PATCH 09/14] Docs update for reconnect --- .../00200-quickstarts/00300-nodejs.md | 35 ++--- .../00200-quickstarts/00400-typescript.md | 6 + .../00600-clients/00300-connection.md | 30 ++-- .../00700-typescript-reference.md | 143 ++++++++++++++++-- 4 files changed, 177 insertions(+), 37 deletions(-) diff --git a/docs/docs/00100-intro/00200-quickstarts/00300-nodejs.md b/docs/docs/00100-intro/00200-quickstarts/00300-nodejs.md index 0c75b7d8905..65081a59555 100644 --- a/docs/docs/00100-intro/00200-quickstarts/00300-nodejs.md +++ b/docs/docs/00100-intro/00200-quickstarts/00300-nodejs.md @@ -134,35 +134,32 @@ spacetime sql "SELECT * FROM person" - Open `src/main.ts` to see the Node.js client. It uses `DbConnection.builder()` to connect to SpacetimeDB, subscribes to tables, and registers callbacks for insert/delete events. Unlike browser apps, Node.js stores the authentication token in a file instead of localStorage. + Open `src/main.ts` to see the Node.js client. It uses `DbConnection.builder()` to connect to SpacetimeDB, subscribes to tables, and registers callbacks for insert/delete events. Unlike browser apps, Node.js stores the authentication token in a file instead of localStorage. Enable `withAutomaticReconnect()` to recover after connection loss. Register subscriptions and row callbacks once, outside `onConnect`, which fires again on reconnect. ```typescript import { DbConnection } from './module_bindings/index.js'; -DbConnection.builder() +const conn = DbConnection.builder() .withUri(HOST) .withDatabaseName(DB_NAME) - .withToken(loadToken()) // Load saved token from file - .onConnect((conn, identity, token) => { + .withToken(loadToken()) + .withAutomaticReconnect() + .onConnect((_conn, identity, token) => { console.log('Connected! Identity:', identity.toHexString()); - saveToken(token); // Save token for future connections - - // Subscribe to all tables - conn.subscriptionBuilder() - .onApplied((ctx) => { - // Show current people - const people = [...ctx.db.person.iter()]; - console.log('Current people:', people.length); - }) - .subscribeToAllTables(); - - // Listen for table changes - conn.db.person.onInsert((ctx, person) => { - console.log(`[Added] ${person.name}`); - }); + saveToken(token); }) .build(); + +conn.subscriptionBuilder() + .onApplied(ctx => { + console.log('Current people:', [...ctx.db.person.iter()].length); + }) + .subscribeToAllTables(); + +conn.db.person.onInsert((_ctx, person) => { + console.log(`[Added] ${person.name}`); +}); ```` diff --git a/docs/docs/00100-intro/00200-quickstarts/00400-typescript.md b/docs/docs/00100-intro/00200-quickstarts/00400-typescript.md index 3c83bde433b..5b2be5a7766 100644 --- a/docs/docs/00100-intro/00200-quickstarts/00400-typescript.md +++ b/docs/docs/00100-intro/00200-quickstarts/00400-typescript.md @@ -120,6 +120,12 @@ spacetime logs +## Reconnect after connection loss + +For browser or Node.js clients, enable `.withAutomaticReconnect()` on your generated `DbConnection` builder. Subscriptions and row callbacks survive reconnects; register them once rather than inside `onConnect`. If your auth tokens expire, also provide an initial token with `.withToken(initialToken)` and a refresh callback with `.withTokenProvider(() => auth.getAccessToken())`. + +See [automatic reconnection](../../00200-core-concepts/00600-clients/00700-typescript-reference.md#method-withautomaticreconnect) for lifecycle callbacks and framework behavior. + ## Next steps - See the [Chat App Tutorial](../00300-tutorials/00100-chat-app.md) for a complete example diff --git a/docs/docs/00200-core-concepts/00600-clients/00300-connection.md b/docs/docs/00200-core-concepts/00600-clients/00300-connection.md index 009053f2aaf..32be97adb1e 100644 --- a/docs/docs/00200-core-concepts/00600-clients/00300-connection.md +++ b/docs/docs/00200-core-concepts/00600-clients/00300-connection.md @@ -29,6 +29,7 @@ import { DbConnection } from './module_bindings'; const conn = DbConnection.builder() .withUri("https://maincloud.spacetimedb.com") .withDatabaseName("my_database") + .withAutomaticReconnect() .build(); ``` @@ -84,6 +85,7 @@ To connect to a database hosted on MainCloud: const conn = DbConnection.builder() .withUri("https://maincloud.spacetimedb.com") .withDatabaseName("my_database") + .withAutomaticReconnect() .build(); ``` @@ -245,17 +247,25 @@ const TOKEN_KEY = `${HOST}/${DB_NAME}/auth_token`; const conn = DbConnection.builder() .withUri(HOST) .withDatabaseName(DB_NAME) + .withToken(localStorage.getItem(TOKEN_KEY) ?? undefined) + .withAutomaticReconnect() .onConnect((conn, identity, token) => { console.log(`Connected! Identity: ${identity.toHexString()}`); // Save token for reconnection — keyed per server/database localStorage.setItem(TOKEN_KEY, token); }) - .onConnectError((_ctx, error) => { - console.error(`Connection failed:`, error); + .onConnectError((_ctx, error, attempt, delayMs) => { + console.error('Connection failed:', error); + if (attempt !== undefined) console.log(`Retry ${attempt} in ${delayMs} ms`); }) - .onDisconnect(() => { - console.log('Disconnected from SpacetimeDB'); - }); + .onDisconnect((_ctx, error, attempt, delayMs) => { + if (attempt !== undefined) { + console.log(`Connection lost; retry ${attempt} in ${delayMs} ms`, error); + } else { + console.log('Connection ended', error); + } + }) + .build(); ``` @@ -397,13 +407,15 @@ Conn->Disconnect(); ### Reconnection Behavior -:::note[Reconnection behavior] +For TypeScript, add `.withAutomaticReconnect()` to the builder to recover after an established connection drops. The connection object, cache, table handles, and callbacks remain usable. Cache reads serve the last known data while `isReconnecting` is `true`; subscriptions are replayed and reconciled after reconnecting. Register subscriptions and row callbacks once, outside `onConnect`, because `onConnect` fires after every successful reconnect. -Lower-level `DbConnection` objects do not reconnect themselves. If you create a `DbConnection` directly and the connection is interrupted, create a new `DbConnection` to re-establish connectivity. We recommend implementing reconnection logic in your application if reliable connectivity is critical. +`onDisconnect` and `onConnectError` receive optional `nextReconnectAttempt` and `nextReconnectDelayMs` parameters. Both are `undefined` when the core SDK will not retry. Initial connection failures are not retried, and `disconnect()` cancels recovery. Reducer and procedure calls made during an outage fail immediately; in-flight calls fail with an unknown-result error because they may have executed. -The TypeScript React, Solid, and Svelte providers manage their connections through the SDK's shared connection manager. While a provider is mounted, that manager automatically rebuilds unexpectedly closed connections with exponential backoff and re-checks connection liveness when the page becomes visible, regains focus, returns online, or is restored from the back-forward cache. +Use `.withTokenProvider(() => auth.getAccessToken())` alongside `.withToken(initialToken)` for expiring credentials. The provider runs before reconnect attempts when the retained token needs refreshing, not periodically while connected. -::: +The TypeScript React, Solid, and Svelte providers enable automatic reconnection through their shared connection manager. Vue and Angular require `.withAutomaticReconnect()` on the provider's builder. See the [TypeScript reference](./00700-typescript-reference.md#method-withautomaticreconnect) for retry policy, token refresh, and framework behavior. This feature requires a server that supports session IDs and batch subscriptions. + +For TypeScript connections without this option, and for other SDKs described on this page, create a new connection if you need to recover after a connection loss. ## Connection Identity diff --git a/docs/docs/00200-core-concepts/00600-clients/00700-typescript-reference.md b/docs/docs/00200-core-concepts/00600-clients/00700-typescript-reference.md index 4980dc1b9f5..b6a04328131 100644 --- a/docs/docs/00200-core-concepts/00600-clients/00700-typescript-reference.md +++ b/docs/docs/00200-core-concepts/00600-clients/00700-typescript-reference.md @@ -129,6 +129,8 @@ Construct a `DbConnection` by calling `DbConnection.builder()` and chaining conf | [`onConnectError` callback](#callback-onconnecterror) | Register a callback to run if the connection is rejected or the host is unreachable. | | [`onDisconnect` callback](#callback-ondisconnect) | Register a callback to run when the connection ends. | | [`withToken` method](#method-withtoken) | Supply a token to authenticate with the remote database. | +| [`withAutomaticReconnect` method](#method-withautomaticreconnect) | Keep the connection and subscriptions usable across connection loss. | +| [`withTokenProvider` method](#method-withtokenprovider) | Refresh credentials before reconnecting. | | [`build` method](#method-build) | Finalize configuration and connect. | #### Method `withUri` @@ -177,29 +179,41 @@ class DbConnectionBuilder { Chain a call to `.onConnect(callback)` to your builder to register a callback to run when your new `DbConnection` successfully initiates its connection to the remote database. The callback accepts three arguments: a reference to the `DbConnection`, the `Identity` by which SpacetimeDB identifies this connection, and a private access token which can be saved and later passed to [`withToken`](#method-withtoken) to authenticate the same user in future connections. +`onConnect` fires again after each successful automatic reconnect. Register row callbacks and subscriptions once, outside this callback, to avoid accumulating duplicate listeners or subscriptions. Use `onConnect` for work needed on every connection, such as saving the token. + #### Callback `onConnectError` ```typescript class DbConnectionBuilder { public onConnectError( - callback: (ctx: ErrorContext, error: Error) => void + callback: ( + ctx: ErrorContext, + error: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void ): DbConnectionBuilder; } ``` -Chain a call to `.onConnectError(callback)` to your builder to register a callback to run when your connection fails. +Called when the initial connection or a reconnect attempt fails before `onConnect`. When another attempt is scheduled, `nextReconnectAttempt` is its one-based number and `nextReconnectDelayMs` is the delay in milliseconds. Both are `undefined` when the SDK will not retry. Initial connection failures are never retried by the core SDK. #### Callback `onDisconnect` ```typescript class DbConnectionBuilder { public onDisconnect( - callback: (ctx: ErrorContext, error: Error | null) => void + callback: ( + ctx: ErrorContext, + error?: Error, + nextReconnectAttempt?: number, + nextReconnectDelayMs?: number + ) => void ): DbConnectionBuilder; } ``` -Chain a call to `.onDisconnect(callback)` to your builder to register a callback to run when your `DbConnection` disconnects from the remote database, either as a result of a call to [`disconnect`](#method-disconnect) or due to an error. +Called when an established connection is lost, or when the application calls [`disconnect`](#method-disconnect). With automatic reconnection enabled, the trailing parameters describe the next attempt and delay in milliseconds. Both are `undefined` when no retry is scheduled. An explicit `disconnect()` also passes `undefined` for `error`, including when called between reconnect attempts. #### Method `withToken` @@ -209,7 +223,91 @@ class DbConnectionBuilder { } ``` -Chain a call to `.withToken(token)` to your builder to provide an OpenID Connect compliant JSON Web Token to authenticate with, or to explicitly select an anonymous connection. If this method is not called or `null` is passed, SpacetimeDB will generate a new `Identity` and sign a new private access token for the connection. +Chain a call to `.withToken(token)` to your builder to provide an OpenID Connect compliant JSON Web Token to authenticate with, or to explicitly select an anonymous connection. If this method is not called or `undefined` is passed, SpacetimeDB will generate a new `Identity` and sign a new private access token for the connection. + +#### Method `withAutomaticReconnect` + +```typescript +class DbConnectionBuilder { + public withAutomaticReconnect(): this; +} +``` + +Enable automatic reconnection after an established connection drops. This is opt-in for plain TypeScript connections and requires a server with session IDs and batch subscription support. + +The SDK keeps the same connection object, table handles, subscription handles, and registered callbacks. Each new connection has a fresh `ConnectionId`, while the SDK retains the authentication token to preserve the client's `Identity`, including for anonymous clients. + +Retries continue until `disconnect()` or a recognized terminal failure, such as a changed identity, rejected credentials with no remaining refresh attempt, or a fatal protocol error. The base delays are 1, 2, 4, 8, 16, and 30 seconds, with ±50% jitter and a final 30-second cap. A successful connection resets the backoff. Initial connection failures do not retry. Browser resume events also check for silently closed sockets and bring scheduled retries forward. + +```typescript +import { DbConnection, tables } from './module_bindings'; + +const conn = DbConnection.builder() + .withUri('http://localhost:3000') + .withDatabaseName('my-database') + .withAutomaticReconnect() + .onConnect((_conn, identity) => { + console.log('Connected:', identity.toHexString()); + }) + .onDisconnect((_ctx, error, attempt, delayMs) => { + if (attempt !== undefined) { + console.log(`Disconnected; retry ${attempt} in ${delayMs} ms`, error); + } else { + console.log('Connection ended', error); + } + }) + .onConnectError((_ctx, error, attempt, delayMs) => { + if (attempt !== undefined) { + console.log(`Connect failed; retry ${attempt} in ${delayMs} ms`, error); + } else { + console.error('Connection failed; no retry scheduled', error); + } + }) + .build(); + +conn.db.player.onInsert((_ctx, player) => console.log(player)); +conn.subscriptionBuilder() + .onApplied(() => console.log('Subscription applied')) + .subscribe(tables.player); +``` + +While reconnecting: + +- `conn.isActive` is `false` and `conn.isReconnecting` is `true`. +- Cache reads still return the last known rows, which may be stale. +- Reducer and procedure calls fail immediately with `DisconnectedError` and are not sent. Calls already in flight fail with `UnknownCallResultError`: they may have executed, so retrying them could repeat an effect. +- New subscriptions are retained for the next connection. Unsubscribing ends the handle locally and removes it from replay. + +On reconnect, the SDK replays subscriptions in one batch and reconciles the cache before notifying callbacks. Unchanged rows produce no row callbacks; changed rows produce the usual insert, update, or delete callbacks. Each replayed subscription fires `onApplied` again, or `onError` if rejected. `onConnect` precedes subscription replay, so wait for `onApplied` when you need refreshed data. + +#### Method `withTokenProvider` + +```typescript +class DbConnectionBuilder { + public withTokenProvider(provider: () => Promise): this; +} +``` + +Supply fresh credentials for reconnect attempts. This method does not enable automatic reconnection by itself, and the provider is not called for the initial connection. Obtain the initial token first and pass it to `withToken`: + +```typescript +const initialToken = await auth.getAccessToken(); +const conn = DbConnection.builder() + .withUri('https://maincloud.spacetimedb.com') + .withDatabaseName('my-database') + .withToken(initialToken) + .withAutomaticReconnect() + .withTokenProvider(() => auth.getAccessToken()) + .build(); +``` + +Here `auth` is your application's authentication client. Its method should return a usable token, refreshing it through your identity provider when necessary. The returned token must identify the same user; use a new connection for sign-in or account changes. + +Before each reconnect attempt, the SDK reads the retained JWT's `exp` and `iat` claims. It calls the provider when the remaining validity is at most 5% of the token's lifetime, with a minimum margin of 30 seconds. Without `iat`, it uses the 30-second margin. If expiry cannot be read, it calls the provider on every attempt. + +A recognized rejection of the retained token forces a refresh on the next attempt. Rejection of a freshly supplied token is terminal. If the provider throws or rejects, the attempt fails and the SDK retries with backoff. A token exchange service outage is also retryable. + +Refresh happens before reconnecting, not on a periodic timer while connected. The SDK retains tokens in memory; persistence across page reloads remains the application's responsibility. #### Method `build` @@ -292,7 +390,7 @@ interface DbContext { } ``` -Gracefully close the `DbConnection`. Throws an error if the connection is already disconnected. +Close the `DbConnection`. With automatic reconnection enabled, this cancels any scheduled or pending reconnect and fires `onDisconnect` with no error or retry parameters. The connection cannot be restarted; build a new one to connect again. ### Subscribe to queries @@ -563,7 +661,17 @@ interface DbContext { } ``` -`true` if the connection has not yet disconnected. Note that a connection `isActive` when it is constructed, before its [`onConnect` callback](#callback-onconnect) is invoked. +Whether the connection is currently active. With automatic reconnection enabled, this remains `false` until the server's initial connection message arrives, and becomes `false` again during an outage. It does not indicate whether subscriptions have finished applying. + +#### Field `isReconnecting` + +```typescript +class DbConnection { + readonly isReconnecting: boolean; +} +``` + +`true` after losing an established connection while the SDK is waiting for or attempting a reconnect. It is `false` during the initial connection, after a successful reconnect, and after the connection ends. This field is on `DbConnection`, not the general `DbContext` interface. ## Type `EventContext` @@ -1034,7 +1142,11 @@ The SpacetimeDB TypeScript SDK includes React bindings under the `spacetimedb/re The React integration is fully compatible with React StrictMode and correctly handles the double-mount behavior (only one WebSocket connection is created). -While a `SpacetimeDBProvider` is mounted, the shared connection manager also replaces the managed `DbConnection` if the underlying WebSocket closes or reports a connection error. Reconnect attempts use exponential backoff, starting at 1 second and doubling after each consecutive failure up to a 30 second maximum; the backoff resets after a successful connection. In browser environments, the manager also re-checks connection liveness when the page becomes visible, regains focus, returns online, or is restored from the back-forward cache, so a stalled reconnect or silently closed socket can be rebuilt promptly after a suspended tab resumes. Hooks such as `useTable` observe the provider state, receive the fresh connection, and establish their subscriptions again; while the replacement connection is being established, `useTable` reports `isReady` as `false` until its subscription is applied on the new connection. This provider-level recovery does not change the lower-level `DbConnection` contract: applications that create a `DbConnection` directly are still responsible for creating a new connection if they need reconnection behavior. +The React provider enables core automatic reconnection. After an established connection drops, it retains the same `DbConnection` and reflects its lifecycle events in provider state. `useSpacetimeDB().isActive` is `false` during the outage; `useTable` reports `isReady` as `false` until its subscription applies again. Cached rows can remain visible while stale, so use these flags for a connection-status indicator. + +The shared connection manager still creates a replacement connection when the core SDK reports that it will not retry, including initial connection failures. Those replacements use the manager's existing exponential backoff. Calling `disconnect()` explicitly prevents this recovery. + +Pass `withTokenProvider` on the provider's builder when your credentials expire. Keep the builder stable across renders, as in the example below. | Name | Description | | ----------------------------------------------------------- | --------------------------------------------------------- | @@ -1175,7 +1287,20 @@ An opaque identifier for a client connection to a database, intended to differen ## Framework Integrations -The SpacetimeDB TypeScript SDK includes built-in integrations for React, SolidJS, Vue, and Svelte. These provide reactive hooks that automatically subscribe to queries and re-render when data changes. +The SpacetimeDB TypeScript SDK includes built-in integrations for React, SolidJS, Vue, Svelte, and Angular. + +React, Solid, and Svelte use the shared connection manager, which enables automatic reconnection on their builders. Vue and Angular build connections directly: add `.withAutomaticReconnect()` to the builder passed to their provider. All integrations accept `.withTokenProvider(...)` on that builder. These settings belong on the builder, not on individual table hooks. + +For example, configure a Vue or Angular connection builder with: + +```typescript +const connectionBuilder = DbConnection.builder() + .withUri('http://localhost:3000') + .withDatabaseName('my-database') + .withAutomaticReconnect(); +``` + +For expiring credentials, obtain the initial token with your auth client, then add `.withToken(initialToken).withTokenProvider(() => auth.getAccessToken())`. See [token refresh](#method-withtokenprovider) for when the provider runs. ### React From bdc2d015798e8a6fde4a0f447d7bed18abf59a04 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Tue, 15 Sep 2026 10:03:23 +0200 Subject: [PATCH 10/14] re-ran code gen --- crates/bindings-typescript/src/sdk/client_api/index.ts | 7 +++++-- crates/bindings-typescript/src/sdk/client_api/types.ts | 1 + 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/bindings-typescript/src/sdk/client_api/index.ts b/crates/bindings-typescript/src/sdk/client_api/index.ts index 85be7b0c365..bcccb6c46dd 100644 --- a/crates/bindings-typescript/src/sdk/client_api/index.ts +++ b/crates/bindings-typescript/src/sdk/client_api/index.ts @@ -1,7 +1,7 @@ // THIS FILE IS AUTOMATICALLY GENERATED BY SPACETIMEDB. EDITS TO THIS FILE // WILL NOT BE SAVED. MODIFY TABLES IN YOUR MODULE SOURCE CODE INSTEAD. -// This was generated using spacetimedb cli version 2.0.0 (commit 098afaf1a5ed935bce5a32c88620e829506effe7). +// This was generated using spacetimedb cli version 2.8.3 (commit 3f9d3285a3e677d928a068785e6d53dc1f7a53d1). /* eslint-disable */ /* tslint:disable */ @@ -53,7 +53,7 @@ const proceduresSchema = __procedures(); /** The remote SpacetimeDB module schema, both runtime and type information. */ const REMOTE_MODULE = { versionInfo: { - cliVersion: '2.0.0' as const, + cliVersion: '2.8.3' as const, }, tables: tablesSchema.schemaType.tables, reducers: reducersSchema.reducersType.reducers, @@ -73,6 +73,9 @@ export const reducers = __convertToAccessorMap( reducersSchema.reducersType.reducers ); +/** The procedures available in this remote SpacetimeDB module. */ +export const procedures = __convertToAccessorMap(proceduresSchema.procedures); + /** The context type returned in callbacks for all possible events. */ export type EventContext = __EventContextInterface; /** The context type returned in callbacks for reducer events. */ diff --git a/crates/bindings-typescript/src/sdk/client_api/types.ts b/crates/bindings-typescript/src/sdk/client_api/types.ts index ce4f4ed0680..94c1501f6d6 100644 --- a/crates/bindings-typescript/src/sdk/client_api/types.ts +++ b/crates/bindings-typescript/src/sdk/client_api/types.ts @@ -253,6 +253,7 @@ export const SubscribeSet = __t.object('SubscribeSet', { }); export type SubscribeSet = __Infer; +// The tagged union or sum type for the algebraic type `SubscribeSetOutcome`. export const SubscribeSetOutcome = __t.enum('SubscribeSetOutcome', { get Applied() { return QueryRows; From 3603a6ddc7f5d0a7a0384d1931073d528aab9be2 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 16 Sep 2026 13:49:40 +0200 Subject: [PATCH 11/14] Update crates/bindings-typescript/src/sdk/db_connection_impl.ts Co-authored-by: Jason Larabie Signed-off-by: Alessandro Asoni --- crates/bindings-typescript/src/sdk/db_connection_impl.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/bindings-typescript/src/sdk/db_connection_impl.ts b/crates/bindings-typescript/src/sdk/db_connection_impl.ts index 0ab7a2e9ba4..c6cf761877c 100644 --- a/crates/bindings-typescript/src/sdk/db_connection_impl.ts +++ b/crates/bindings-typescript/src/sdk/db_connection_impl.ts @@ -754,8 +754,11 @@ export class DbConnectionImpl const attempt = this.#reconnectAttempt + 1; const delayMs = computeReconnectDelayMs(attempt); +try { this.#emitter.emit('connectError', this, error, attempt, delayMs); +} finally { this.#scheduleReconnect(attempt, delayMs); +} } #scheduleReconnect(attempt: number, delayMs: number): void { From 26cd607ee435497620817eac15a80962c4b1f7e7 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 16 Sep 2026 13:49:55 +0200 Subject: [PATCH 12/14] Update crates/bindings-typescript/src/sdk/db_connection_impl.ts Co-authored-by: Jason Larabie Signed-off-by: Alessandro Asoni --- crates/bindings-typescript/src/sdk/db_connection_impl.ts | 3 +++ 1 file changed, 3 insertions(+) diff --git a/crates/bindings-typescript/src/sdk/db_connection_impl.ts b/crates/bindings-typescript/src/sdk/db_connection_impl.ts index c6cf761877c..561159ec773 100644 --- a/crates/bindings-typescript/src/sdk/db_connection_impl.ts +++ b/crates/bindings-typescript/src/sdk/db_connection_impl.ts @@ -720,8 +720,11 @@ export class DbConnectionImpl const attempt = this.#reconnectAttempt + 1; const delayMs = computeReconnectDelayMs(attempt); +try { this.#emitter.emit('disconnect', this, error, attempt, delayMs); +} finally { this.#scheduleReconnect(attempt, delayMs); +} } #handleAttemptFailure(error: Error): void { From e2c0e9a3df0ecb9eb7c1c51f10323745518e4c97 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 16 Sep 2026 14:01:45 +0200 Subject: [PATCH 13/14] Performance improvement for deepEqual --- crates/bindings-typescript/src/lib/util.ts | 30 +++-- .../tests/deep_equal.test.ts | 126 ++++++++++++++++++ 2 files changed, 145 insertions(+), 11 deletions(-) create mode 100644 crates/bindings-typescript/tests/deep_equal.test.ts diff --git a/crates/bindings-typescript/src/lib/util.ts b/crates/bindings-typescript/src/lib/util.ts index 143d54bdfe0..fe25558251c 100644 --- a/crates/bindings-typescript/src/lib/util.ts +++ b/crates/bindings-typescript/src/lib/util.ts @@ -6,11 +6,8 @@ import type { ParamsObj } from './reducers'; import type { ColumnBuilder, TypeBuilder } from './type_builders'; import type { CamelCase, SnakeCase } from './type_util'; -export function deepEqual(obj1: any, obj2: any): boolean { - // If both are strictly equal (covers primitives and reference equality), return true +export function deepEqual(obj1: unknown, obj2: unknown): boolean { if (obj1 === obj2) return true; - - // If either is a primitive type or one is null, return false since we already checked for strict equality if ( typeof obj1 !== 'object' || obj1 === null || @@ -20,20 +17,31 @@ export function deepEqual(obj1: any, obj2: any): boolean { return false; } - // Get keys of both objects + let firstKey = 0; + if (obj1 instanceof Uint8Array && obj2 instanceof Uint8Array) { + if (obj1.length !== obj2.length) return false; + for (let i = 0; i < obj1.length; i++) { + if (obj1[i] !== obj2[i]) return false; + } + // Typed-array indices precede other enumerable keys and are already equal. + firstKey = obj1.length; + } + const keys1 = Object.keys(obj1); const keys2 = Object.keys(obj2); - - // If number of keys is different, return false if (keys1.length !== keys2.length) return false; - // Check all keys and compare values recursively - for (const key of keys1) { - if (!keys2.includes(key) || !deepEqual(obj1[key], obj2[key])) { + const values1 = obj1 as Record; + const values2 = obj2 as Record; + for (let i = firstKey; i < keys1.length; i++) { + const key = keys1[i]; + if ( + !Object.prototype.propertyIsEnumerable.call(obj2, key) || + !deepEqual(values1[key], values2[key]) + ) { return false; } } - return true; } diff --git a/crates/bindings-typescript/tests/deep_equal.test.ts b/crates/bindings-typescript/tests/deep_equal.test.ts new file mode 100644 index 00000000000..f6eb9ecf604 --- /dev/null +++ b/crates/bindings-typescript/tests/deep_equal.test.ts @@ -0,0 +1,126 @@ +import { describe, expect, test } from 'vitest'; +import { deepEqual } from '../src/lib/util'; + +function expectEqual(left: unknown, right: unknown, equal: boolean): void { + expect(deepEqual(left, right)).toBe(equal); + expect(deepEqual(right, left)).toBe(equal); +} + +describe('deepEqual', () => { + test('preserves primitive and reference equality', () => { + const object = { value: 1 }; + expectEqual(object, object, true); + expectEqual(null, null, true); + expectEqual(null, {}, false); + expectEqual(undefined, null, false); + expectEqual(1, '1', false); + expectEqual(1n, 1n, true); + expectEqual(0, -0, true); + expectEqual(NaN, NaN, false); + }); + + test('compares large byte arrays by value', () => { + const left = new Uint8Array(100_000).fill(7); + const right = left.slice(); + expectEqual(left, right, true); + right[right.length - 1] = 8; + expectEqual(left, right, false); + expectEqual(left, left.subarray(1), false); + expectEqual(new Uint8Array(), new Uint8Array(), true); + }); + + test('compares views using their own offsets and lengths', () => { + const buffer = new Uint8Array([9, 1, 2, 3, 1, 2, 3, 8]); + expectEqual(buffer.subarray(1, 4), buffer.subarray(4, 7), true); + expectEqual(buffer.subarray(1, 4), new Uint8Array([1, 2, 3]), true); + expectEqual(buffer.subarray(1, 4), buffer.subarray(2, 5), false); + }); + + test('compares blobs nested inside rows and arrays', () => { + const left = { id: 1n, items: [{ data: new Uint8Array([1, 2, 3]) }] }; + const right = { id: 1n, items: [{ data: new Uint8Array([1, 2, 3]) }] }; + expectEqual(left, right, true); + right.items[0].data[2] = 4; + expectEqual(left, right, false); + }); + + test('compares enumerable properties attached to byte arrays', () => { + const left = Object.assign(new Uint8Array([1, 2]), { label: 'a' }); + const right = Object.assign(new Uint8Array([1, 2]), { label: 'a' }); + expectEqual(left, right, true); + right.label = 'b'; + expectEqual(left, right, false); + expectEqual(left, new Uint8Array([1, 2]), false); + expectEqual( + left, + Object.assign(new Uint8Array([1, 2]), { other: 'a' }), + false + ); + }); + + test('compares large ordinary arrays and nested elements', () => { + const left = Array.from({ length: 10_000 }, (_, i) => ({ value: i })); + const right = Array.from({ length: 10_000 }, (_, i) => ({ value: i })); + expectEqual(left, right, true); + right[right.length - 1].value = -1; + expectEqual(left, right, false); + expectEqual([1, 2], [1, 2, 3], false); + }); + + test('preserves sparse-array enumerable-property semantics', () => { + expectEqual(new Array(2), new Array(4), true); + expectEqual(new Array(1), [undefined], false); + const sparse = new Array(2); + sparse[1] = 1; + expectEqual(sparse, [undefined, 1], false); + const left = [1]; + left.length = 100; + expectEqual(left, [1], true); + expectEqual( + Object.assign([1], { label: 'a' }), + Object.assign([1], { label: 'b' }), + false + ); + }); + + test('preserves structural equality across object kinds', () => { + expectEqual([1, 2], { 0: 1, 1: 2 }, true); + expectEqual(new Uint8Array([1, 2]), [1, 2], true); + expectEqual(new Uint8Array([1, 2]), new Uint16Array([1, 2]), true); + }); + + test('ignores property insertion order', () => { + expectEqual( + { first: 1, second: { value: 2 } }, + { second: { value: 2 }, first: 1 }, + true + ); + }); + + test('requires matching own enumerable keys, even for undefined values', () => { + expectEqual({ first: undefined }, { second: undefined }, false); + expectEqual({ first: undefined }, {}, false); + const inherited = Object.create({ first: undefined }) as Record< + string, + unknown + >; + inherited.second = undefined; + expectEqual({ first: undefined }, inherited, false); + const hidden = Object.defineProperty({ second: undefined }, 'first', { + value: undefined, + }); + expectEqual({ first: undefined }, hidden, false); + }); + + test('supports null prototypes and shadowed property-check methods', () => { + const left = Object.assign(Object.create(null) as Record, { + value: 1, + }); + expectEqual(left, { value: 1 }, true); + expectEqual( + { hasOwnProperty: 1, propertyIsEnumerable: 2 }, + { propertyIsEnumerable: 2, hasOwnProperty: 1 }, + true + ); + }); +}); From 1185628601a12d10cf906ad3349eebeb2c7e1b13 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Thu, 17 Sep 2026 15:22:36 +0200 Subject: [PATCH 14/14] Update agent skill files --- .../skills/typescript-client/SKILL.md | 103 ++++++++++++------ skills/typescript-client/SKILL.md | 103 ++++++++++++------ 2 files changed, 144 insertions(+), 62 deletions(-) diff --git a/codex-plugin/plugins/spacetimedb/skills/typescript-client/SKILL.md b/codex-plugin/plugins/spacetimedb/skills/typescript-client/SKILL.md index 3a31183133f..dc0bb44fd50 100644 --- a/codex-plugin/plugins/spacetimedb/skills/typescript-client/SKILL.md +++ b/codex-plugin/plugins/spacetimedb/skills/typescript-client/SKILL.md @@ -18,19 +18,23 @@ Generated bindings convert snake_case names to camelCase, including row fields: ## React: main.tsx ```typescript -import React, { useEffect, useMemo } from 'react'; +import React, { useMemo } from 'react'; import ReactDOM from 'react-dom/client'; import { SpacetimeDBProvider } from 'spacetimedb/react'; import { DbConnection } from './module_bindings'; import { MODULE_NAME, SPACETIMEDB_URI } from './config'; import App from './App'; +const TOKEN_KEY = `${SPACETIMEDB_URI}/${MODULE_NAME}/auth_token`; + function Root() { const connectionBuilder = useMemo(() => DbConnection.builder() .withUri(SPACETIMEDB_URI) .withDatabaseName(MODULE_NAME) - .withToken(localStorage.getItem('auth_token') || undefined), + .withToken(localStorage.getItem(TOKEN_KEY) ?? undefined) + .withAutomaticReconnect() + .onConnect((_conn, _identity, token) => localStorage.setItem(TOKEN_KEY, token)), [] ); return ( @@ -49,42 +53,42 @@ ReactDOM.createRoot(document.getElementById('root')!).render(); import { useTable, useSpacetimeDB } from 'spacetimedb/react'; import { DbConnection, tables } from './module_bindings'; -function App() { - const { isActive, identity: myIdentity, token, getConnection } = useSpacetimeDB(); +export default function App() { + const { isActive, identity: myIdentity, getConnection } = useSpacetimeDB(); const conn = getConnection() as DbConnection | null; - // Save auth token - useEffect(() => { if (token) localStorage.setItem('auth_token', token); }, [token]); - - // Subscribe when connected. Prefer typed query builders over raw SQL - useEffect(() => { - if (!conn || !isActive) return; - conn.subscriptionBuilder() - .onApplied(() => setSubscribed(true)) - .subscribe([tables.entity, tables.record]); - // Or with filters: tables.entity.where(r => r.active.eq(true)) - // Or raw SQL: 'SELECT * FROM entity' - }, [conn, isActive]); - - // Reactive data. Returns [rows, isReady] + // useTable owns subscriptions and reports readiness after replay. const [entities, entitiesReady] = useTable(tables.entity); const [records, recordsReady] = useTable(tables.record); - // useTable with row callbacks const [onlineUsers] = useTable( tables.entity.where(r => r.active.eq(true)), { - onInsert: (user) => console.log('User connected:', user.name), - onDelete: (user) => console.log('User disconnected:', user.name), + onInsert: user => console.log('User connected:', user.name), + onDelete: user => console.log('User disconnected:', user.name), onUpdate: (oldUser, newUser) => console.log('Updated:', newUser.name), } ); - // Call reducers with object syntax - conn?.reducers.addRecord({ data }).catch(console.error); + const addRecord = (data: string) => { + if (!conn || !isActive) return; + void conn.reducers.addRecord({ data }).catch(console.error); + }; + const ownsEntity = entities.some( + row => row.owner.toHexString() === myIdentity?.toHexString() + ); - // Compare identities - const isMe = row.owner.toHexString() === myIdentity?.toHexString(); + return ( +
+

{isActive ? 'Connected' : 'Not connected'}

+

{entitiesReady && recordsReady ? 'Data ready' : 'Waiting for current data'}

+

{onlineUsers.length} users online, {records.length} records

+

{ownsEntity ? 'You own an entity' : 'No owned entity'}

+ +
+ ); } ``` @@ -93,22 +97,59 @@ function App() { ```typescript import { DbConnection, tables } from './module_bindings'; +const HOST = 'wss://maincloud.spacetimedb.com'; +const DATABASE = 'my_module'; +const TOKEN_KEY = `${HOST}/${DATABASE}/auth_token`; + const conn = DbConnection.builder() - .withUri('wss://maincloud.spacetimedb.com') - .withDatabaseName('my_module') - .onConnect((ctx) => { - ctx.subscriptionBuilder() - .onApplied(() => console.log('Ready')) - .subscribe([tables.user, tables.message]); + .withUri(HOST) + .withDatabaseName(DATABASE) + .withToken(localStorage.getItem(TOKEN_KEY) ?? undefined) + .withAutomaticReconnect() + .onConnect((_conn, identity, token) => { + localStorage.setItem(TOKEN_KEY, token); + console.log('Connected as:', identity.toHexString()); + }) + .onDisconnect((_ctx, error, nextAttempt, delayMs) => { + if (nextAttempt !== undefined) { + console.warn(`Reconnect attempt ${nextAttempt} in ${delayMs} ms`, error); + } else { + console.log('Connection ended', error); + } + }) + .onConnectError((_ctx, error, nextAttempt, delayMs) => { + console.error('Connection failed:', error); + if (nextAttempt !== undefined) { + console.log(`Retry ${nextAttempt} in ${delayMs} ms`); + } }) .build(); +// Register once; the SDK replays this subscription after reconnecting. +const subscription = conn.subscriptionBuilder() + .onApplied(() => console.log('Ready')) + .subscribe([tables.user, tables.message]); + // Row callbacks conn.db.user.onInsert((ctx, user) => console.log('Joined:', user.name)); conn.db.user.onDelete((ctx, user) => console.log('Left:', user.name)); conn.db.user.onUpdate((ctx, oldUser, newUser) => console.log('Updated:', newUser.name)); ``` +## Automatic Reconnect and Token Refresh + +The React provider's connection manager enables automatic reconnect; the explicit `.withAutomaticReconnect()` above also shows the setting to use for direct connections. Without it, a direct connection does not recover automatically. For direct connections, initial failures are not retried. After an established connection drops, retries use exponential backoff and jitter with a 30-second cap, until `disconnect()` or a terminal failure. While mounted, the React provider also preserves its separate connection-manager retries: it builds a replacement connection after an initial or terminal failure that the core connection will not retry. It respects an explicit `disconnect()`. + +The same connection, identity, table handles, subscriptions, and row callbacks survive recovery. Each attempt gets a fresh connection ID. `onConnect` runs again before subscription replay, so register subscriptions and row callbacks once, outside that callback. The SDK replays subscriptions in one batch, retains readable but stale cached rows during outages, and emits net row changes after reconciliation. Subscription `onApplied` runs again after replay; keep one-time setup separate. + +In React, let `useTable` manage its own subscriptions. Do not add a second subscription effect for the same queries or recreate subscriptions whenever `isActive` changes. Use the hook's `isReady` result for data readiness: a successful reconnect handshake does not mean replay has completed. Invoke reducers from event handlers, not during rendering. For a manually created subscription, retain its handle and call `unsubscribe()` when it is no longer needed, including during an outage. + +For direct connections, `conn.isReconnecting` reports recovery before the next successful handshake. `onDisconnect` and `onConnectError` receive `(ctx, error, nextReconnectAttempt, nextReconnectDelayMs)`. The last two arguments are `undefined` when no retry is scheduled. Do not build a replacement connection or run your own retry timer while automatic recovery is pending. An explicit `conn.disconnect()` stops recovery, including a pending token refresh result. + +For expiring credentials, also call `.withTokenProvider(() => refreshTokenAsync())`, where your authentication integration supplies `refreshTokenAsync(): Promise`. Supply the initial token with `.withToken(initialToken)`; the provider is used only for reconnect attempts. It must return a non-empty token for the same identity. The SDK calls it when remaining validity is at most 30 seconds or 5% of the original lifetime, whichever is greater, when expiry cannot be read, or after a reused token is rejected. Provider failures retry; rejection of a freshly provided token is terminal. No periodic refresh runs while connected, and disconnecting does not cancel the provider's own asynchronous work. + +Calls made while disconnected fail immediately. Pending reducer and procedure promises reject with `UnknownCallResultError` (exported from `spacetimedb`) when the connection is lost before a result arrives. The server may have executed the operation; the SDK never replays it. Do not automatically retry non-idempotent calls on that error. + ## Gotchas - **`useTable` rows are `readonly`.** Copy before sorting/mutating, or it fails to type-check: diff --git a/skills/typescript-client/SKILL.md b/skills/typescript-client/SKILL.md index 3a31183133f..dc0bb44fd50 100644 --- a/skills/typescript-client/SKILL.md +++ b/skills/typescript-client/SKILL.md @@ -18,19 +18,23 @@ Generated bindings convert snake_case names to camelCase, including row fields: ## React: main.tsx ```typescript -import React, { useEffect, useMemo } from 'react'; +import React, { useMemo } from 'react'; import ReactDOM from 'react-dom/client'; import { SpacetimeDBProvider } from 'spacetimedb/react'; import { DbConnection } from './module_bindings'; import { MODULE_NAME, SPACETIMEDB_URI } from './config'; import App from './App'; +const TOKEN_KEY = `${SPACETIMEDB_URI}/${MODULE_NAME}/auth_token`; + function Root() { const connectionBuilder = useMemo(() => DbConnection.builder() .withUri(SPACETIMEDB_URI) .withDatabaseName(MODULE_NAME) - .withToken(localStorage.getItem('auth_token') || undefined), + .withToken(localStorage.getItem(TOKEN_KEY) ?? undefined) + .withAutomaticReconnect() + .onConnect((_conn, _identity, token) => localStorage.setItem(TOKEN_KEY, token)), [] ); return ( @@ -49,42 +53,42 @@ ReactDOM.createRoot(document.getElementById('root')!).render(); import { useTable, useSpacetimeDB } from 'spacetimedb/react'; import { DbConnection, tables } from './module_bindings'; -function App() { - const { isActive, identity: myIdentity, token, getConnection } = useSpacetimeDB(); +export default function App() { + const { isActive, identity: myIdentity, getConnection } = useSpacetimeDB(); const conn = getConnection() as DbConnection | null; - // Save auth token - useEffect(() => { if (token) localStorage.setItem('auth_token', token); }, [token]); - - // Subscribe when connected. Prefer typed query builders over raw SQL - useEffect(() => { - if (!conn || !isActive) return; - conn.subscriptionBuilder() - .onApplied(() => setSubscribed(true)) - .subscribe([tables.entity, tables.record]); - // Or with filters: tables.entity.where(r => r.active.eq(true)) - // Or raw SQL: 'SELECT * FROM entity' - }, [conn, isActive]); - - // Reactive data. Returns [rows, isReady] + // useTable owns subscriptions and reports readiness after replay. const [entities, entitiesReady] = useTable(tables.entity); const [records, recordsReady] = useTable(tables.record); - // useTable with row callbacks const [onlineUsers] = useTable( tables.entity.where(r => r.active.eq(true)), { - onInsert: (user) => console.log('User connected:', user.name), - onDelete: (user) => console.log('User disconnected:', user.name), + onInsert: user => console.log('User connected:', user.name), + onDelete: user => console.log('User disconnected:', user.name), onUpdate: (oldUser, newUser) => console.log('Updated:', newUser.name), } ); - // Call reducers with object syntax - conn?.reducers.addRecord({ data }).catch(console.error); + const addRecord = (data: string) => { + if (!conn || !isActive) return; + void conn.reducers.addRecord({ data }).catch(console.error); + }; + const ownsEntity = entities.some( + row => row.owner.toHexString() === myIdentity?.toHexString() + ); - // Compare identities - const isMe = row.owner.toHexString() === myIdentity?.toHexString(); + return ( +
+

{isActive ? 'Connected' : 'Not connected'}

+

{entitiesReady && recordsReady ? 'Data ready' : 'Waiting for current data'}

+

{onlineUsers.length} users online, {records.length} records

+

{ownsEntity ? 'You own an entity' : 'No owned entity'}

+ +
+ ); } ``` @@ -93,22 +97,59 @@ function App() { ```typescript import { DbConnection, tables } from './module_bindings'; +const HOST = 'wss://maincloud.spacetimedb.com'; +const DATABASE = 'my_module'; +const TOKEN_KEY = `${HOST}/${DATABASE}/auth_token`; + const conn = DbConnection.builder() - .withUri('wss://maincloud.spacetimedb.com') - .withDatabaseName('my_module') - .onConnect((ctx) => { - ctx.subscriptionBuilder() - .onApplied(() => console.log('Ready')) - .subscribe([tables.user, tables.message]); + .withUri(HOST) + .withDatabaseName(DATABASE) + .withToken(localStorage.getItem(TOKEN_KEY) ?? undefined) + .withAutomaticReconnect() + .onConnect((_conn, identity, token) => { + localStorage.setItem(TOKEN_KEY, token); + console.log('Connected as:', identity.toHexString()); + }) + .onDisconnect((_ctx, error, nextAttempt, delayMs) => { + if (nextAttempt !== undefined) { + console.warn(`Reconnect attempt ${nextAttempt} in ${delayMs} ms`, error); + } else { + console.log('Connection ended', error); + } + }) + .onConnectError((_ctx, error, nextAttempt, delayMs) => { + console.error('Connection failed:', error); + if (nextAttempt !== undefined) { + console.log(`Retry ${nextAttempt} in ${delayMs} ms`); + } }) .build(); +// Register once; the SDK replays this subscription after reconnecting. +const subscription = conn.subscriptionBuilder() + .onApplied(() => console.log('Ready')) + .subscribe([tables.user, tables.message]); + // Row callbacks conn.db.user.onInsert((ctx, user) => console.log('Joined:', user.name)); conn.db.user.onDelete((ctx, user) => console.log('Left:', user.name)); conn.db.user.onUpdate((ctx, oldUser, newUser) => console.log('Updated:', newUser.name)); ``` +## Automatic Reconnect and Token Refresh + +The React provider's connection manager enables automatic reconnect; the explicit `.withAutomaticReconnect()` above also shows the setting to use for direct connections. Without it, a direct connection does not recover automatically. For direct connections, initial failures are not retried. After an established connection drops, retries use exponential backoff and jitter with a 30-second cap, until `disconnect()` or a terminal failure. While mounted, the React provider also preserves its separate connection-manager retries: it builds a replacement connection after an initial or terminal failure that the core connection will not retry. It respects an explicit `disconnect()`. + +The same connection, identity, table handles, subscriptions, and row callbacks survive recovery. Each attempt gets a fresh connection ID. `onConnect` runs again before subscription replay, so register subscriptions and row callbacks once, outside that callback. The SDK replays subscriptions in one batch, retains readable but stale cached rows during outages, and emits net row changes after reconciliation. Subscription `onApplied` runs again after replay; keep one-time setup separate. + +In React, let `useTable` manage its own subscriptions. Do not add a second subscription effect for the same queries or recreate subscriptions whenever `isActive` changes. Use the hook's `isReady` result for data readiness: a successful reconnect handshake does not mean replay has completed. Invoke reducers from event handlers, not during rendering. For a manually created subscription, retain its handle and call `unsubscribe()` when it is no longer needed, including during an outage. + +For direct connections, `conn.isReconnecting` reports recovery before the next successful handshake. `onDisconnect` and `onConnectError` receive `(ctx, error, nextReconnectAttempt, nextReconnectDelayMs)`. The last two arguments are `undefined` when no retry is scheduled. Do not build a replacement connection or run your own retry timer while automatic recovery is pending. An explicit `conn.disconnect()` stops recovery, including a pending token refresh result. + +For expiring credentials, also call `.withTokenProvider(() => refreshTokenAsync())`, where your authentication integration supplies `refreshTokenAsync(): Promise`. Supply the initial token with `.withToken(initialToken)`; the provider is used only for reconnect attempts. It must return a non-empty token for the same identity. The SDK calls it when remaining validity is at most 30 seconds or 5% of the original lifetime, whichever is greater, when expiry cannot be read, or after a reused token is rejected. Provider failures retry; rejection of a freshly provided token is terminal. No periodic refresh runs while connected, and disconnecting does not cancel the provider's own asynchronous work. + +Calls made while disconnected fail immediately. Pending reducer and procedure promises reject with `UnknownCallResultError` (exported from `spacetimedb`) when the connection is lost before a result arrives. The server may have executed the operation; the SDK never replays it. Do not automatically retry non-idempotent calls on that error. + ## Gotchas - **`useTable` rows are `readonly`.** Copy before sorting/mutating, or it fails to type-check: