Alessandro/auto reconnect server side

crates/client-api-messages/src/websocket/v2.rs

@@ -26,6 +26,8 @@ pub enum ClientMessage {
     CallReducer(CallReducer),
     /// Invoke a procedure, a non-transactional side-effecting function which runs in the database.
     CallProcedure(CallProcedure),
+    /// Add multiple sets of subscribed queries in one atomic step.
+    SubscribeBatch(SubscribeBatch),
 }
 
 /// Sent by client to register a subscription to a new query set
@@ -92,6 +94,42 @@ pub enum UnsubscribeFlags {
     SendDroppedRows = 1,
 }
 
+/// Sent by client to register multiple subscriptions in one atomic step.
+///
+/// The server registers every subscription set under a single subscription-manager
+/// lock and evaluates all of them at a single transaction snapshot,
+/// then responds with one [`SubscribeBatchApplied`] message carrying a result per set.
+/// No [`TransactionUpdate`] is delivered between the registration of the first set
+/// and the [`SubscribeBatchApplied`] response,
+/// and updates for the new sets resume after it.
+///
+/// A set whose queries are invalid or fail to compute reports an error in its
+/// [`SubscribeSetResult`]. The remaining sets still apply.
+#[derive(SpacetimeType)]
+#[sats(crate = spacetimedb_lib)]
+pub struct SubscribeBatch {
+    /// An identifier for a client request.
+    pub request_id: u32,
+
+    /// The subscription sets to register.
+    ///
+    /// Each [`QuerySetId`] must be distinct,
+    /// and must not be used by any other subscription on the same connection.
+    pub sets: Box<[SubscribeSet]>,
+}
+
+/// One subscription set within a [`SubscribeBatch`].
+#[derive(SpacetimeType)]
+#[sats(crate = spacetimedb_lib)]
+pub struct SubscribeSet {
+    /// An identifier for this subscription,
+    /// which should not be used for any other subscriptions on the same connection.
+    pub query_set_id: QuerySetId,
+
+    /// A set of queries to subscribe to, each a single SQL `SELECT` statement.
+    pub query_strings: Box<[Box<str>]>,
+}
+
 /// 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<str>,
 }
 
+/// Response to [`SubscribeBatch`], carrying one result per registered query set.
+///
+/// This message's `request_id` matches the one the client provided in the [`SubscribeBatch`] message,
+/// and `results` contains exactly one entry per received [`SubscribeSet`], in the same order.
+///
+/// Every applied set's rows are evaluated at a single transaction snapshot.
+#[derive(SpacetimeType, Debug)]
+#[sats(crate = spacetimedb_lib)]
+pub struct SubscribeBatchApplied {
+    /// The request_id of the corresponding [`SubscribeBatch`] message.
+    pub request_id: u32,
+    /// One result per query set, in the order the sets appeared in the request.
+    pub results: Box<[SubscribeSetResult]>,
+}
+
+/// The result for one query set within a [`SubscribeBatchApplied`].
+#[derive(SpacetimeType, Debug)]
+#[sats(crate = spacetimedb_lib)]
+pub struct SubscribeSetResult {
+    /// The [`QuerySetId`] the client provided for this set.
+    pub query_set_id: QuerySetId,
+    /// The outcome for this set.
+    pub outcome: SubscribeSetOutcome,
+}
+
+/// The outcome for one query set within a [`SubscribeBatchApplied`].
+#[derive(SpacetimeType, Debug)]
+#[sats(crate = spacetimedb_lib)]
+pub enum SubscribeSetOutcome {
+    /// The set was applied; contains its initial matching rows.
+    /// The set behaves like one registered with an individual [`Subscribe`] afterwards.
+    Applied(QueryRows),
+    /// The set failed to compile or compute.
+    /// The set is not registered; its [`QuerySetId`] may be re-used.
+    /// The error string follows the conventions of [`SubscriptionError`]'s `error` field.
+    Error(Box<str>),
+}
+
 /// 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.

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<ConnectionIdForUrl>,
+    /// 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<SessionIdForUrl>,
     #[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<bool>,
 }
 
+/// 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<D: serde::Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
+        let hex = <std::borrow::Cow<'de, str>>::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<SessionIdForUrl> for SessionId {
+    fn from(session_id: SessionIdForUrl) -> Self {
+        session_id.0
+    }
+}
+
 fn resolve_confirmed_reads_default(version: WsVersion, confirmed: Option<bool>) -> bool {
     if let Some(confirmed) = confirmed {
         return confirmed;
@@ -119,6 +149,7 @@ pub async fn handle_websocket<S>(
     Path(SubscribeParams { name_or_identity }): Path<SubscribeParams>,
     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<SessionId> = session_id.map(Into::into);
+    let sessions = ctx.client_actor_index().sessions();
 
     let ws_config = WebSocketConfig::default()
         .max_message_size(Some(0x2000000))
@@ -255,6 +288,29 @@ where
 
         log::debug!("websocket: New client connected from {client_log_string}");
 
+        // If this connection resumes a session which a live connection still
+        // holds, that connection is taken over. Stop its actor and run its
+        // module-side disconnect to completion, so the module observes
+        // `client_disconnected` for it strictly before `client_connected` for
+        // this one, and never two live connections for one session.
+        if let Some(session_id) = session_id
+            && let Some(superseded) = sessions.claim_session(client_id, session_id, None)
+        {
+            log::debug!(
+                "websocket: Connection {} supersedes {} for session {session_id}",
+                client_id.connection_id,
+                superseded.client_id.connection_id,
+            );
+            if let Some(sender) = &superseded.sender {
+                sender.kick(ClientDisconnectCause::ConnectionSuperseded);
+            }
+            // Awaiting this is what orders the two lifecycle reducers:
+            // both run on the module's main instance, and this one is
+            // enqueued first.
+            let module = module_rx.borrow().clone();
+            module.disconnect_client(superseded.client_id).await;
+        }
+
         let connected = match ClientConnection::call_client_connected_maybe_reject(
             &mut module_rx,
             client_id,
@@ -284,6 +340,11 @@ where
                     }
                 };
                 record_client_rejection(db_identity, cause);
+                // The session claim is only meaningful for a connection which
+                // exists, so give it up again.
+                if let Some(session_id) = session_id {
+                    sessions.release_session(client_id, session_id);
+                }
                 return;
             }
         };
@@ -292,7 +353,16 @@ where
             "websocket: Database accepted connection from {client_log_string}; spawning ws_client_actor and ClientConnection"
         );
 
-        let actor = |client, receiver| ws_client_actor(ws_opts, client, ws, receiver);
+        // Release the session claim when the actor ends, including when it is
+        // aborted: dropping the future drops the guard.
+        let session_guard = session_id.map(|session_id| {
+            let sessions = sessions.clone();
+            scopeguard::guard((), move |()| sessions.release_session(client_id, session_id))
+        });
+        let actor = |client, receiver| async move {
+            let _session_guard = session_guard;
+            ws_client_actor(ws_opts, client, ws, receiver).await;
+        };
         let client = ClientConnection::spawn(
             client_id,
             auth.into(),
@@ -305,6 +375,12 @@ where
         )
         .await;
 
+        // Now that the actor exists, register its sender so that a later
+        // connection resuming this session can stop it.
+        if let Some(session_id) = session_id {
+            sessions.attach_sender(client_id, session_id, &client.sender());
+        }
+
         // Send the client their identity token message as the first message
         // NOTE: We're adding this to the protocol because some client libraries are
         // unable to access the http response headers.

crates/core/src/client.rs

@@ -3,6 +3,7 @@ use std::fmt;
 
 mod client_connection;
 mod client_connection_index;
+mod client_session_index;
 pub mod consume_each_list;
 mod message_handlers;
 mod message_handlers_v1;
@@ -16,6 +17,7 @@ pub use client_connection::{
     WsVersion,
 };
 pub use client_connection_index::ClientActorIndex;
+pub use client_session_index::{ClientSessionIndex, SessionId, SupersededConnection};
 pub use message_handlers::MessageHandleError;
 pub use message_handlers_v1::MessageExecutionError;
 pub use messages::OutboundMessage;

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<Option<ExecutionMetrics>, 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,