Add a drop safe async version of scope guard

crates/client-api/src/routes/subscribe.rs

@@ -20,7 +20,7 @@ use derive_more::From;
 use futures::{pin_mut, Sink, SinkExt, Stream, StreamExt};
 use http::{HeaderValue, StatusCode};
 use prometheus::{Histogram, IntGauge};
-use scopeguard::{defer, ScopeGuard};
+use scopeguard::defer;
 use serde::Deserialize;
 use spacetimedb::client::messages::{
     serialize, serialize_v3, IdentityTokenMessage, InUseSerializeBuffer, SerializeBuffer, SwitchedServerMessage,
@@ -58,7 +58,7 @@ use crate::util::serde::humantime_duration;
 use crate::util::websocket::{
     CloseCode, CloseFrame, Message as WsMessage, WebSocketConfig, WebSocketStream, WebSocketUpgrade, WsError,
 };
-use crate::util::{NameOrIdentity, XForwardedFor};
+use crate::util::{async_cleanup_guard, NameOrIdentity, XForwardedFor};
 use crate::{log_and_500, Authorization, ControlStateDelegate, NodeDelegate};
 
 #[allow(clippy::declare_interior_mutable_const)]
@@ -298,7 +298,10 @@ where
             Some(session_id) => match sessions.try_reserve(db_identity, client_id, session_id) {
                 Ok(reservation) => Some(reservation),
                 Err(SessionBusy) => {
-                    WORKER_METRICS.ws_clients_session_busy.with_label_values(&db_identity).inc();
+                    WORKER_METRICS
+                        .ws_clients_session_busy
+                        .with_label_values(&db_identity)
+                        .inc();
                     log::debug!("websocket: Refusing connection for {client_log_string}: session {session_id} is busy");
                     let close = CloseFrame {
                         code: CloseCode::from(SESSION_BUSY_CLOSE_CODE),
@@ -575,14 +578,15 @@ async fn ws_client_actor(
     session: Option<SessionReservation>,
 ) {
     // Runs the module-side disconnect even if this task gets cancelled.
-    let mut client = scopeguard::guard((client, session), |(client, session)| {
-        tokio::spawn(ws_client_teardown(client, session));
+    let mut client = async_cleanup_guard((client, session), |(client, session)| {
+        ws_client_teardown(client, session)
     });
 
     ws_client_actor_inner(&mut client.0, options, ws, sendrx).await;
 
-    let (client, session) = ScopeGuard::into_inner(client);
-    ws_client_teardown(client, session).await;
+    if let Err(e) = client.cleanup().await {
+        log::error!("websocket client teardown task failed: {e}");
+    }
 }
 
 /// Run the module-side disconnect, then free the connection's session.

crates/client-api/src/util.rs

@@ -3,7 +3,10 @@ pub(crate) mod serde;
 pub mod websocket;
 
 use core::fmt;
+use std::future::Future;
+use std::marker::PhantomData;
 use std::net::IpAddr;
+use std::ops::{Deref, DerefMut};
 
 use axum::body::Bytes;
 use axum::extract::{FromRequest, Request};
@@ -15,10 +18,117 @@ use http::{HeaderName, HeaderValue, StatusCode};
 use hyper::body::Body;
 use spacetimedb::Identity;
 use spacetimedb_client_api_messages::name::DatabaseName;
+use tokio::task::{JoinError, JoinHandle};
 
 use crate::routes::identity::IdentityForUrl;
 use crate::{log_and_500, ControlStateReadAccess};
 
+/// Returns a guard that runs async cleanup for `value` when dropped.
+///
+/// This is cancel-safe with respect to cancellation of the task holding the
+/// guard: dropping the guard spawns the cleanup future in its own task instead
+/// of trying to run async cleanup from `Drop`.
+///
+/// This does not guarantee that cleanup survives shutdown of the Tokio runtime,
+/// process exit, or explicit abortion of the spawned cleanup task.
+///
+/// Dropping this guard calls [`tokio::spawn`], so it must be dropped from
+/// within a Tokio runtime.
+pub(crate) fn async_cleanup_guard<T, F, Fut>(value: T, cleanup: F) -> AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    AsyncCleanupGuard {
+        value: Some(value),
+        cleanup: Some(cleanup),
+        future: PhantomData,
+    }
+}
+
+/// Scope guard for values that require async cleanup.
+///
+/// Dropping the guard is cancel-safe for the guarded task: it moves the guarded
+/// value into a newly spawned cleanup task. Drop does not wait for cleanup to
+/// complete.
+///
+/// Call [`Self::cleanup`] on the normal path when the current task should wait
+/// for cleanup. That method starts cleanup in a spawned task before awaiting it,
+/// so cancelling the waiter does not cancel the cleanup task.
+pub(crate) struct AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    value: Option<T>,
+    cleanup: Option<F>,
+    future: PhantomData<fn() -> Fut>,
+}
+
+impl<T, F, Fut> AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    fn spawn_cleanup(&mut self) -> JoinHandle<()> {
+        let value = self.value.take().expect("cleanup value already taken");
+        let cleanup = self.cleanup.take().expect("cleanup function already taken");
+        tokio::spawn(cleanup(value))
+    }
+
+    /// Starts cleanup and waits for the cleanup task to finish.
+    ///
+    /// This is cancel-safe with respect to cancellation of the caller: cleanup
+    /// is spawned before this method awaits, so dropping this future after its
+    /// first poll drops only the wait for completion, not the cleanup itself.
+    ///
+    /// This is not cancel-safe against explicit abortion of the returned
+    /// cleanup task by the runtime or against runtime shutdown.
+    pub(crate) async fn cleanup(mut self) -> Result<(), JoinError> {
+        self.spawn_cleanup().await
+    }
+}
+
+impl<T, F, Fut> Deref for AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    type Target = T;
+
+    fn deref(&self) -> &Self::Target {
+        self.value.as_ref().expect("cleanup value already taken")
+    }
+}
+
+impl<T, F, Fut> DerefMut for AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        self.value.as_mut().expect("cleanup value already taken")
+    }
+}
+
+impl<T, F, Fut> Drop for AsyncCleanupGuard<T, F, Fut>
+where
+    T: Send + 'static,
+    F: FnOnce(T) -> Fut + Send + 'static,
+    Fut: Future<Output = ()> + Send + 'static,
+{
+    fn drop(&mut self) {
+        if let (Some(value), Some(cleanup)) = (self.value.take(), self.cleanup.take()) {
+            tokio::spawn(cleanup(value));
+        }
+    }
+}
+
 pub struct ByteStringBody(pub ByteString);
 
 #[async_trait::async_trait]
@@ -192,6 +302,9 @@ impl<S> FromRequest<S> for EmptyBody {
 mod tests {
     use super::*;
     use headers::Header;
+    use std::sync::atomic::{AtomicBool, Ordering};
+    use std::sync::Arc;
+    use tokio::sync::{oneshot, Notify};
 
     fn decode_one(raw: &str) -> Result<XForwardedFor, headers::Error> {
         let val = HeaderValue::from_str(raw).unwrap();
@@ -224,4 +337,76 @@ mod tests {
         assert!(decode_one("not-an-ip").is_err());
         assert!(decode_one("not-an-ip, 10.0.0.1").is_err());
     }
+
+    #[tokio::test]
+    async fn async_cleanup_guard_runs_cleanup_when_dropped() {
+        let (tx, rx) = oneshot::channel();
+
+        drop(async_cleanup_guard((), move |()| async move {
+            tx.send(()).unwrap();
+        }));
+
+        rx.await.expect("cleanup should run");
+    }
+
+    #[tokio::test]
+    async fn async_cleanup_guard_cleanup_waits_for_cleanup() {
+        let cleaned_up = Arc::new(AtomicBool::new(false));
+        let cleanup_started = Arc::new(Notify::new());
+        let finish_cleanup = Arc::new(Notify::new());
+        let cleaned_up_for_guard = Arc::clone(&cleaned_up);
+        let cleanup_started_for_guard = Arc::clone(&cleanup_started);
+        let finish_cleanup_for_guard = Arc::clone(&finish_cleanup);
+
+        let cleanup = tokio::spawn(
+            async_cleanup_guard((), move |()| async move {
+                cleanup_started_for_guard.notify_one();
+                finish_cleanup_for_guard.notified().await;
+                cleaned_up_for_guard.store(true, Ordering::Release);
+            })
+            .cleanup(),
+        );
+
+        cleanup_started.notified().await;
+        assert!(!cleaned_up.load(Ordering::Acquire));
+
+        finish_cleanup.notify_one();
+        cleanup
+            .await
+            .expect("cleanup join task should not panic")
+            .expect("cleanup task should not panic");
+        assert!(cleaned_up.load(Ordering::Acquire));
+    }
+
+    #[tokio::test]
+    async fn async_cleanup_guard_cleanup_continues_if_waiter_is_aborted() {
+        let cleaned_up = Arc::new(AtomicBool::new(false));
+        let cleanup_started = Arc::new(Notify::new());
+        let finish_cleanup = Arc::new(Notify::new());
+        let cleaned_up_for_guard = Arc::clone(&cleaned_up);
+        let cleanup_started_for_guard = Arc::clone(&cleanup_started);
+        let finish_cleanup_for_guard = Arc::clone(&finish_cleanup);
+
+        let cleanup = tokio::spawn(
+            async_cleanup_guard((), move |()| async move {
+                cleanup_started_for_guard.notify_one();
+                finish_cleanup_for_guard.notified().await;
+                cleaned_up_for_guard.store(true, Ordering::Release);
+            })
+            .cleanup(),
+        );
+
+        cleanup_started.notified().await;
+        cleanup.abort();
+        assert!(cleanup.await.unwrap_err().is_cancelled());
+
+        finish_cleanup.notify_one();
+        tokio::time::timeout(std::time::Duration::from_secs(1), async {
+            while !cleaned_up.load(Ordering::Acquire) {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("cleanup should continue after waiter abort");
+    }
 }

crates/core/src/client/client_session_index.rs

@@ -129,7 +129,10 @@ impl ClientSessionIndex {
         let mut sessions = self.sessions.lock().expect("session index poisoned");
         // Connections are told apart by their name, the host's per-connection
         // counter, rather than by their connection id, which a client may repeat.
-        if sessions.get(&key).is_some_and(|holder| holder.client.name == client.name) {
+        if sessions
+            .get(&key)
+            .is_some_and(|holder| holder.client.name == client.name)
+        {
             sessions.remove(&key);
         }
     }
@@ -238,7 +241,9 @@ mod tests {
         client: ClientActorId,
         session: SessionId,
     ) -> (SessionReservation, Arc<ClientConnectionSender>) {
-        let reservation = index.try_reserve(database, client, session).expect("session should be free");
+        let reservation = index
+            .try_reserve(database, client, session)
+            .expect("session should be free");
         let sender = sender(client);
         reservation.establish(&sender);
         (reservation, sender)
@@ -291,7 +296,9 @@ mod tests {
             .try_reserve(a_database(), client(an_identity(), 1), session())
             .unwrap();
 
-        assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err());
+        assert!(index
+            .try_reserve(a_database(), client(an_identity(), 2), session())
+            .is_err());
         assert_eq!(index.len(), 1);
     }
 
@@ -332,7 +339,9 @@ mod tests {
     async fn dropping_the_reservation_frees_the_session() {
         let index = index();
         let (first, _) = connect(&index, a_database(), client(an_identity(), 1), session());
-        assert!(index.try_reserve(a_database(), client(an_identity(), 2), session()).is_err());
+        assert!(index
+            .try_reserve(a_database(), client(an_identity(), 2), session())
+            .is_err());
 
         drop(first);
 

crates/smoketests/tests/cluster/connection_session.rs

@@ -97,7 +97,13 @@ impl TestConnection {
                 Ok(Some(connection))
             }
             Ok(other) => bail!("expected InitialConnection, got {other:?}"),
-            Err(err) if err.downcast_ref::<Closed>().is_some_and(|closed| closed.is_session_busy()) => Ok(None),
+            Err(err)
+                if err
+                    .downcast_ref::<Closed>()
+                    .is_some_and(|closed| closed.is_session_busy()) =>
+            {
+                Ok(None)
+            }
             Err(err) => Err(err),
         }
     }