From 7731764f63c415c40c3430d1a40fcda1ae6c44f8 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Wed, 8 Jul 2026 12:01:26 +0530 Subject: [PATCH 1/4] http disconnect cancel smoketest --- crates/smoketests/modules/Cargo.lock | 8 ++ crates/smoketests/modules/Cargo.toml | 1 + .../client-connection-http-cancel/Cargo.toml | 12 ++ .../client-connection-http-cancel/src/lib.rs | 30 +++++ .../smoketests/client_connection_errors.rs | 119 +++++++++++++++++- 5 files changed, 169 insertions(+), 1 deletion(-) create mode 100644 crates/smoketests/modules/client-connection-http-cancel/Cargo.toml create mode 100644 crates/smoketests/modules/client-connection-http-cancel/src/lib.rs diff --git a/crates/smoketests/modules/Cargo.lock b/crates/smoketests/modules/Cargo.lock index 2ce4b735807..e7adacae94a 100644 --- a/crates/smoketests/modules/Cargo.lock +++ b/crates/smoketests/modules/Cargo.lock @@ -609,6 +609,14 @@ dependencies = [ "spacetimedb", ] +[[package]] +name = "smoketest-module-client-connection-http-cancel" +version = "0.1.0" +dependencies = [ + "log", + "spacetimedb", +] + [[package]] name = "smoketest-module-client-connection-reject" version = "0.1.0" diff --git a/crates/smoketests/modules/Cargo.toml b/crates/smoketests/modules/Cargo.toml index c23de0030f6..c5a3faffebb 100644 --- a/crates/smoketests/modules/Cargo.toml +++ b/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 diff --git a/crates/smoketests/modules/client-connection-http-cancel/Cargo.toml b/crates/smoketests/modules/client-connection-http-cancel/Cargo.toml new file mode 100644 index 00000000000..46ecd40b509 --- /dev/null +++ b/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 diff --git a/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs b/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs new file mode 100644 index 00000000000..be0548aa6be --- /dev/null +++ b/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs @@ -0,0 +1,30 @@ +use spacetimedb::{log, ReducerContext, Table}; + +#[spacetimedb::table(accessor = marker, public)] +pub struct Marker { + id: u32, +} + +#[spacetimedb::reducer(init)] +pub fn init(ctx: &ReducerContext) { + ctx.db.marker().insert(Marker { id: 0 }); +} + +#[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"); +} diff --git a/crates/smoketests/tests/smoketests/client_connection_errors.rs b/crates/smoketests/tests/smoketests/client_connection_errors.rs index c937d750ce4..baeb5172225 100644 --- a/crates/smoketests/tests/smoketests/client_connection_errors.rs +++ b/crates/smoketests/tests/smoketests/client_connection_errors.rs @@ -1,4 +1,8 @@ -use spacetimedb_smoketests::Smoketest; +use spacetimedb_smoketests::{require_local_server, Smoketest}; +use std::io::Write; +use std::net::TcpStream; +use std::thread; +use std::time::{Duration, Instant}; /// Test that client_connected returning an error rejects the connection #[test] @@ -57,3 +61,116 @@ fn test_client_disconnected_error_still_deletes_st_client() { "Expected at most 1 st_client row (the SQL connection itself), got {row_count}: {sql_out}", ); } + +#[test] +fn test_http_reducer_call_cancel_still_deletes_st_client() { + require_local_server!(); + + let test = Smoketest::builder() + .precompiled_module("client-connection-http-cancel") + .build(); + + let streams = send_http_reducer_calls_and_hold_connections(&test, "slow", 64); + wait_for_log( + &test, + "http_cancel_repro: slow reducer started", + Duration::from_secs(30), + ); + drop(streams); + + wait_for_log( + &test, + "http_cancel_repro: slow reducer finished", + Duration::from_secs(30), + ); + + let (st_client_count, st_connection_credentials_count) = wait_for_client_rows_to_clear(&test); + assert!( + st_client_count <= 1, + "Expected at most 1 st_client row (the SQL connection itself) after dropping 64 HTTP calls post-connect, got {st_client_count}", + ); + assert!( + st_connection_credentials_count <= 1, + "Expected at most 1 st_connection_credentials row (the SQL connection itself) after dropping 64 HTTP calls post-connect, got {st_connection_credentials_count}", + ); +} + +fn send_http_reducer_calls_and_hold_connections(test: &Smoketest, reducer: &str, count: usize) -> Vec { + let host = test.server_host(); + let request = http_reducer_call_request(test, reducer); + + (0..count) + .map(|_| { + let mut stream = TcpStream::connect(host).expect("Failed to connect to local server"); + stream.set_nodelay(true).expect("Failed to set TCP_NODELAY"); + stream + .write_all(request.as_bytes()) + .expect("Failed to write HTTP request"); + stream.flush().expect("Failed to flush HTTP request"); + stream + }) + .collect() +} + +fn http_reducer_call_request(test: &Smoketest, reducer: &str) -> String { + let identity = test.database_identity.as_ref().expect("No database published"); + let host = test.server_host(); + let token = test.read_token().expect("Failed to read auth token"); + let path = format!("/v1/database/{identity}/call/{reducer}"); + + format!( + "POST {path} HTTP/1.1\r\n\ + Host: {host}\r\n\ + Authorization: Bearer {token}\r\n\ + Content-Type: application/json\r\n\ + Content-Length: 2\r\n\ + Connection: close\r\n\ + \r\n\ + []" + ) +} + +fn wait_for_log(test: &Smoketest, needle: &str, timeout: Duration) -> Vec { + let start = Instant::now(); + loop { + let logs = test.logs(200).unwrap(); + if logs.iter().any(|l| l.contains(needle)) { + return logs; + } + assert!( + start.elapsed() < timeout, + "Timed out waiting for log containing {needle:?}: {:?}", + logs + ); + thread::sleep(Duration::from_millis(100)); + } +} + +fn wait_for_client_rows_to_clear(test: &Smoketest) -> (usize, usize) { + let start = Instant::now(); + let mut last = (usize::MAX, usize::MAX); + + while start.elapsed() < Duration::from_secs(20) { + last = ( + sql_count(test, "st_client"), + sql_count(test, "st_connection_credentials"), + ); + if last.0 <= 1 && last.1 <= 1 { + return last; + } + thread::sleep(Duration::from_millis(250)); + } + + last +} + +fn sql_count(test: &Smoketest, table: &str) -> usize { + let output = test + .sql(&format!("SELECT COUNT(*) AS count FROM {table}")) + .unwrap_or_else(|err| panic!("Failed to count rows in {table}: {err:#}")); + + output + .lines() + .find_map(|line| line.trim().parse::().ok()) + .unwrap_or_else(|| panic!("Failed to parse COUNT(*) output for {table}: {output}")) +} From bd27fd81c90f6017f29e89825a9a7b9df779f943 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Wed, 8 Jul 2026 17:51:21 +0530 Subject: [PATCH 2/4] make smoketest simpler for single client to hide false positives --- Cargo.lock | 1 + crates/smoketests/Cargo.toml | 1 + .../client-connection-http-cancel/src/lib.rs | 7 +-- .../smoketests/client_connection_errors.rs | 46 +++++++++++-------- 4 files changed, 29 insertions(+), 26 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 40ddf5ea619..dd6f6c6ca43 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8850,6 +8850,7 @@ dependencies = [ "regex", "reqwest", "serde_json", + "socket2 0.5.10", "spacetimedb-core", "spacetimedb-guard", "tempfile", diff --git a/crates/smoketests/Cargo.toml b/crates/smoketests/Cargo.toml index 473be886c55..fa2461b13c5 100644 --- a/crates/smoketests/Cargo.toml +++ b/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 diff --git a/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs b/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs index be0548aa6be..1830b7f862d 100644 --- a/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs +++ b/crates/smoketests/modules/client-connection-http-cancel/src/lib.rs @@ -1,15 +1,10 @@ -use spacetimedb::{log, ReducerContext, Table}; +use spacetimedb::{log, ReducerContext}; #[spacetimedb::table(accessor = marker, public)] pub struct Marker { id: u32, } -#[spacetimedb::reducer(init)] -pub fn init(ctx: &ReducerContext) { - ctx.db.marker().insert(Marker { id: 0 }); -} - #[spacetimedb::reducer(client_connected)] pub fn connected(ctx: &ReducerContext) { log::info!("http_cancel_repro: connected {}", ctx.sender()); diff --git a/crates/smoketests/tests/smoketests/client_connection_errors.rs b/crates/smoketests/tests/smoketests/client_connection_errors.rs index baeb5172225..ec1fb641cb6 100644 --- a/crates/smoketests/tests/smoketests/client_connection_errors.rs +++ b/crates/smoketests/tests/smoketests/client_connection_errors.rs @@ -1,3 +1,4 @@ +use socket2::SockRef; use spacetimedb_smoketests::{require_local_server, Smoketest}; use std::io::Write; use std::net::TcpStream; @@ -70,13 +71,13 @@ fn test_http_reducer_call_cancel_still_deletes_st_client() { .precompiled_module("client-connection-http-cancel") .build(); - let streams = send_http_reducer_calls_and_hold_connections(&test, "slow", 64); + let stream = send_http_reducer_call_and_hold_connection(&test, "slow"); wait_for_log( &test, "http_cancel_repro: slow reducer started", Duration::from_secs(30), ); - drop(streams); + abort_tcp_stream(stream); wait_for_log( &test, @@ -87,29 +88,32 @@ fn test_http_reducer_call_cancel_still_deletes_st_client() { let (st_client_count, st_connection_credentials_count) = wait_for_client_rows_to_clear(&test); assert!( st_client_count <= 1, - "Expected at most 1 st_client row (the SQL connection itself) after dropping 64 HTTP calls post-connect, got {st_client_count}", + "Expected at most 1 st_client row (the SQL connection itself) after dropping an HTTP call post-connect, got {st_client_count}", ); assert!( st_connection_credentials_count <= 1, - "Expected at most 1 st_connection_credentials row (the SQL connection itself) after dropping 64 HTTP calls post-connect, got {st_connection_credentials_count}", + "Expected at most 1 st_connection_credentials row (the SQL connection itself) after dropping an HTTP call post-connect, got {st_connection_credentials_count}", ); } -fn send_http_reducer_calls_and_hold_connections(test: &Smoketest, reducer: &str, count: usize) -> Vec { +fn send_http_reducer_call_and_hold_connection(test: &Smoketest, reducer: &str) -> TcpStream { let host = test.server_host(); let request = http_reducer_call_request(test, reducer); - (0..count) - .map(|_| { - let mut stream = TcpStream::connect(host).expect("Failed to connect to local server"); - stream.set_nodelay(true).expect("Failed to set TCP_NODELAY"); - stream - .write_all(request.as_bytes()) - .expect("Failed to write HTTP request"); - stream.flush().expect("Failed to flush HTTP request"); - stream - }) - .collect() + let mut stream = TcpStream::connect(host).expect("Failed to connect to local server"); + stream.set_nodelay(true).expect("Failed to set TCP_NODELAY"); + stream + .write_all(request.as_bytes()) + .expect("Failed to write HTTP request"); + stream.flush().expect("Failed to flush HTTP request"); + stream +} + +fn abort_tcp_stream(stream: TcpStream) { + SockRef::from(&stream) + .set_linger(Some(Duration::ZERO)) + .expect("Failed to set SO_LINGER=0"); + drop(stream); } fn http_reducer_call_request(test: &Smoketest, reducer: &str) -> String { @@ -118,13 +122,15 @@ fn http_reducer_call_request(test: &Smoketest, reducer: &str) -> String { let token = test.read_token().expect("Failed to read auth token"); let path = format!("/v1/database/{identity}/call/{reducer}"); + // Keep the HTTP/1.1 connection alive so dropping the TCP stream is an + // unexpected client disconnect while the reducer call is still active. format!( "POST {path} HTTP/1.1\r\n\ Host: {host}\r\n\ Authorization: Bearer {token}\r\n\ Content-Type: application/json\r\n\ Content-Length: 2\r\n\ - Connection: close\r\n\ + Connection: keep-alive\r\n\ \r\n\ []" ) @@ -142,7 +148,7 @@ fn wait_for_log(test: &Smoketest, needle: &str, timeout: Duration) -> Vec (usize, usize) { let start = Instant::now(); let mut last = (usize::MAX, usize::MAX); - while start.elapsed() < Duration::from_secs(20) { + while start.elapsed() < Duration::from_secs(5) { last = ( sql_count(test, "st_client"), sql_count(test, "st_connection_credentials"), @@ -158,7 +164,7 @@ fn wait_for_client_rows_to_clear(test: &Smoketest) -> (usize, usize) { if last.0 <= 1 && last.1 <= 1 { return last; } - thread::sleep(Duration::from_millis(250)); + thread::sleep(Duration::from_millis(50)); } last From 083fcbbc4b9d6d6ddfbbb43cc9ab21369f89fdb3 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Thu, 9 Jul 2026 12:28:09 +0530 Subject: [PATCH 3/4] move whole body in tokio spawn' --- crates/client-api/src/routes/database.rs | 181 ++++++++++++----------- 1 file changed, 95 insertions(+), 86 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index 85698f02a52..c98165a0dab 100644 --- a/crates/client-api/src/routes/database.rs +++ b/crates/client-api/src/routes/database.rs @@ -43,7 +43,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 +152,67 @@ pub async fn call( 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 |module: ModuleHost, 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, auth.claims.identity, fut).await } - #[derive(Deserialize)] pub struct HttpRouteRootParams { name_or_identity: NameOrIdentity, @@ -709,6 +698,42 @@ pub struct SqlQueryParams { pub confirmed: Option, } +async fn with_connection( + module: ModuleHost, + caller_auth: ConnectionAuthCtx, + caller_identity: Identity, + func: F, +) -> axum::response::Result +where + F: FnOnce(ModuleHost, ConnectionId) -> Fut + Send + 'static, + Fut: Future> + Send + 'static, + R: Send + 'static, +{ + tokio::spawn(async move { + // HTTP callers always need a connection ID to provide to connect/disconnect, + // so generate one. + 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 = func(module.clone(), 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( worker_ctx: S, SqlParams { name_or_identity }: SqlParams, @@ -718,21 +743,12 @@ pub async fn sql_direct( sql: String, ) -> axum::response::Result>> 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| { let sql_auth = worker_ctx .authorize_sql(caller_identity, database.database_identity) .await?; @@ -744,16 +760,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( @@ -764,7 +773,7 @@ pub async fn sql( body: String, ) -> axum::response::Result where - S: NodeDelegate + ControlStateDelegate + Authorization, + S: NodeDelegate + ControlStateDelegate + Authorization + 'static, { let caller_identity = auth.claims.identity; let caller_auth: ConnectionAuthCtx = auth.into(); From 297c01da674e4a7dc15a4f0d4406c825fd3f78e1 Mon Sep 17 00:00:00 2001 From: Shubham Mishra Date: Thu, 9 Jul 2026 13:23:18 +0530 Subject: [PATCH 4/4] delegate http handler work to tokio spawn tasks --- crates/client-api/src/routes/database.rs | 23 ++++++++++++++--------- crates/pg/src/pg_server.rs | 6 +++--- 2 files changed, 17 insertions(+), 12 deletions(-) diff --git a/crates/client-api/src/routes/database.rs b/crates/client-api/src/routes/database.rs index c98165a0dab..862597d289c 100644 --- a/crates/client-api/src/routes/database.rs +++ b/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; @@ -156,7 +157,7 @@ pub async fn call( let (module, Database { owner_identity, .. }) = find_module_and_database(&worker_ctx, name_or_identity).await?; - let fut = async |module: ModuleHost, connection_id: ConnectionId| { + let fut = async move |module: ModuleHost, caller_identity: Identity, connection_id: ConnectionId| { let result = match module .call_reducer( caller_identity, @@ -211,7 +212,7 @@ pub async fn call( } }; - with_connection(module, caller_auth, auth.claims.identity, fut).await + with_connection(module, caller_auth, caller_identity, fut).await } #[derive(Deserialize)] pub struct HttpRouteRootParams { @@ -698,29 +699,33 @@ pub struct SqlQueryParams { pub confirmed: Option, } +/// 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( module: ModuleHost, caller_auth: ConnectionAuthCtx, caller_identity: Identity, - func: F, + fut: F, ) -> axum::response::Result where - F: FnOnce(ModuleHost, ConnectionId) -> Fut + Send + 'static, + F: FnOnce(ModuleHost, Identity, ConnectionId) -> Fut + Send + 'static, Fut: Future> + Send + 'static, R: Send + 'static, { tokio::spawn(async move { - // HTTP callers always need a connection ID to provide to connect/disconnect, - // so generate one. 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`. + // 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 = func(module.clone(), connection_id).await; + let result = fut(module.clone(), caller_identity, connection_id).await; // Always disconnect, even if authorization or execution failed. module @@ -748,7 +753,7 @@ where let (host, database) = find_leader_and_database(&worker_ctx, name_or_identity).await?; let module = host.module().await.map_err(log_and_500)?; - let fut = async move |_module: ModuleHost, caller_identity: Identity| { + 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?; diff --git a/crates/pg/src/pg_server.rs b/crates/pg/src/pg_server.rs index ebca2b61b95..15170a35866 100644 --- a/crates/pg/src/pg_server.rs +++ b/crates/pg/src/pg_server.rs @@ -150,7 +150,7 @@ struct PgSpacetimeDB { impl PgSpacetimeDB where - T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone, + T: ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static, { async fn exe_sql(&self, query: String) -> PgWireResult> { let params = self.cached.lock().await.clone().unwrap(); @@ -314,7 +314,7 @@ impl SimpleQueryHandler for PgSpacetimeDB where - T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone, + T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static, { async fn do_query(&self, _client: &mut C, query: &str) -> PgWireResult> where @@ -347,7 +347,7 @@ impl PgSpacetimeDBFactory { impl PgWireServerHandlers for PgSpacetimeDBFactory where - T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone, + T: Sync + Send + ControlStateReadAccess + ControlStateWriteAccess + NodeDelegate + Authorization + Clone + 'static, { fn simple_query_handler(&self) -> Arc { self.handler.clone()