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", 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..f9579473046 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,16 @@ 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: dropping the future drops the guard. + 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 +375,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..b71ab159c4a 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 registry 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..ac15ee28015 100644 --- a/crates/core/src/subscription/module_subscription_actor.rs +++ b/crates/core/src/subscription/module_subscription_actor.rs @@ -1251,6 +1251,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 +1295,35 @@ 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 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 +1363,20 @@ 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( @@ -1322,6 +1391,49 @@ impl ModuleSubscriptions { self.add_multi_subscription_inner(Some(instance), sender, auth, request, timer, _assert) } + /// Check the row limit and evaluate the initial matching rows of one query + /// set that was just registered with the subscription manager, removing the + /// registration again on failure. Shared between the single- and + /// batch-subscribe v2 paths. + /// + /// The outer `Err` is an internal error which aborts the whole request; + /// the inner `Err` is a per-set error to report to the client. + fn eval_registered_query_set( + &self, + sender: &Arc, + auth: &AuthCtx, + queries: &[Arc], + physical_plans: &HashMap>, + tx: &mut TxId, + query_set_id: ws_v2::QuerySetId, + ) -> Result>, DBError> { + let subscription_metrics = &self.metrics.subscribe; + + if let Err(err) = self.check_new_query_row_limit(queries, physical_plans, tx, auth) { + self.remove_failed_subscription(subscription_metrics, sender.id, FailedSubscription::V2(query_set_id))?; + return Ok(Err(err.to_string().into())); + } + + let Ok((update, metrics)) = self.evaluate_queries(sender.clone(), queries, tx, TableUpdateType::Subscribe) + else { + self.remove_failed_subscription(subscription_metrics, sender.id, FailedSubscription::V2(query_set_id))?; + return Ok(Err("Internal error evaluating queries".into())); + }; + 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" + ))) + } + }; + Ok(Ok((rows, metrics))) + } + fn add_v2_subscription_inner( &self, instance: Option<&mut RefInstance>, @@ -1387,29 +1499,18 @@ impl ModuleSubscriptions { 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 (rows, metrics) = match self.eval_registered_query_set( + &sender, + &auth, + &queries, + &physical_plans, + &mut tx, + request.query_set_id, + )? { + Ok(ok) => ok, + Err(err) => { + send_err_msg(err); + return Ok((None, trapped)); } }; @@ -1419,12 +1520,152 @@ impl ModuleSubscriptions { ws_v2::SubscribeApplied { request_id: request.request_id, query_set_id: request.query_set_id, - rows: ws_v2::QueryRows { tables }, + rows, }, ); Ok((Some(metrics), trapped)) } + + /// Implementation of [`Self::add_batch_subscription`]. + /// + /// Each set is compiled and evaluated by the same code as an individual + /// `Subscribe` ([`Self::compile_queries`], [`Self::eval_registered_query_set`]). + /// What this function adds is the batch-atomic phasing that a plain loop + /// over [`Self::add_v2_subscription_inner`] could not provide: every set is + /// registered under one subscription-manager lock, evaluated against one + /// 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 _); + + // Compile every set. Each set gets its own transaction for compilation, + // exactly as an individual `Subscribe` would; the snapshot that matters + // for the rows is the one taken below, after registration. + for (index, set) in request.sets.iter().enumerate() { + match self.compile_queries( + sender.id.identity, + auth.clone(), + &set.query_strings, + set.query_strings.len(), + subscription_metrics, + ) { + Ok(CompiledQueryBatch { + queries, + physical_plans: set_physical_plans, + auth: _, + mut_tx, + compile_timer: _compile_timer, + }) => { + // Release the compilation transaction; the batch takes its + // own transaction for registration and evaluation. + drop(self.guard_mut_tx(mut_tx, <_>::default())); + physical_plans.extend(set_physical_plans); + compiled_sets.push((index, set.query_set_id, queries)); + } + Err(err) => { + outcomes[index] = Some(ws_v2::SubscribeSetOutcome::Error(err.to_string().into())); + } + } + } + + // We always get the db lock before the subscription lock to avoid deadlocks. + let (mut_tx, _tx_offset) = self.begin_mut_tx(Workload::Subscribe); + + // 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, + // exactly as an individual `Subscribe` would. + let mut total_metrics = ExecutionMetrics::default(); + for (index, query_set_id, queries) in registered { + outcomes[index] = Some( + match self.eval_registered_query_set( + &sender, + &auth, + &queries, + &physical_plans, + &mut tx, + query_set_id, + )? { + Ok((rows, metrics)) => { + total_metrics.merge(metrics); + ws_v2::SubscribeSetOutcome::Applied(rows) + } + Err(err) => ws_v2::SubscribeSetOutcome::Error(err), + }, + ); + } + + // 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>, 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/spacetime.json b/crates/smoketests/spacetime.json new file mode 100644 index 00000000000..fc8abe551ec --- /dev/null +++ b/crates/smoketests/spacetime.json @@ -0,0 +1,3 @@ +{ + "server": "maincloud" +} \ No newline at end of file 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 {