Http handler disconnect fix.

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

@@ -1,4 +1,5 @@
 use std::borrow::Cow;
+use std::future::Future;
 use std::num::NonZeroU8;
 use std::str::FromStr;
 use std::time::Duration;
@@ -43,7 +44,7 @@ use spacetimedb_client_api_messages::name::{
 };
 use spacetimedb_lib::db::raw_def::v10::RawModuleDefV10;
 use spacetimedb_lib::db::raw_def::v9::RawModuleDefV9;
-use spacetimedb_lib::http as st_http;
+use spacetimedb_lib::{http as st_http, ConnectionId};
 use spacetimedb_lib::{sats, AlgebraicValue, Hash, ProductValue, Timestamp};
 use spacetimedb_schema::auto_migrate::{
     MigrationPolicy as SchemaMigrationPolicy, MigrationToken, PrettyPrintStyle as AutoMigratePrettyPrintStyle,
@@ -152,78 +153,67 @@ pub async fn call<S: ControlStateDelegate + NodeDelegate>(
 
     let args = FunctionArgs::Json(body);
 
-    // HTTP callers always need a connection ID to provide to connect/disconnect,
-    // so generate one.
-    let connection_id = generate_random_connection_id();
+    let caller_auth: ConnectionAuthCtx = auth.into();
 
     let (module, Database { owner_identity, .. }) = find_module_and_database(&worker_ctx, name_or_identity).await?;
 
-    // Call the database's `client_connected` reducer, if any.
-    // If it fails or rejects the connection, bail.
-    module
-        .call_identity_connected(auth.into(), connection_id)
-        .await
-        .map_err(client_connected_error_to_response)?;
-
-    let result = match module
-        .call_reducer(
-            caller_identity,
-            Some(connection_id),
-            None,
-            None,
-            None,
-            &reducer,
-            args.clone(),
-        )
-        .await
-    {
-        Ok(rcr) => Ok(CallResult::Reducer(rcr)),
-        Err(ReducerCallError::NoSuchReducer | ReducerCallError::ScheduleReducerNotFound) => {
-            // Not a reducer — try procedure instead
-            match module
-                .call_procedure(caller_identity, Some(connection_id), None, &reducer, args)
-                .await
-                .result
-            {
-                Ok(res) => Ok(CallResult::Procedure(res)),
-                Err(e) => Err(map_procedure_error(e, &reducer)),
+    let fut = async move |module: ModuleHost, caller_identity: Identity, connection_id: ConnectionId| {
+        let result = match module
+            .call_reducer(
+                caller_identity,
+                Some(connection_id),
+                None,
+                None,
+                None,
+                &reducer,
+                args.clone(),
+            )
+            .await
+        {
+            Ok(rcr) => Ok(CallResult::Reducer(rcr)),
+            Err(ReducerCallError::NoSuchReducer | ReducerCallError::ScheduleReducerNotFound) => {
+                // Not a reducer — try procedure instead
+                match module
+                    .call_procedure(caller_identity, Some(connection_id), None, &reducer, args)
+                    .await
+                    .result
+                {
+                    Ok(res) => Ok(CallResult::Procedure(res)),
+                    Err(e) => Err(map_procedure_error(e, &reducer)),
+                }
             }
+            Err(e) => Err(map_reducer_error(e, &reducer)),
+        };
+
+        match result {
+            Ok(CallResult::Reducer(result)) => {
+                let (status, body) = reducer_outcome_response(&owner_identity, &reducer, result.outcome);
+                Ok((
+                    status,
+                    TypedHeader(SpacetimeEnergyUsed(result.execution_budget_used)),
+                    TypedHeader(SpacetimeExecutionDurationMicros(result.execution_duration)),
+                    body,
+                )
+                    .into_response())
+            }
+            Ok(CallResult::Procedure(result)) => {
+                // Procedures don't assign a special meaning to error returns, unlike reducers,
+                // as there's no transaction for them to automatically abort.
+                // Instead, we just pass on their return value with the OK status so long as we successfully invoked the procedure.
+                let (status, body) = procedure_outcome_response(result.return_val);
+                Ok((
+                    status,
+                    TypedHeader(SpacetimeExecutionDurationMicros(result.execution_duration)),
+                    body,
+                )
+                    .into_response())
+            }
+            Err(e) => Err((e.0, e.1).into()),
         }
-        Err(e) => Err(map_reducer_error(e, &reducer)),
     };
 
-    module
-        .call_identity_disconnected(caller_identity, connection_id)
-        .await
-        .map_err(client_disconnected_error_to_response)?;
-
-    match result {
-        Ok(CallResult::Reducer(result)) => {
-            let (status, body) = reducer_outcome_response(&owner_identity, &reducer, result.outcome);
-            Ok((
-                status,
-                TypedHeader(SpacetimeEnergyUsed(result.execution_budget_used)),
-                TypedHeader(SpacetimeExecutionDurationMicros(result.execution_duration)),
-                body,
-            )
-                .into_response())
-        }
-        Ok(CallResult::Procedure(result)) => {
-            // Procedures don't assign a special meaning to error returns, unlike reducers,
-            // as there's no transaction for them to automatically abort.
-            // Instead, we just pass on their return value with the OK status so long as we successfully invoked the procedure.
-            let (status, body) = procedure_outcome_response(result.return_val);
-            Ok((
-                status,
-                TypedHeader(SpacetimeExecutionDurationMicros(result.execution_duration)),
-                body,
-            )
-                .into_response())
-        }
-        Err(e) => Err((e.0, e.1).into()),
-    }
+    with_connection(module, caller_auth, caller_identity, fut).await
 }
-
 #[derive(Deserialize)]
 pub struct HttpRouteRootParams {
     name_or_identity: NameOrIdentity,
@@ -709,6 +699,46 @@ pub struct SqlQueryParams {
     pub confirmed: Option<bool>,
 }
 
+/// Runs `fut` inside an HTTP client connection lifecycle.
+///
+/// The lifecycle is detached from the request task, so once `client_connected`
+/// succeeds, `client_disconnected` still runs if the handler is dropped or
+/// `fut` returns an error.
+async fn with_connection<F, Fut, R>(
+    module: ModuleHost,
+    caller_auth: ConnectionAuthCtx,
+    caller_identity: Identity,
+    fut: F,
+) -> axum::response::Result<R>
+where
+    F: FnOnce(ModuleHost, Identity, ConnectionId) -> Fut + Send + 'static,
+    Fut: Future<Output = axum::response::Result<R>> + Send + 'static,
+    R: Send + 'static,
+{
+    tokio::spawn(async move {
+        let connection_id = generate_random_connection_id();
+
+        // Run the module's client_connected reducer, if any.
+        // If it rejects the connection, bail before executing `fut`
+        module
+            .call_identity_connected(caller_auth, connection_id)
+            .await
+            .map_err(client_connected_error_to_response)?;
+
+        let result = fut(module.clone(), caller_identity, connection_id).await;
+
+        // Always disconnect, even if authorization or execution failed.
+        module
+            .call_identity_disconnected(caller_identity, connection_id)
+            .await
+            .map_err(client_disconnected_error_to_response)?;
+
+        result
+    })
+    .await
+    .map_err(log_and_500)?
+}
+
 pub async fn sql_direct<S>(
     worker_ctx: S,
     SqlParams { name_or_identity }: SqlParams,
@@ -718,21 +748,12 @@ pub async fn sql_direct<S>(
     sql: String,
 ) -> axum::response::Result<Vec<SqlStmtResult<ProductValue>>>
 where
-    S: NodeDelegate + ControlStateDelegate + Authorization,
+    S: NodeDelegate + ControlStateDelegate + Authorization + 'static,
 {
-    let connection_id = generate_random_connection_id();
-
     let (host, database) = find_leader_and_database(&worker_ctx, name_or_identity).await?;
 
-    // Run the module's client_connected reducer, if any.
-    // If it rejects the connection, bail before executing SQL.
     let module = host.module().await.map_err(log_and_500)?;
-    module
-        .call_identity_connected(caller_auth, connection_id)
-        .await
-        .map_err(client_connected_error_to_response)?;
-
-    let result = async {
+    let fut = async move |_module: ModuleHost, caller_identity: Identity, _connection_id: ConnectionId| {
         let sql_auth = worker_ctx
             .authorize_sql(caller_identity, database.database_identity)
             .await?;
@@ -744,16 +765,9 @@ where
             sql,
         )
         .await
-    }
-    .await;
-
-    // Always disconnect, even if authorization or execution failed.
-    module
-        .call_identity_disconnected(caller_identity, connection_id)
-        .await
-        .map_err(client_disconnected_error_to_response)?;
+    };
 
-    result
+    with_connection(module, caller_auth, caller_identity, fut).await
 }
 
 pub async fn sql<S>(
@@ -764,7 +778,7 @@ pub async fn sql<S>(
     body: String,
 ) -> axum::response::Result<impl IntoResponse>
 where
-    S: NodeDelegate + ControlStateDelegate + Authorization,
+    S: NodeDelegate + ControlStateDelegate + Authorization + 'static,
 {
     let caller_identity = auth.claims.identity;
     let caller_auth: ConnectionAuthCtx = auth.into();

crates/pg/src/pg_server.rs

@@ -150,7 +150,7 @@ struct PgSpacetimeDB<T> {
 
 impl<T> PgSpacetimeDB<T>
 where
-    T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone,
+    T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static,
 {
     async fn exe_sql(&self, query: String) -> PgWireResult<Vec<Response>> {
         let params = self.cached.lock().await.clone().unwrap();
@@ -314,7 +314,7 @@ impl<T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDel
 #[async_trait]
 impl<T> SimpleQueryHandler for PgSpacetimeDB<T>
 where
-    T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone,
+    T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static,
 {
     async fn do_query<C>(&self, _client: &mut C, query: &str) -> PgWireResult<Vec<Response>>
     where
@@ -347,7 +347,7 @@ impl<T> PgSpacetimeDBFactory<T> {
 
 impl<T> PgWireServerHandlers for PgSpacetimeDBFactory<T>
 where
-    T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone,
+    T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static,
 {
     fn simple_query_handler(&self) -> Arc<impl SimpleQueryHandler> {
         self.handler.clone()

crates/smoketests/Cargo.toml

@@ -21,6 +21,7 @@ spacetimedb-core.workspace = true
 cargo_metadata.workspace = true
 assert_cmd = "2"
 predicates = "3"
+socket2.workspace = true
 tokio.workspace = true
 tokio-postgres.workspace = true
 xmltree.workspace = true

crates/smoketests/modules/Cargo.toml

@@ -82,6 +82,7 @@ members = [
     "delete-database",
     "client-connection-reject",
     "client-connection-disconnect-panic",
+    "client-connection-http-cancel",
     "sql-connect-hook",
 
     # Log filtering tests

crates/smoketests/modules/client-connection-http-cancel/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "smoketest-module-client-connection-http-cancel"
+version = "0.1.0"
+edition = "2021"
+publish = false
+
+[lib]
+crate-type = ["cdylib"]
+
+[dependencies]
+spacetimedb.workspace = true
+log.workspace = true

crates/smoketests/modules/client-connection-http-cancel/src/lib.rs

@@ -0,0 +1,25 @@
+use spacetimedb::{log, ReducerContext};
+
+#[spacetimedb::table(accessor = marker, public)]
+pub struct Marker {
+    id: u32,
+}
+
+#[spacetimedb::reducer(client_connected)]
+pub fn connected(ctx: &ReducerContext) {
+    log::info!("http_cancel_repro: connected {}", ctx.sender());
+}
+
+#[spacetimedb::reducer(client_disconnected)]
+pub fn disconnected(ctx: &ReducerContext) {
+    log::info!("http_cancel_repro: disconnected {}", ctx.sender());
+}
+
+#[spacetimedb::reducer]
+pub fn slow(_ctx: &ReducerContext) {
+    log::info!("http_cancel_repro: slow reducer started");
+    for i in 0..300_000_000u64 {
+        core::hint::black_box(i);
+    }
+    log::info!("http_cancel_repro: slow reducer finished");
+}