From 92262e7ce6a31fc15ff49d3ed3242bb96d6325f2 Mon Sep 17 00:00:00 2001 From: Stuart Douglas Date: Thu, 6 Aug 2026 12:34:20 +1000 Subject: [PATCH 1/2] fix(relay): clean up in-flight websocket requests Signed-off-by: Stuart Douglas Amp-Thread-ID: https://ampcode.com/threads/T-019fd4c7-3d09-71af-b091-05920b3b0957 Co-authored-by: Amp --- crates/buzz-relay/src/connection.rs | 114 ++++++++++++++++++++++++-- crates/buzz-relay/src/handlers/req.rs | 101 ++++++++++++++--------- 2 files changed, 169 insertions(+), 46 deletions(-) diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index 72a7eb9126..03543bbb77 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::time::Duration; use axum::extract::ws::{Message as WsMessage, WebSocket}; +use dashmap::DashMap; use futures_util::{Sink, SinkExt, StreamExt}; use tokio::sync::{mpsc, Mutex, RwLock}; use tokio_util::sync::CancellationToken; @@ -29,6 +30,34 @@ const AUTH_TIMEOUT: Duration = Duration::from_secs(5); /// Shared mutable subscription map for a single WebSocket connection. pub(crate) type ConnectionSubscriptions = Arc>>>; +/// In-flight REQ handlers keyed by client subscription ID. +type PendingReqs = Arc>; + +fn start_pending_req( + pending_reqs: &PendingReqs, + sub_id: &str, + connection_cancel: &CancellationToken, +) -> (Uuid, CancellationToken) { + let request_id = Uuid::new_v4(); + let request_cancel = connection_cancel.child_token(); + if let Some((_, previous_cancel)) = + pending_reqs.insert(sub_id.to_owned(), (request_id, request_cancel.clone())) + { + previous_cancel.cancel(); + } + (request_id, request_cancel) +} + +fn finish_pending_req(pending_reqs: &PendingReqs, sub_id: &str, request_id: Uuid) { + pending_reqs.remove_if(sub_id, |_, (active_id, _)| *active_id == request_id); +} + +fn cancel_pending_req(pending_reqs: &PendingReqs, sub_id: &str) { + if let Some((_, (_, request_cancel))) = pending_reqs.remove(sub_id) { + request_cancel.cancel(); + } +} + /// Request for the writer to flush a restart close and report the result. pub(crate) struct RestartClose { pub(crate) flushed: tokio::sync::oneshot::Sender, @@ -273,7 +302,15 @@ async fn handle_active_connection( let _ = heartbeat_task.await; let _ = auth_timeout_task.await; - for removed in state.sub_registry.remove_connection(conn.conn_id) { + // Synchronize teardown with REQ registration. A handler that acquired the + // mutex before cancellation finishes its registration first and is removed + // here; one that acquires it afterward observes its cancelled request token. + let removed_subscriptions = { + let mut subscriptions = conn.subscriptions.lock().await; + subscriptions.clear(); + state.sub_registry.remove_connection(conn.conn_id) + }; + for removed in removed_subscriptions { state .pubsub .release_topic(&conn.tenant, topic_for_subscription(removed.channel_id)) @@ -437,6 +474,7 @@ async fn recv_loop( missed_pongs: Arc, cancel: CancellationToken, ) { + let pending_reqs = Arc::new(DashMap::new()); loop { tokio::select! { msg = ws_recv.next() => { @@ -458,7 +496,12 @@ async fn recv_loop( break; } trace!(len = text.len(), "frame received"); - handle_text_message(text.to_string(), Arc::clone(&conn), Arc::clone(&state)).await; + handle_text_message( + text.to_string(), + Arc::clone(&conn), + Arc::clone(&state), + Arc::clone(&pending_reqs), + ).await; } Some(Ok(WsMessage::Binary(bytes))) => { let max_frame_bytes = state.config.max_frame_bytes; @@ -480,7 +523,12 @@ async fn recv_loop( // (notably certain Nostr libraries) send text payloads in binary frames. // NIP-01 is text-only, but accepting binary is a common relay extension. if let Ok(text) = String::from_utf8(bytes.to_vec()) { - handle_text_message(text, Arc::clone(&conn), Arc::clone(&state)).await; + handle_text_message( + text, + Arc::clone(&conn), + Arc::clone(&state), + Arc::clone(&pending_reqs), + ).await; } } Some(Ok(WsMessage::Pong(_))) => { @@ -511,7 +559,12 @@ async fn recv_loop( } } -async fn handle_text_message(text: String, conn: Arc, state: Arc) { +async fn handle_text_message( + text: String, + conn: Arc, + state: Arc, + pending_reqs: PendingReqs, +) { let msg = match ClientMessage::parse(&text) { Ok(m) => m, Err(e) => { @@ -573,10 +626,20 @@ async fn handle_text_message(text: String, conn: Arc, state: Ar return; } }; + let (request_id, request_cancel) = + start_pending_req(&pending_reqs, &sub_id, &conn.cancel); let span = tracing::info_span!("ws.req", conn_id = %conn.conn_id, sub_id = %sub_id); tokio::spawn( async move { - handlers::req::handle_req(sub_id, filters, conn, state).await; + handlers::req::handle_req( + sub_id.clone(), + filters, + Arc::clone(&conn), + state, + request_cancel, + ) + .await; + finish_pending_req(&pending_reqs, &sub_id, request_id); drop(permit); } .instrument(span), @@ -604,6 +667,7 @@ async fn handle_text_message(text: String, conn: Arc, state: Ar ); } ClientMessage::Close(sub_id) => { + cancel_pending_req(&pending_reqs, &sub_id); handlers::close::handle_close(sub_id, Arc::clone(&conn), Arc::clone(&state)).await; } } @@ -797,6 +861,46 @@ mod tests { .collect() } + #[test] + fn close_cancels_an_in_flight_req() { + let pending_reqs = Arc::new(DashMap::new()); + let connection_cancel = CancellationToken::new(); + let (_, request_cancel) = start_pending_req(&pending_reqs, "history", &connection_cancel); + + cancel_pending_req(&pending_reqs, "history"); + + assert!(request_cancel.is_cancelled()); + assert!(pending_reqs.is_empty()); + } + + #[test] + fn replacement_req_cancels_only_the_previous_generation() { + let pending_reqs = Arc::new(DashMap::new()); + let connection_cancel = CancellationToken::new(); + let (old_id, old_cancel) = start_pending_req(&pending_reqs, "live", &connection_cancel); + let (new_id, new_cancel) = start_pending_req(&pending_reqs, "live", &connection_cancel); + + assert!(old_cancel.is_cancelled()); + assert!(!new_cancel.is_cancelled()); + + finish_pending_req(&pending_reqs, "live", old_id); + assert!(pending_reqs.contains_key("live")); + + finish_pending_req(&pending_reqs, "live", new_id); + assert!(pending_reqs.is_empty()); + } + + #[test] + fn connection_close_cancels_an_in_flight_req() { + let pending_reqs = Arc::new(DashMap::new()); + let connection_cancel = CancellationToken::new(); + let (_, request_cancel) = start_pending_req(&pending_reqs, "history", &connection_cancel); + + connection_cancel.cancel(); + + assert!(request_cancel.is_cancelled()); + } + #[test] fn req_rejections_are_subscription_scoped() { let reason = "rate-limited: too many concurrent requests"; diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index fd7deadf51..26910e824e 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -15,6 +15,7 @@ use buzz_db::EventQuery; use buzz_pubsub::EventTopic; use hex; use nostr::Filter; +use tokio_util::sync::CancellationToken; use buzz_auth::Scope; @@ -45,6 +46,7 @@ pub async fn handle_req( filters: Vec, conn: Arc, state: Arc, + request_cancel: CancellationToken, ) { let (conn_id, pubkey_bytes, token_channel_ids) = { let auth = conn.auth_state.read().await; @@ -60,16 +62,6 @@ pub async fn handle_req( } let pk_bytes = ctx.pubkey.to_bytes().to_vec(); - - let subs = conn.subscriptions.lock().await; - if !subs.contains_key(&sub_id) && subs.len() >= MAX_SUBSCRIPTIONS { - conn.send(RelayMessage::closed( - &sub_id, - "error: too many subscriptions", - )); - return; - } - (conn.conn_id, pk_bytes, ctx.channel_ids.clone()) } _ => { @@ -216,45 +208,58 @@ pub async fn handle_req( )); return; } - handle_search_req( - &sub_id, - &filters, - &accessible_channels, - token_channel_ids.is_none(), - &conn.tenant, - &pubkey_bytes, - &conn, - &state, - trace_state.as_ref(), - ) - .await; + tokio::select! { + biased; + _ = request_cancel.cancelled() => {} + _ = handle_search_req( + &sub_id, + &filters, + &accessible_channels, + token_channel_ids.is_none(), + &conn.tenant, + &pubkey_bytes, + &conn, + &state, + trace_state.as_ref(), + ) => {} + } return; } - { + let replaced = { let mut subs = conn.subscriptions.lock().await; + if request_cancel.is_cancelled() { + return; + } + if !subs.contains_key(&sub_id) && subs.len() >= MAX_SUBSCRIPTIONS { + conn.send(RelayMessage::closed( + &sub_id, + "error: too many subscriptions", + )); + return; + } subs.insert(sub_id.clone(), filters.clone()); - } - - let replaced = state.sub_registry.register_scoped( - conn.tenant.community(), - conn_id, - sub_id.clone(), - filters.clone(), - channel_id, - ); - if let Some(replaced) = replaced { + let replaced = state.sub_registry.register_scoped( + conn.tenant.community(), + conn_id, + sub_id.clone(), + filters.clone(), + channel_id, + ); + if let Some(replaced) = replaced { + state + .pubsub + .release_topic(&conn.tenant, topic_for_subscription(replaced.channel_id)) + .await; + } state .pubsub - .release_topic(&conn.tenant, topic_for_subscription(replaced.channel_id)) + .retain_topic(&conn.tenant, topic_for_subscription(channel_id)) .await; - } - state - .pubsub - .retain_topic(&conn.tenant, topic_for_subscription(channel_id)) - .await; + replaced + }; - debug!(conn_id = %conn_id, sub_id = %sub_id, "Subscription registered"); + debug!(conn_id = %conn_id, sub_id = %sub_id, replaced = replaced.is_some(), "Subscription registered"); // NIP-01 OR semantics: execute one DB query per filter and deduplicate results // by event ID. Collapsing all filters into a single query would merge their @@ -317,7 +322,15 @@ pub async fn handle_req( .buffered(FILTER_QUERY_CONCURRENCY); // Phase 3 — post-processing, strictly in filter order. - while let Some((idx, per_filter_channel, filter_events)) = results.next().await { + loop { + let next = tokio::select! { + biased; + _ = request_cancel.cancelled() => return, + next = results.next() => next, + }; + let Some((idx, per_filter_channel, filter_events)) = next else { + break; + }; let filter = &filters[idx]; let events = match filter_events { Ok(evs) => evs, @@ -370,6 +383,9 @@ pub async fn handle_req( } for stored in &events { + if request_cancel.is_cancelled() { + return; + } // Per-filter NIP-01 matching — use the current filter only, not the // full filter set. OR semantics across filters are handled by the outer // loop (each filter gets its own DB query). @@ -410,6 +426,9 @@ pub async fn handle_req( } } + if request_cancel.is_cancelled() { + return; + } conn.send(RelayMessage::eose(&sub_id)); debug!( From cd62a82676772e718b3c9825821780e27f6994e1 Mon Sep 17 00:00:00 2001 From: Stuart Douglas Date: Thu, 6 Aug 2026 13:14:34 +1000 Subject: [PATCH 2/2] fix(relay): close request lifecycle races Signed-off-by: Stuart Douglas Amp-Thread-ID: https://ampcode.com/threads/T-019fd4c7-3d09-71af-b091-05920b3b0957 Co-authored-by: Amp --- crates/buzz-relay/src/connection.rs | 322 ++++++++++++++++++++++++-- crates/buzz-relay/src/handlers/req.rs | 277 ++++++++++++++++------ 2 files changed, 516 insertions(+), 83 deletions(-) diff --git a/crates/buzz-relay/src/connection.rs b/crates/buzz-relay/src/connection.rs index 03543bbb77..b09a93e8aa 100644 --- a/crates/buzz-relay/src/connection.rs +++ b/crates/buzz-relay/src/connection.rs @@ -31,7 +31,7 @@ const AUTH_TIMEOUT: Duration = Duration::from_secs(5); pub(crate) type ConnectionSubscriptions = Arc>>>; /// In-flight REQ handlers keyed by client subscription ID. -type PendingReqs = Arc>; +pub(crate) type PendingReqs = Arc>; fn start_pending_req( pending_reqs: &PendingReqs, @@ -302,20 +302,7 @@ async fn handle_active_connection( let _ = heartbeat_task.await; let _ = auth_timeout_task.await; - // Synchronize teardown with REQ registration. A handler that acquired the - // mutex before cancellation finishes its registration first and is removed - // here; one that acquires it afterward observes its cancelled request token. - let removed_subscriptions = { - let mut subscriptions = conn.subscriptions.lock().await; - subscriptions.clear(); - state.sub_registry.remove_connection(conn.conn_id) - }; - for removed in removed_subscriptions { - state - .pubsub - .release_topic(&conn.tenant, topic_for_subscription(removed.channel_id)) - .await; - } + cleanup_connection_subscriptions(&conn, &state).await; state.conn_manager.deregister(conn.conn_id); if let AuthState::Authenticated(ref auth_ctx) = *conn.auth_state.read().await { let remaining = state.conn_manager.connection_ids_for_pubkey_in_community( @@ -335,6 +322,26 @@ async fn handle_active_connection( drop(permit); } +/// Remove all subscription state owned by one connection. +/// +/// Taking the connection subscription lock synchronizes teardown with REQ +/// registration: a handler already in its commit section finishes first and is +/// removed here; one entering afterward observes cancellation and does not +/// register. +async fn cleanup_connection_subscriptions(conn: &ConnectionState, state: &AppState) { + let removed_subscriptions = { + let mut subscriptions = conn.subscriptions.lock().await; + subscriptions.clear(); + state.sub_registry.remove_connection(conn.conn_id) + }; + for removed in removed_subscriptions { + state + .pubsub + .release_topic(&conn.tenant, topic_for_subscription(removed.channel_id)) + .await; + } +} + /// Outbound send loop with control-frame priority. /// /// Control frames (Pong, Close) are drained first on every iteration, @@ -637,6 +644,8 @@ async fn handle_text_message( Arc::clone(&conn), state, request_cancel, + Arc::clone(&pending_reqs), + request_id, ) .await; finish_pending_req(&pending_reqs, &sub_id, request_id); @@ -778,6 +787,7 @@ fn topic_for_subscription(channel_id: Option) -> EventTopic { mod tests { use super::*; use std::sync::{Arc, Mutex}; + use tokio::sync::{oneshot, Notify}; #[derive(Debug, Default)] struct MockSinkState { @@ -861,6 +871,114 @@ mod tests { .collect() } + async fn subscription_test_state() -> Arc { + let mut config = crate::config::Config::from_env().expect("default config loads"); + config.require_relay_membership = false; + config.redis_url = "redis://127.0.0.1:1".to_string(); + let pool = sqlx::PgPool::connect_lazy(&config.database_url).expect("lazy pg pool"); + let db = buzz_db::Db::from_pool(pool.clone()); + let redis_pool = deadpool_redis::Config::from_url(&config.redis_url) + .create_pool(Some(deadpool_redis::Runtime::Tokio1)) + .expect("redis pool"); + let pubsub = Arc::new( + buzz_pubsub::PubSubManager::new(&config.redis_url, redis_pool.clone()) + .await + .expect("pubsub manager"), + ); + let audit = buzz_audit::AuditService::new(pool.clone()); + let auth = buzz_auth::AuthService::new(config.auth.clone()); + let search = buzz_search::SearchService::new(pool.clone()); + let workflow_engine = Arc::new(buzz_workflow::WorkflowEngine::new( + db.clone(), + buzz_workflow::WorkflowConfig::default(), + )); + let media_storage = buzz_media::MediaStorage::new(&config.media).expect("media storage"); + let (state, _audit_shutdown) = AppState::new( + config, + db, + redis_pool, + audit, + pubsub, + auth, + search, + workflow_engine, + nostr::Keys::generate(), + media_storage, + ); + Arc::new(state) + } + + fn subscription_test_conn_with_receiver() -> (Arc, mpsc::Receiver) { + let (send_tx, send_rx) = mpsc::channel(8); + let (ctrl_tx, _ctrl_rx) = mpsc::channel(8); + let conn = Arc::new(ConnectionState { + conn_id: Uuid::new_v4(), + tenant: buzz_core::TenantContext::resolved( + buzz_core::CommunityId::from_uuid(Uuid::new_v4()), + "relay.example", + ), + remote_addr: "127.0.0.1:1234".parse().expect("socket addr"), + auth_state: RwLock::new(AuthState::Pending { + challenge: "test".to_string(), + }), + subscriptions: Arc::new(tokio::sync::Mutex::new(HashMap::new())), + send_tx, + ctrl_tx, + cancel: CancellationToken::new(), + backpressure_count: Arc::new(AtomicU8::new(0)), + grace_limit: 3, + }); + (conn, send_rx) + } + + fn subscription_test_conn() -> Arc { + subscription_test_conn_with_receiver().0 + } + + fn spawn_blocked_registration( + state: Arc, + conn: Arc, + sub_id: &'static str, + filters: Vec, + request_cancel: CancellationToken, + ) -> ( + oneshot::Receiver<()>, + Arc, + tokio::task::JoinHandle>, + ) { + let (started_tx, started_rx) = oneshot::channel(); + let resume = Arc::new(Notify::new()); + let task_resume = Arc::clone(&resume); + let task = tokio::spawn(async move { + crate::handlers::req::register_subscription_after_check_for_test( + sub_id, + &filters, + None, + &conn, + &state, + &request_cancel, + async move { + let _ = started_tx.send(()); + task_resume.notified().await; + }, + ) + .await + }); + (started_rx, resume, task) + } + + async fn assert_subscription_state_empty(state: &AppState, conn: &ConnectionState) { + assert!(conn.subscriptions.lock().await.is_empty()); + assert_eq!(state.sub_registry.total_subscriptions(), 0); + assert_eq!( + state + .pubsub + .topic_refcount(&conn.tenant, EventTopic::Global) + .await, + 0 + ); + } + #[test] fn close_cancels_an_in_flight_req() { let pending_reqs = Arc::new(DashMap::new()); @@ -901,6 +1019,180 @@ mod tests { assert!(request_cancel.is_cancelled()); } + #[tokio::test] + async fn close_prevents_a_blocked_req_from_registering_late() { + let state = subscription_test_state().await; + let conn = subscription_test_conn(); + let pending_reqs = Arc::new(DashMap::new()); + let (_, request_cancel) = start_pending_req(&pending_reqs, "history", &conn.cancel); + let filters = vec![Filter::new().kind(nostr::Kind::TextNote)]; + let (started, resume, task) = spawn_blocked_registration( + Arc::clone(&state), + Arc::clone(&conn), + "history", + filters, + request_cancel, + ); + started.await.expect("blocked REQ started"); + + cancel_pending_req(&pending_reqs, "history"); + let mut close = Box::pin(crate::handlers::close::handle_close( + "history".to_string(), + Arc::clone(&conn), + Arc::clone(&state), + )); + assert!(futures_util::poll!(&mut close).is_pending()); + resume.notify_one(); + + assert_eq!(task.await.expect("blocked REQ task"), Some(false)); + close.await; + assert_subscription_state_empty(&state, &conn).await; + } + + #[tokio::test] + async fn disconnect_prevents_a_blocked_req_from_registering_late() { + let state = subscription_test_state().await; + let conn = subscription_test_conn(); + let pending_reqs = Arc::new(DashMap::new()); + let (_, request_cancel) = start_pending_req(&pending_reqs, "history", &conn.cancel); + let filters = vec![Filter::new().kind(nostr::Kind::TextNote)]; + let (started, resume, task) = spawn_blocked_registration( + Arc::clone(&state), + Arc::clone(&conn), + "history", + filters, + request_cancel, + ); + started.await.expect("blocked REQ started"); + + conn.cancel.cancel(); + let mut cleanup = Box::pin(cleanup_connection_subscriptions(&conn, &state)); + assert!(futures_util::poll!(&mut cleanup).is_pending()); + resume.notify_one(); + + assert_eq!(task.await.expect("blocked REQ task"), Some(false)); + cleanup.await; + assert_subscription_state_empty(&state, &conn).await; + } + + #[tokio::test] + async fn replacement_prevents_an_older_blocked_req_from_overwriting_state() { + let state = subscription_test_state().await; + let conn = subscription_test_conn(); + let pending_reqs = Arc::new(DashMap::new()); + let (_, old_cancel) = start_pending_req(&pending_reqs, "live", &conn.cancel); + let old_filters = vec![Filter::new().kind(nostr::Kind::TextNote)]; + let (started, resume, old_task) = spawn_blocked_registration( + Arc::clone(&state), + Arc::clone(&conn), + "live", + old_filters, + old_cancel, + ); + started.await.expect("old REQ started"); + + let (_, new_cancel) = start_pending_req(&pending_reqs, "live", &conn.cancel); + let new_filters = vec![Filter::new().kind(nostr::Kind::Reaction)]; + let mut replacement = Box::pin(crate::handlers::req::register_subscription_if_active( + "live", + &new_filters, + None, + &conn, + &state, + &new_cancel, + )); + assert!(futures_util::poll!(&mut replacement).is_pending()); + resume.notify_one(); + + assert_eq!(old_task.await.expect("old REQ task"), Some(false)); + assert_eq!(replacement.await, Some(true)); + assert_eq!( + conn.subscriptions.lock().await.get("live"), + Some(&new_filters) + ); + assert_eq!( + state.sub_registry.get_filters(conn.conn_id, "live"), + Some(new_filters) + ); + assert_eq!(state.sub_registry.total_subscriptions(), 1); + assert_eq!( + state + .pubsub + .topic_refcount(&conn.tenant, EventTopic::Global) + .await, + 1 + ); + + cancel_pending_req(&pending_reqs, "live"); + crate::handlers::close::handle_close( + "live".to_string(), + Arc::clone(&conn), + Arc::clone(&state), + ) + .await; + assert_subscription_state_empty(&state, &conn).await; + } + + #[tokio::test] + async fn close_and_replacement_suppress_cancelled_search_output() { + let state = subscription_test_state().await; + let (conn, mut send_rx) = subscription_test_conn_with_receiver(); + let pending_reqs = Arc::new(DashMap::new()); + let (close_id, close_cancel) = + start_pending_req(&pending_reqs, "search-close", &conn.cancel); + + cancel_pending_req(&pending_reqs, "search-close"); + crate::handlers::close::handle_close( + "search-close".to_string(), + Arc::clone(&conn), + Arc::clone(&state), + ) + .await; + assert!(!crate::handlers::req::send_search_frame_if_active( + &conn, + &pending_reqs, + "search-close", + close_id, + &close_cancel, + "stale EVENT after CLOSE".to_string(), + )); + assert!(!crate::handlers::req::send_search_frame_if_active( + &conn, + &pending_reqs, + "search-close", + close_id, + &close_cancel, + RelayMessage::eose("search-close"), + )); + + let closed = send_rx.try_recv().expect("CLOSED acknowledgement"); + let WsMessage::Text(closed) = closed else { + panic!("expected CLOSED text frame"); + }; + assert!(closed.contains(r#"["CLOSED","search-close""#)); + assert!(send_rx.try_recv().is_err(), "no stale search output"); + + let (old_id, old_cancel) = start_pending_req(&pending_reqs, "search-replace", &conn.cancel); + let (_, _new_cancel) = start_pending_req(&pending_reqs, "search-replace", &conn.cancel); + assert!(!crate::handlers::req::send_search_frame_if_active( + &conn, + &pending_reqs, + "search-replace", + old_id, + &old_cancel, + "stale EVENT after replacement".to_string(), + )); + assert!(!crate::handlers::req::send_search_frame_if_active( + &conn, + &pending_reqs, + "search-replace", + old_id, + &old_cancel, + RelayMessage::eose("search-replace"), + )); + assert!(send_rx.try_recv().is_err(), "no replaced-generation output"); + } + #[test] fn req_rejections_are_subscription_scoped() { let reason = "rate-limited: too many concurrent requests"; diff --git a/crates/buzz-relay/src/handlers/req.rs b/crates/buzz-relay/src/handlers/req.rs index 26910e824e..183431fe31 100644 --- a/crates/buzz-relay/src/handlers/req.rs +++ b/crates/buzz-relay/src/handlers/req.rs @@ -19,7 +19,7 @@ use tokio_util::sync::CancellationToken; use buzz_auth::Scope; -use crate::connection::{AuthState, ConnectionState}; +use crate::connection::{AuthState, ConnectionState, PendingReqs}; use crate::protocol::RelayMessage; use crate::state::AppState; @@ -47,6 +47,8 @@ pub async fn handle_req( conn: Arc, state: Arc, request_cancel: CancellationToken, + pending_reqs: PendingReqs, + request_id: uuid::Uuid, ) { let (conn_id, pubkey_bytes, token_channel_ids) = { let auth = conn.auth_state.read().await; @@ -208,58 +210,38 @@ pub async fn handle_req( )); return; } - tokio::select! { - biased; - _ = request_cancel.cancelled() => {} - _ = handle_search_req( - &sub_id, - &filters, - &accessible_channels, - token_channel_ids.is_none(), - &conn.tenant, - &pubkey_bytes, - &conn, - &state, - trace_state.as_ref(), - ) => {} - } + handle_search_req( + &sub_id, + &filters, + &accessible_channels, + token_channel_ids.is_none(), + &conn.tenant, + &pubkey_bytes, + &conn, + &state, + trace_state.as_ref(), + &request_cancel, + &pending_reqs, + request_id, + ) + .await; return; } - let replaced = { - let mut subs = conn.subscriptions.lock().await; - if request_cancel.is_cancelled() { - return; - } - if !subs.contains_key(&sub_id) && subs.len() >= MAX_SUBSCRIPTIONS { - conn.send(RelayMessage::closed( - &sub_id, - "error: too many subscriptions", - )); - return; - } - subs.insert(sub_id.clone(), filters.clone()); - let replaced = state.sub_registry.register_scoped( - conn.tenant.community(), - conn_id, - sub_id.clone(), - filters.clone(), - channel_id, - ); - if let Some(replaced) = replaced { - state - .pubsub - .release_topic(&conn.tenant, topic_for_subscription(replaced.channel_id)) - .await; - } - state - .pubsub - .retain_topic(&conn.tenant, topic_for_subscription(channel_id)) - .await; - replaced + let Some(replaced) = register_subscription_if_active( + &sub_id, + &filters, + channel_id, + &conn, + &state, + &request_cancel, + ) + .await + else { + return; }; - debug!(conn_id = %conn_id, sub_id = %sub_id, replaced = replaced.is_some(), "Subscription registered"); + debug!(conn_id = %conn_id, sub_id = %sub_id, replaced, "Subscription registered"); // NIP-01 OR semantics: execute one DB query per filter and deduplicate results // by event ID. Collapsing all filters into a single query would merge their @@ -439,6 +421,104 @@ pub async fn handle_req( ); } +/// Commit one regular subscription unless its REQ generation was cancelled. +/// +/// The connection subscription lock is the lifecycle barrier shared with +/// CLOSE and disconnect cleanup. Once the cancellation check passes, local +/// state, fan-out indexes, and Redis topic refcounts transition together before +/// cleanup can acquire the lock. +pub(crate) async fn register_subscription_if_active( + sub_id: &str, + filters: &[Filter], + channel_id: Option, + conn: &ConnectionState, + state: &AppState, + request_cancel: &CancellationToken, +) -> Option { + register_subscription_after_check( + sub_id, + filters, + channel_id, + conn, + state, + request_cancel, + std::future::ready(()), + ) + .await +} + +async fn register_subscription_after_check( + sub_id: &str, + filters: &[Filter], + channel_id: Option, + conn: &ConnectionState, + state: &AppState, + request_cancel: &CancellationToken, + after_check: F, +) -> Option +where + F: std::future::Future, +{ + let mut subs = conn.subscriptions.lock().await; + if request_cancel.is_cancelled() { + return None; + } + after_check.await; + if !subs.contains_key(sub_id) && subs.len() >= MAX_SUBSCRIPTIONS { + conn.send(RelayMessage::closed( + sub_id, + "error: too many subscriptions", + )); + return None; + } + subs.insert(sub_id.to_owned(), filters.to_vec()); + let replaced = state.sub_registry.register_scoped( + conn.tenant.community(), + conn.conn_id, + sub_id.to_owned(), + filters.to_vec(), + channel_id, + ); + if let Some(replaced) = replaced { + state + .pubsub + .release_topic(&conn.tenant, topic_for_subscription(replaced.channel_id)) + .await; + } + state + .pubsub + .retain_topic(&conn.tenant, topic_for_subscription(channel_id)) + .await; + Some(replaced.is_some()) +} + +/// Test seam that pauses registration after its cancellation check while the +/// lifecycle mutex remains held. +#[cfg(test)] +pub(crate) async fn register_subscription_after_check_for_test( + sub_id: &str, + filters: &[Filter], + channel_id: Option, + conn: &ConnectionState, + state: &AppState, + request_cancel: &CancellationToken, + after_check: F, +) -> Option +where + F: std::future::Future, +{ + register_subscription_after_check( + sub_id, + filters, + channel_id, + conn, + state, + request_cancel, + after_check, + ) + .await +} + /// FTS candidate hits fetched per page. Pages are always full regardless of /// the requested limit — post-filtering discards an unpredictable share of /// hits, so the scan fetches candidates in full pages rather than sizing @@ -538,6 +618,20 @@ pub(crate) fn build_search_channel_scope_filter( }) } +pub(crate) fn send_search_frame_if_active( + conn: &ConnectionState, + pending_reqs: &PendingReqs, + sub_id: &str, + request_id: uuid::Uuid, + request_cancel: &CancellationToken, + frame: String, +) -> bool { + let Some(active) = pending_reqs.get(sub_id) else { + return false; + }; + active.0 == request_id && !request_cancel.is_cancelled() && conn.send(frame) +} + /// Handle a NIP-50 search REQ: query Postgres FTS, fetch full events, deliver results, EOSE. /// Search subscriptions are one-shot — no persistent subscription is registered. #[allow(clippy::too_many_arguments)] @@ -551,7 +645,13 @@ async fn handle_search_req( conn: &ConnectionState, state: &AppState, trace_state: Option<&crate::conformance::AbstractState>, + request_cancel: &CancellationToken, + pending_reqs: &PendingReqs, + request_id: uuid::Uuid, ) { + if request_cancel.is_cancelled() { + return; + } // The community-wide channel scope (no #h tag on the filter). `None` means // "no accessible channels and no global access" → EOSE, exactly as the // legacy string-filter helper short-circuited. @@ -559,7 +659,14 @@ async fn handle_search_req( match build_search_channel_scope_filter(accessible_channels, include_global) { Some(scope) => scope, None => { - conn.send(RelayMessage::eose(sub_id)); + send_search_frame_if_active( + conn, + pending_reqs, + sub_id, + request_id, + request_cancel, + RelayMessage::eose(sub_id), + ); return; } }; @@ -567,6 +674,9 @@ async fn handle_search_req( let mut seen_ids: HashSet = HashSet::new(); for filter in filters { + if request_cancel.is_cancelled() { + return; + } let search_text = match &filter.search { Some(s) if !s.is_empty() => s.clone(), _ => continue, @@ -629,6 +739,9 @@ async fn handle_search_req( let mut emitted: u32 = 0; for page in 1..=MAX_SEARCH_PAGES { + if request_cancel.is_cancelled() { + return; + } if emitted >= limit { break; } @@ -646,7 +759,11 @@ async fn handle_search_req( mode: buzz_search::SearchMode::FullText, }; - let search_result = match state.search.search(&search_query).await { + let search_result = match tokio::select! { + biased; + _ = request_cancel.cancelled() => return, + result = state.search.search(&search_query) => result, + } { Ok(r) => r, Err(e) => { warn!(sub_id = %sub_id, "NIP-50 search failed: {e}"); @@ -664,11 +781,15 @@ async fn handle_search_req( if !hit_ids.is_empty() { let id_refs: Vec<&[u8]> = hit_ids.iter().map(|b| b.as_slice()).collect(); - let events = match state - .db - .get_events_by_ids_routed("req_search_hydrate", tenant.community(), &id_refs) - .await - { + let events = match tokio::select! { + biased; + _ = request_cancel.cancelled() => return, + result = state.db.get_events_by_ids_routed( + "req_search_hydrate", + tenant.community(), + &id_refs, + ) => result, + } { Ok(evs) => evs, Err(e) => { warn!(sub_id = %sub_id, "NIP-50 batch fetch failed: {e}"); @@ -696,17 +817,20 @@ async fn handle_search_req( } s.into_iter().collect() }; - let channel_communities = - match state.db.communities_of_channels(&distinct).await { - Ok(m) => m, - Err(e) => { - warn!( - sub_id = %sub_id, - "conformance row-community lookup failed: {e}" - ); - std::collections::HashMap::new() - } - }; + let channel_communities = match tokio::select! { + biased; + _ = request_cancel.cancelled() => return, + result = state.db.communities_of_channels(&distinct) => result, + } { + Ok(m) => m, + Err(e) => { + warn!( + sub_id = %sub_id, + "conformance row-community lookup failed: {e}" + ); + std::collections::HashMap::new() + } + }; crate::conformance::record_read_by_id_rows( &state.tracer, state_snap, @@ -723,6 +847,9 @@ async fn handle_search_req( .collect(); for id_array in &hit_ids { + if request_cancel.is_cancelled() { + return; + } if emitted >= limit { break; } @@ -749,7 +876,14 @@ async fn handle_search_req( if !seen_ids.insert(stored.event.id) { continue; } - if !conn.send(RelayMessage::event(sub_id, &stored.event)) { + if !send_search_frame_if_active( + conn, + pending_reqs, + sub_id, + request_id, + request_cancel, + RelayMessage::event(sub_id, &stored.event), + ) { return; } emitted += 1; @@ -762,7 +896,14 @@ async fn handle_search_req( } } - conn.send(RelayMessage::eose(sub_id)); + send_search_frame_if_active( + conn, + pending_reqs, + sub_id, + request_id, + request_cancel, + RelayMessage::eose(sub_id), + ); } /// Convert a single NIP-01 filter into an [`EventQuery`] for the database.