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/common.rs b/crates/client-api-messages/src/websocket/common.rs index 45e6b4b381f..688baa9eab5 100644 --- a/crates/client-api-messages/src/websocket/common.rs +++ b/crates/client-api-messages/src/websocket/common.rs @@ -53,6 +53,14 @@ pub const SERVER_MSG_COMPRESSION_TAG_BROTLI: u8 = 1; /// The tag recognized by the host and SDKs to mean gzip compression of a `ServerMessage`. pub const SERVER_MSG_COMPRESSION_TAG_GZIP: u8 = 2; +/// Websocket close code sent when a connection supplies a `session_id` +/// which another live connection of the same identity still holds. +/// +/// The server stops that connection and tears it down. The refused client +/// should retry after a short delay, without counting this as a failed +/// connection attempt for backoff purposes. +pub const SESSION_BUSY_CLOSE_CODE: u16 = 4000; + pub type RowSize = u16; pub type RowOffset = u64; diff --git a/crates/client-api-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..e1293a3997a 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -28,7 +28,8 @@ use spacetimedb::client::messages::{ }; use spacetimedb::client::{ ClientActorId, ClientConfig, ClientConnection, ClientConnectionReceiver, DataMessage, MessageExecutionError, - MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, WsVersion, + MessageHandleError, MeteredReceiver, MeteredSender, OutboundMessage, Protocol, SessionBusy, SessionId, + SessionReservation, WsVersion, }; use spacetimedb::host::module_host::ClientConnectedError; use spacetimedb::host::NoSuchModule; @@ -38,6 +39,7 @@ use spacetimedb::worker_metrics::{ record_client_rejection, ClientDisconnectCause, ClientDisconnectRecorder, ClientRejectCause, WORKER_METRICS, }; use spacetimedb::Identity; +use spacetimedb_client_api_messages::websocket::common::SESSION_BUSY_CLOSE_CODE; use spacetimedb_client_api_messages::websocket::v1 as ws_v1; use spacetimedb_client_api_messages::websocket::v2 as ws_v2; use spacetimedb_client_api_messages::websocket::v3 as ws_v3; @@ -86,6 +88,18 @@ 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 still held by a live + /// connection of the same identity, this connection is refused with + /// [`SESSION_BUSY_CLOSE_CODE`] and the old one is torn down. A retry + /// succeeds once the old connection's `client_disconnected` has run, so the + /// module never observes two live connections for one session. + /// See [`spacetimedb::client::ClientSessionIndex`]. + /// + /// Connections which do not supply one behave exactly as before. + pub session_id: Option, #[serde(default)] pub compression: ws_v1::Compression, /// Whether we want "light" responses, tailored to network bandwidth constrained clients. @@ -100,6 +114,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 +152,7 @@ pub async fn handle_websocket( Path(SubscribeParams { name_or_identity }): Path, Query(SubscribeQueryParams { connection_id, + session_id, compression, light, confirmed, @@ -228,6 +262,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)) @@ -236,7 +272,7 @@ where let ws_opts = ctx.websocket_options(); tokio::spawn(async move { - let ws = match ws_upgrade.upgrade(ws_config).await { + let mut ws = match ws_upgrade.upgrade(ws_config).await { Ok(ws) => ws, Err(err) => { record_client_rejection(db_identity, ClientRejectCause::WebsocketUpgradeError); @@ -255,6 +291,28 @@ where log::debug!("websocket: New client connected from {client_log_string}"); + // Reserved before `client_connected` so that no two connections of one + // session run it. Released by the actor's teardown after the + // module-side disconnect, so a retry finds `client_disconnected` run. + let session = match session_id { + Some(session_id) => match sessions.try_reserve(db_identity, client_id, session_id) { + Ok(reservation) => Some(reservation), + Err(SessionBusy) => { + WORKER_METRICS.ws_clients_session_busy.with_label_values(&db_identity).inc(); + log::debug!("websocket: Refusing connection for {client_log_string}: session {session_id} is busy"); + let close = CloseFrame { + code: CloseCode::from(SESSION_BUSY_CLOSE_CODE), + reason: "session busy".into(), + }; + if let Err(e) = ws.close(Some(close)).await { + log::debug!("websocket: Error refusing connection for {client_log_string}: {e}"); + } + return; + } + }, + None => None, + }; + let connected = match ClientConnection::call_client_connected_maybe_reject( &mut module_rx, client_id, @@ -292,7 +350,12 @@ 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); + let actor = |client: ClientConnection, receiver| { + if let Some(session) = &session { + session.establish(&client.sender()); + } + ws_client_actor(ws_opts, client, ws, receiver, session) + }; let client = ClientConnection::spawn( client_id, auth.into(), @@ -509,15 +572,26 @@ async fn ws_client_actor( client: ClientConnection, ws: WebSocketStream, sendrx: ClientConnectionReceiver, + session: Option, ) { - // ensure that even if this task gets cancelled, we always cleanup the connection - let mut client = scopeguard::guard(client, |client| { - tokio::spawn(client.disconnect()); + // Runs the module-side disconnect even if this task gets cancelled. + let mut client = scopeguard::guard((client, session), |(client, session)| { + tokio::spawn(ws_client_teardown(client, session)); }); - ws_client_actor_inner(&mut client, options, ws, sendrx).await; + ws_client_actor_inner(&mut client.0, options, ws, sendrx).await; - ScopeGuard::into_inner(client).disconnect().await; + let (client, session) = ScopeGuard::into_inner(client); + ws_client_teardown(client, session).await; +} + +/// Run the module-side disconnect, then free the connection's session. +/// +/// The session is released only after `client_disconnected` has run, so that +/// a connection retrying for the same session observes it in order. +async fn ws_client_teardown(client: ClientConnection, session: Option) { + client.disconnect().await; + drop(session); } async fn ws_client_actor_inner( diff --git a/crates/core/src/client.rs b/crates/core/src/client.rs index 812d03c0701..7281f72c0f9 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, SessionBusy, SessionId, SessionReservation}; pub use message_handlers::MessageHandleError; pub use message_handlers_v1::MessageExecutionError; pub use messages::OutboundMessage; diff --git a/crates/core/src/client/client_connection.rs b/crates/core/src/client/client_connection.rs index 23e801e4b30..821d303e601 100644 --- a/crates/core/src/client/client_connection.rs +++ b/crates/core/src/client/client_connection.rs @@ -408,6 +408,20 @@ impl ClientConnectionSender { self.cancelled.load(Ordering::Relaxed) } + /// Stop this connection's websocket actor. + /// + /// Used when a newer connection arrives for this connection's session + /// (see [`super::ClientSessionIndex`]), and when a client exceeds its + /// outgoing queue capacity. The actor's teardown runs the module-side + /// disconnect. + pub fn kick(&self, cause: ClientDisconnectCause) { + if let Some(metrics) = &self.metrics { + metrics.disconnect_recorder.record(cause); + } + 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 +469,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 +1188,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..6d968d774f2 --- /dev/null +++ b/crates/core/src/client/client_session_index.rs @@ -0,0 +1,358 @@ +//! Tracking of live client sessions, so that at most one connection serves a +//! session at a time. +//! +//! A client which reconnects automatically sends the same client-generated +//! session id on every connection attempt. Each connection still receives its +//! own [`ConnectionId`] and its own `client_connected` / `client_disconnected` +//! events. The session id only identifies which earlier connection a new one +//! replaces. +//! +//! A client frequently notices a dropped connection before the server does, +//! as the server needs up to its idle timeout to notice an idle peer. When a +//! connection arrives for a session which a live connection still holds, the +//! new connection is refused and the old one is stopped. The session frees up +//! once the old connection is fully closed, including its module-side +//! disconnect, so a retry finds `client_disconnected` already run. +//! +//! [`ConnectionId`]: spacetimedb_lib::ConnectionId + +use std::collections::HashMap; +use std::sync::{Arc, Mutex, Weak}; + +use spacetimedb_lib::Identity; + +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 belongs to one client on one database, so it can only ever be +/// replaced by the same client on the same database. +#[derive(PartialEq, Eq, Clone, Copy, Hash, Debug)] +struct SessionKey { + database_identity: Identity, + client_identity: Identity, + session_id: SessionId, +} + +/// The connection holding a session. +struct SessionHolder { + client: ClientActorId, + /// Stops the connection's actor when a newer connection arrives. + /// + /// Empty while the connection is still running `client_connected` and has + /// no actor yet. Weak so that the map never keeps a sender alive. + sender: Option>, +} + +/// The session is held by another connection. +#[derive(Debug, PartialEq, Eq)] +pub struct SessionBusy; + +/// The live sessions of one host, each mapped to the connection holding it. +#[derive(Default)] +pub struct ClientSessionIndex { + sessions: Mutex>, +} + +impl ClientSessionIndex { + pub fn new() -> Self { + Self::default() + } + + /// Reserve a session for `client`. + /// + /// If another connection holds the session, its actor is stopped and + /// [`SessionBusy`] is returned. The session frees up once that connection + /// has been torn down, so the caller should refuse `client` and let it + /// retry. + /// + /// The returned reservation releases the session when dropped. The caller + /// should complete it with [`SessionReservation::establish`] once the + /// connection has an actor, and drop it only after the connection is fully + /// closed. + pub fn try_reserve( + self: &Arc, + database_identity: Identity, + client: ClientActorId, + session_id: SessionId, + ) -> Result { + let key = SessionKey { + database_identity, + client_identity: client.identity, + session_id, + }; + let mut sessions = self.sessions.lock().expect("session index poisoned"); + if let Some(holder) = sessions.get(&key) { + log::debug!( + "websocket: Connection {} refused, session {session_id} still held by {}", + client.connection_id, + holder.client.connection_id, + ); + if let Some(sender) = holder.sender.as_ref().and_then(Weak::upgrade) { + sender.kick(ClientDisconnectCause::ConnectionSuperseded); + } + return Err(SessionBusy); + } + sessions.insert(key, SessionHolder { client, sender: None }); + Ok(SessionReservation { + index: self.clone(), + key, + client, + }) + } + + /// Remove the session if `client` still holds it. + fn release(&self, key: SessionKey, client: ClientActorId) { + let mut sessions = self.sessions.lock().expect("session index poisoned"); + // Connections are told apart by their name, the host's per-connection + // counter, rather than by their connection id, which a client may repeat. + if sessions.get(&key).is_some_and(|holder| holder.client.name == client.name) { + sessions.remove(&key); + } + } + + /// 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 by one connection for its whole lifetime. +/// +/// Dropping it releases the session, so it must be dropped only after the +/// connection is fully closed, including its module-side disconnect. +pub struct SessionReservation { + index: Arc, + key: SessionKey, + client: ClientActorId, +} + +impl SessionReservation { + /// Record the connection's sender, so that a later connection for the + /// same session can stop it. + pub fn establish(&self, sender: &Arc) { + let mut sessions = self.index.sessions.lock().expect("session index poisoned"); + if let Some(holder) = sessions.get_mut(&self.key) { + holder.sender = Some(Arc::downgrade(sender)); + } + } +} + +impl Drop for SessionReservation { + fn drop(&mut self) { + self.index.release(self.key, self.client); + } +} + +#[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; + + /// 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` matches its connection id, as the websocket + /// handler would assign distinct names to distinct connections. + fn client(identity: Identity, connection_id: u128) -> ClientActorId { + ClientActorId { + identity, + 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]) + } + + fn session() -> SessionId { + SessionId::from_u128(7) + } + + /// Reserve a session and establish it, as a connected client does. + fn connect( + index: &Arc, + database: Identity, + client: ClientActorId, + session: SessionId, + ) -> (SessionReservation, Arc) { + let reservation = index.try_reserve(database, client, session).expect("session should be free"); + let sender = sender(client); + reservation.establish(&sender); + (reservation, sender) + } + + #[tokio::test] + async fn first_connection_reserves_the_session() { + let index = index(); + + let (_first, _) = connect(&index, a_database(), client(an_identity(), 1), session()); + + assert_eq!(index.len(), 1); + } + + #[tokio::test] + async fn reconnect_is_refused_and_stops_the_holder() { + let index = index(); + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); + + let refused = index.try_reserve(a_database(), client(an_identity(), 2), session()); + + assert!(refused.is_err()); + assert!(first_sender.is_cancelled()); + assert_eq!(index.len(), 1); + } + + /// A client may repeat a connection id across connections, so the holder + /// is told apart by its name rather than by that id. + #[tokio::test] + async fn reconnect_reusing_connection_id_is_refused() { + let index = index(); + let first = ClientActorId { + name: ClientName(1), + ..client(an_identity(), 1) + }; + let second = ClientActorId { + name: ClientName(2), + ..client(an_identity(), 1) + }; + let (_first, first_sender) = connect(&index, a_database(), first, session()); + + assert!(index.try_reserve(a_database(), second, session()).is_err()); + assert!(first_sender.is_cancelled()); + } + + #[tokio::test] + async fn holder_without_a_sender_still_refuses() { + let index = index(); + let _first = index + .try_reserve(a_database(), client(an_identity(), 1), session()) + .unwrap(); + + assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); + assert_eq!(index.len(), 1); + } + + #[tokio::test] + async fn different_identity_does_not_conflict() { + let index = index(); + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); + + let (_second, _) = connect(&index, a_database(), client(another_identity(), 2), session()); + + assert!(!first_sender.is_cancelled()); + assert_eq!(index.len(), 2); + } + + #[tokio::test] + async fn different_session_does_not_conflict() { + let index = index(); + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); + + let (_second, _) = connect(&index, a_database(), client(an_identity(), 2), SessionId::from_u128(8)); + + assert!(!first_sender.is_cancelled()); + assert_eq!(index.len(), 2); + } + + #[tokio::test] + async fn different_database_does_not_conflict() { + let index = index(); + let (_first, first_sender) = connect(&index, a_database(), client(an_identity(), 1), session()); + + let (_second, _) = connect(&index, another_database(), client(an_identity(), 2), session()); + + assert!(!first_sender.is_cancelled()); + assert_eq!(index.len(), 2); + } + + #[tokio::test] + async fn dropping_the_reservation_frees_the_session() { + let index = index(); + let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session()); + assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err()); + + drop(first); + + assert!(index.is_empty()); + let retry = index.try_reserve(a_database(), client(an_identity(), 2), session()); + assert!(retry.is_ok()); + assert_eq!(index.len(), 1); + } + + /// A connection which never came to be, because `client_connected` + /// rejected it, frees the session without ever establishing it. + #[tokio::test] + async fn dropping_an_unestablished_reservation_frees_the_session() { + let index = index(); + let reservation = index + .try_reserve(a_database(), client(an_identity(), 1), session()) + .unwrap(); + + drop(reservation); + + assert!(index.is_empty()); + } +} 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..a0d5f45f717 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", } } @@ -276,6 +280,11 @@ metrics_group!( #[labels(database_identity: Identity)] pub ws_clients_idle_timed_out: IntCounterVec, + #[name = spacetime_worker_ws_clients_session_busy_total] + #[help = "The cumulative number of ws connections refused because their session was still held by a connection being torn down"] + #[labels(database_identity: Identity)] + pub ws_clients_session_busy: IntCounterVec, + // Compatibility counters above continue to be emitted for existing dashboards. // Accepted-client disconnection `cause` label values are: // client_close, idle_timeout, incoming_queue_full, outgoing_queue_full, diff --git a/crates/smoketests/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..be6ca0f48fe --- /dev/null +++ b/crates/smoketests/tests/cluster/connection_session.rs @@ -0,0 +1,508 @@ +//! 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 is +//! refused with `SESSION_BUSY_CLOSE_CODE` and the old one is torn down through +//! the normal disconnect sequence. A retry succeeds once the old connection's +//! `client_disconnected` has run. + +use std::time::{Duration, Instant}; + +use anyhow::{bail, Context, Result}; +use futures::{SinkExt, StreamExt}; +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. + /// + /// Retries while the server refuses the connection because its session + /// is still held by a connection being torn down. + async fn open(test: &Smoketest, connection_id: &str, session_id: Option<&str>) -> Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + if let Some(connection) = Self::open_once(test, connection_id, session_id).await? { + return Ok(connection); + } + if Instant::now() > deadline { + bail!("timed out retrying a connection refused as session busy"); + } + tokio::time::sleep(Duration::from_millis(100)).await; + } + } + + /// Open a connection once. Returns `None` if the server refused it with + /// `SESSION_BUSY_CLOSE_CODE`. + async fn open_once(test: &Smoketest, connection_id: &str, session_id: Option<&str>) -> Result> { + let token = test.read_token()?; + let host = test.server_host(); + let database = test + .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 { + Ok(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 + ); + } + Ok(Some(connection)) + } + Ok(other) => bail!("expected InitialConnection, got {other:?}"), + Err(err) if err.downcast_ref::().is_some_and(|closed| closed.is_session_busy()) => Ok(None), + Err(err) => Err(err), + } + } + + /// 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) => return Err(Closed(frame.map(|frame| frame.code.into())).into()), + 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 stopped connection's actor never answers a request. The server does + /// not send a close frame, so the peer's socket stays half-open until it + /// writes, which is what this does. + async fn is_still_served(&mut self) -> bool { + if self + .send(ws_v2::ClientMessage::Subscribe(ws_v2::Subscribe { + 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(_)) + ) + } +} + +/// The server closed the websocket, with the close code it sent, if any. +#[derive(Debug)] +struct Closed(Option); + +impl Closed { + fn is_session_busy(&self) -> bool { + self.0 == Some(ws_common::SESSION_BUSY_CLOSE_CODE) + } +} + +impl std::fmt::Display for Closed { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "websocket closed with code {:?}", self.0) + } +} + +impl std::error::Error for Closed {} + +fn runtime() -> tokio::runtime::Runtime { + tokio::runtime::Builder::new_current_thread() + .enable_all() + .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 is refused while the first is +/// still live, and the first is stopped. Its `client_disconnected` runs +/// strictly before the retried connection's `client_connected`. +#[test] +fn test_reconnect_with_same_session_replaces_connection() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + 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 refused = TestConnection::open_once(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("second connection failed"); + assert!( + refused.is_none(), + "the reconnect should be refused while the first connection is live" + ); + + assert!( + !first.is_still_served().await, + "the first connection should no longer be served" + ); + + let _second = TestConnection::open(&test, CONNECTION_B, Some(SESSION)) + .await + .expect("retried connection failed"); + + let lines = wait_for_log(&test, "connected", CONNECTION_B); + let connected_a = position_of(&lines, "connected", CONNECTION_A).expect("A never connected"); + let disconnected_a = position_of(&lines, "disconnected", CONNECTION_A).expect("A never disconnected"); + 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 is not refused: 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 refuses nor is stopped by a connection with a session id. +#[test] +fn test_connection_without_session_is_not_replaced() { + let test = Smoketest::builder().precompiled_module("connection-session").build(); + + 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 stopped: {lines:?}" + ); + }); +} + +/// Repeated reconnects each replace only the connection immediately before +/// them, leaving exactly one live connection for the session. +#[test] +fn test_repeated_reconnects_leave_one_live_connection() { + 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 stopped"); + + let mut third = TestConnection::open(&test, CONNECTION_C, Some(SESSION)) + .await + .expect("third connection failed"); + wait_for_log(&test, "connected", CONNECTION_C); + assert!(!second.is_still_served().await, "B should have been stopped"); + + let lines = lifecycle_log(&test); + assert!( + position_of(&lines, "disconnected", CONNECTION_B).is_some(), + "B should have been replaced 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 replaces +/// 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 replaced 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 {