diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index e1293a3997a..462bace2d2f 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/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, ) { // 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. diff --git a/crates/client-api/src/util.rs b/crates/client-api/src/util.rs index 0cc87a82bfd..66e5fc91158 100644 --- a/crates/client-api/src/util.rs +++ b/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(value: T, cleanup: F) -> AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + 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 +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + value: Option, + cleanup: Option, + future: PhantomData Fut>, +} + +impl AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + 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 Deref for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + type Target = T; + + fn deref(&self) -> &Self::Target { + self.value.as_ref().expect("cleanup value already taken") + } +} + +impl DerefMut for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + Send + 'static, +{ + fn deref_mut(&mut self) -> &mut Self::Target { + self.value.as_mut().expect("cleanup value already taken") + } +} + +impl Drop for AsyncCleanupGuard +where + T: Send + 'static, + F: FnOnce(T) -> Fut + Send + 'static, + Fut: Future + 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 FromRequest 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 { 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"); + } } diff --git a/crates/core/src/client/client_session_index.rs b/crates/core/src/client/client_session_index.rs index 6d968d774f2..f958ebb9a34 100644 --- a/crates/core/src/client/client_session_index.rs +++ b/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) { - 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); diff --git a/crates/smoketests/tests/cluster/connection_session.rs b/crates/smoketests/tests/cluster/connection_session.rs index be6ca0f48fe..b7ea3758a1d 100644 --- a/crates/smoketests/tests/cluster/connection_session.rs +++ b/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::().is_some_and(|closed| closed.is_session_busy()) => Ok(None), + Err(err) + if err + .downcast_ref::() + .is_some_and(|closed| closed.is_session_busy()) => + { + Ok(None) + } Err(err) => Err(err), } }