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..e56abc1a36f 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 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, +/// 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 the same transaction offset. +#[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 1a4924cb597..5f3d549d4eb 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: 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, + ) + } + None => None, + }; + let connected = match ClientConnection::call_client_connected_maybe_reject( &mut module_rx, client_id, @@ -292,7 +348,17 @@ 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(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 client = ClientConnection::spawn( client_id, auth.into(), @@ -305,6 +371,12 @@ 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. diff --git a/crates/core/src/client.rs b/crates/core/src/client.rs index 812d03c0701..8e14bf5324f 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, 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_connection.rs b/crates/core/src/client/client_connection.rs index 23e801e4b30..e22ba3616d1 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..d303fdf2b02 --- /dev/null +++ b/crates/core/src/client/client_session_index.rs @@ -0,0 +1,546 @@ +//! 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, 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}; + +/// 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 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 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. 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 +/// removed when that connection ends, so a session never outlives its +/// connection. +#[derive(Default)] +pub struct ClientSessionIndex { + sessions: Mutex>>, +} + +impl ClientSessionIndex { + pub fn new() -> Self { + Self::default() + } + + /// Claim a session for `client`, tearing down the connection it supersedes. + /// + /// 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. + /// + /// 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, + teardown: F, + ) -> SessionClaim + where + F: FnOnce(ClientActorId) -> Fut, + Fut: Future, + { + 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) { + 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; + } + + SessionClaim { + index: self.clone(), + key, + holder: Some(holder), + } + } + + /// 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, + }; + 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(); + } + } + + /// 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 + } +} + +/// 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::{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(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]) + } + + fn another_identity() -> Identity { + 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) + } + + /// 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); + + let (_sender, superseded) = connect(&index, a_database(), client(an_identity(), 1), session).await; + + assert!(superseded.is_none()); + assert_eq!(index.len(), 1); + } + + #[tokio::test] + async fn reconnect_supersedes_previous_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) = 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. + 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. + #[tokio::test] + async fn reconnect_reusing_connection_id_supersedes() { + let index = index(); + let session = SessionId::from_u128(7); + 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 (_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); + } + + #[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 (_second, superseded) = + connect(&index, a_database(), client(an_identity(), 2), SessionId::from_u128(8)).await; + + assert!(superseded.is_none()); + 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() { + 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); + let (_sender, _) = connect(&index, a_database(), connection, session).await; + + index.release_session(a_database(), connection, session); + + assert!(index.is_empty()); + } + + #[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); + 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(a_database(), old, session); + + // The replacement still holds the session. + 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. + #[tokio::test] + async fn dropped_claim_releases_the_session() { + let index = index(); + let session = SessionId::from_u128(7); + + 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!( + 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/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 5bef4da58ac..ec0d24dc8e5 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 25eb09e6382..5c60a491d74 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:#}"); @@ -2503,6 +2532,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 aee984fe3a5..c2cc89472ec 100644 --- a/crates/core/src/host/wasm_common/module_host_actor.rs +++ b/crates/core/src/host/wasm_common/module_host_actor.rs @@ -1205,6 +1205,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 c88da72e8d2..f5ba88774b2 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, @@ -218,6 +219,69 @@ 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>, + 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> { + 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 +1133,41 @@ 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, + 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,8 +1185,8 @@ 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, - &auth, + mut_tx, + auth, )? { add_compiled_query( compiled, @@ -1133,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(), @@ -1155,11 +1236,9 @@ 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 +1330,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 +1374,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 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 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. + #[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 +1443,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( @@ -1331,100 +1481,253 @@ 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 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(), + tx_offset, + ws_v2::SubscribeApplied { + request_id, + query_set_id, + rows, + }, + ); + + Ok((Some(metrics), trapped)) + } + + /// Implementation of [`Self::add_batch_subscription`]. + /// + /// 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>>, + sender: Arc, + auth: AuthCtx, + request: ws_v2::SubscribeBatch, + _timer: Instant, + _assert: Option, + ) -> Result<(Option, bool), DBError> { + let ws_v2::SubscribeBatch { request_id, sets } = request; + + 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 _ = 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 = request.query_strings.len(); + + let num_queries: usize = sets.iter().map(|set| set.query_strings.len()).sum(); 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 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::>(); + + // 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. + 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. + 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, subscription_metrics, &mut_tx) { + Ok(CompiledQueries { + queries, + physical_plans: set_physical_plans, + compile_timer: _compile_timer, + }) => { + physical_plans.extend(set_physical_plans); + compiled_sets.push((index, queries)); + } + Err(err) => outcomes[index] = ws_v2::SubscribeSetOutcome::Error(err.to_string().into()), + } + } // 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 = { + // 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, 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()), + } + } + } - subscriptions.add_subscription_v2(sender.clone(), queries, request.query_set_id)? - }; + if registered.is_empty() { + // No set was registered, so there is nothing to evaluate, + // 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. + return Ok(SubscribedQuerySets { + outcomes, + tx: None, + tx_offset: None, + metrics: ExecutionMetrics::default(), + trapped: false, + }); + } 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); + self.materialize_views_and_downgrade_tx(mut_tx, instance, ®istered_queries, auth.caller())?; + + // 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(); + 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] = ws_v2::SubscribeSetOutcome::Error(err.to_string().into()); + continue; + } - subscription_metrics.num_queries_evaluated.inc_by(queries.len() as _); + 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] = ws_v2::SubscribeSetOutcome::Error("Internal error evaluating queries".into()); + continue; + }; + tx.metrics.merge(metrics); + total_metrics.merge(metrics); - 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" - ))) - } - }; + subscription_metrics.num_queries_evaluated.inc_by(queries.len() as _); - let _ = self.broadcast_queue.send_client_message_v2( - sender.clone(), - Some(tx_offset), - ws_v2::SubscribeApplied { - request_id: request.request_id, - query_set_id: request.query_set_id, - rows: ws_v2::QueryRows { tables }, - }, - ); + 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" + ))) + } + }; + outcomes[index] = ws_v2::SubscribeSetOutcome::Applied(rows); + } - Ok((Some(metrics), trapped)) + Ok(SubscribedQuerySets { + outcomes, + tx: Some(tx), + tx_offset: Some(tx_offset), + metrics: total_metrics, + trapped, + }) } + fn add_multi_subscription_inner( &self, instance: Option<&mut RefInstance>, @@ -1894,13 +2197,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)?; @@ -2483,6 +2788,161 @@ 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. + #[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 3f7a49041ed..95a074afcd8 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..1f87311f7de --- /dev/null +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -0,0 +1,459 @@ +//! 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 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] +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 {