Skip to content

Commit a22ea8b

Browse files
Duncanwpfleger96
andcommitted
fix(relay): run connection cleanup alongside the bounded peer wait
Holding the read half through presence cleanup let a stalled Redis (the pool has no wait timeout) keep the socket open indefinitely. Running both concurrently caps the socket at the 1s peer wait while cleanup still never waits on the client. Co-authored-by: Will Pfleger <[email protected]> Signed-off-by: Will Pfleger <[email protected]>
1 parent e951072 commit a22ea8b

4 files changed

Lines changed: 70 additions & 23 deletions

File tree

‎Justfile‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -580,6 +580,7 @@ test-unit:
580580
+ test(=connection::terminal_close_delivery_tests::session_expiry_denial_closes_with_1008_on_real_socket)
581581
+ test(=connection::terminal_close_delivery_tests::server_holds_socket_for_client_close_reply)
582582
+ test(=connection::terminal_close_delivery_tests::server_finishes_promptly_when_client_answers_close)
583+
+ test(=connection::terminal_close_delivery_tests::stalled_cleanup_does_not_hold_socket_open)
583584
+ test(=handlers::auth::tests::b2_pre_cancelled_connection_never_becomes_authenticated)
584585
+ test(=handlers::auth::tests::fi_ban_check_error_emits_terminal_authorization_unavailable)
585586
+ test(=handlers::auth::tests::fi_root_authorization_denied_rows_emit_identical_frames)

‎crates/buzz-relay/src/connection.rs‎

Lines changed: 63 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -756,21 +756,25 @@ async fn handle_active_connection(
756756
}
757757
let _ = auth_cancel_task.await;
758758

759-
crate::handlers::close::release_connection_subscriptions(&conn, &state).await;
760-
state.conn_manager.deregister(conn.conn_id);
761-
if let Some(auth_ctx) = authenticated {
762-
if !state.conn_manager.has_admitted_connection(
763-
conn.tenant.community(),
764-
auth_ctx.pubkey.to_bytes().as_slice(),
765-
) {
766-
let _ = state
767-
.pubsub
768-
.clear_presence(&conn.tenant, &auth_ctx.pubkey)
769-
.await;
759+
// Release application state and wait on the peer concurrently: cleanup
760+
// never waits on the client, and a stalled Redis cannot hold the socket
761+
// open past the bounded peer wait.
762+
let cleanup = async {
763+
crate::handlers::close::release_connection_subscriptions(&conn, &state).await;
764+
state.conn_manager.deregister(conn.conn_id);
765+
if let Some(auth_ctx) = authenticated {
766+
if !state.conn_manager.has_admitted_connection(
767+
conn.tenant.community(),
768+
auth_ctx.pubkey.to_bytes().as_slice(),
769+
) {
770+
let _ = state
771+
.pubsub
772+
.clear_presence(&conn.tenant, &auth_ctx.pubkey)
773+
.await;
774+
}
770775
}
771-
}
772-
// Application state is released above; only the transport waits on the peer.
773-
await_client_close(ws_recv).await;
776+
};
777+
tokio::join!(cleanup, await_client_close(ws_recv));
774778
metrics::gauge!("buzz_ws_connections_active").decrement(1.0);
775779
info!(conn_id = %conn_id, addr = %addr, "WebSocket connection closed");
776780

@@ -3253,6 +3257,16 @@ mod terminal_close_delivery_tests {
32533257
async fn serve(
32543258
assertion: VerifiedAssertion,
32553259
prefire_expiry: bool,
3260+
) -> (std::net::SocketAddr, RelayUrl, Arc<AppState>, Uuid) {
3261+
serve_with_cleanup_pause(assertion, prefire_expiry, None).await
3262+
}
3263+
3264+
/// Like `serve`, but runs the handler inside `cleanup_pause` so a test
3265+
/// can stall the connection's final subscription cleanup.
3266+
async fn serve_with_cleanup_pause(
3267+
assertion: VerifiedAssertion,
3268+
prefire_expiry: bool,
3269+
cleanup_pause: Option<Arc<crate::handlers::close::test_seam::Pause>>,
32563270
) -> (std::net::SocketAddr, RelayUrl, Arc<AppState>, Uuid) {
32573271
let state = crate::state::tests::test_state().await;
32583272
let served_state = Arc::clone(&state);
@@ -3268,8 +3282,12 @@ mod terminal_close_delivery_tests {
32683282
let app = Router::new().route(
32693283
"/",
32703284
get(move |ws: WebSocketUpgrade| {
3271-
let (state, tenant, assertion) =
3272-
(Arc::clone(&state), tenant.clone(), assertion.clone());
3285+
let (state, tenant, assertion, cleanup_pause) = (
3286+
Arc::clone(&state),
3287+
tenant.clone(),
3288+
assertion.clone(),
3289+
cleanup_pause.clone(),
3290+
);
32733291
async move {
32743292
ws.on_upgrade(move |socket| async move {
32753293
let cancel = CancellationToken::new();
@@ -3288,7 +3306,7 @@ mod terminal_close_delivery_tests {
32883306
);
32893307
(gate, tx, rx, Some(expiry))
32903308
});
3291-
handle_active_connection(
3309+
let handler = handle_active_connection(
32923310
socket,
32933311
state,
32943312
"127.0.0.1:1234".parse().expect("client address"),
@@ -3298,8 +3316,15 @@ mod terminal_close_delivery_tests {
32983316
Some(assertion),
32993317
chrono::Utc::now(),
33003318
pre_built,
3301-
)
3302-
.await;
3319+
);
3320+
match cleanup_pause {
3321+
Some(pause) => {
3322+
crate::handlers::close::test_seam::CLEANUP_PAUSE
3323+
.scope(pause, handler)
3324+
.await;
3325+
}
3326+
None => handler.await,
3327+
}
33033328
})
33043329
}
33053330
}),
@@ -3470,7 +3495,7 @@ mod terminal_close_delivery_tests {
34703495
/// (seen live as `1006`). This raw client never replies: the server must
34713496
/// hold the socket for its full bounded wait, which starts only after
34723497
/// connect, so timing from before connect makes the lower bound
3473-
/// independent of scheduler delay. The wait runs after cleanup, so the
3498+
/// independent of scheduler delay. Cleanup runs alongside the wait, so the
34743499
/// connection is already deregistered once the socket ends.
34753500
#[tokio::test]
34763501
async fn server_holds_socket_for_client_close_reply() {
@@ -3490,8 +3515,8 @@ mod terminal_close_delivery_tests {
34903515
}
34913516

34923517
/// The client withholds its Close reply until the connection has
3493-
/// deregistered, so cleanup must run before the bounded peer wait: a
3494-
/// server that cleans up after the wait, or ignores the reply, needs at
3518+
/// deregistered, so cleanup must not wait on the peer: a server that
3519+
/// cleans up after the wait, or ignores the reply, needs at
34953520
/// least `WS_TERMINAL_FLUSH_TIMEOUT` from connect.
34963521
#[tokio::test]
34973522
async fn server_finishes_promptly_when_client_answers_close() {
@@ -3519,4 +3544,20 @@ mod terminal_close_delivery_tests {
35193544
"server must finish on the client's Close, not the timeout; took {elapsed:?}"
35203545
);
35213546
}
3547+
3548+
/// Final cleanup can block on Redis with no deadline (the pool has no
3549+
/// wait timeout). The peer wait runs alongside it, so a stalled cleanup
3550+
/// still lets the socket end once the bounded wait expires.
3551+
#[tokio::test]
3552+
async fn stalled_cleanup_does_not_hold_socket_open() {
3553+
let pause = Arc::new(crate::handlers::close::test_seam::Pause::default());
3554+
let (address, _, _, _) =
3555+
serve_with_cleanup_pause(expired_assertion(), true, Some(Arc::clone(&pause))).await;
3556+
let started = std::time::Instant::now();
3557+
let mut socket = raw_client_after_server_close(address).await;
3558+
3559+
pause.reached.notified().await;
3560+
let elapsed = time_to_eof(&mut socket, started).await;
3561+
assert!(elapsed >= WS_TERMINAL_FLUSH_TIMEOUT);
3562+
}
35223563
}

‎crates/buzz-relay/src/handlers/close.rs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,8 @@ pub(crate) async fn close_if_owner(
8989
/// detached REQ task can register after this runs.
9090
pub(crate) async fn release_connection_subscriptions(conn: &ConnectionState, state: &AppState) {
9191
debug_assert!(conn.cancel.is_cancelled());
92+
#[cfg(test)]
93+
test_seam::pause_at(&test_seam::CLEANUP_PAUSE).await;
9294
let mut subs = conn.subscriptions.lock().await;
9395
subs.clear();
9496
for removed in state.sub_registry.remove_connection(conn.conn_id) {
@@ -116,7 +118,8 @@ pub(crate) async fn release_scope_topics(
116118
/// Test-only pause points, scoped to one task so parallel tests don't
117119
/// interfere: `PAUSE` inside `close_if_owner` between teardown and the
118120
/// terminal frame, `RELEASE_PAUSE` inside `retire_locked` between the
119-
/// registry removal and the topic release.
121+
/// registry removal and the topic release, `CLEANUP_PAUSE` at the start of
122+
/// `release_connection_subscriptions`.
120123
#[cfg(test)]
121124
pub(crate) mod test_seam {
122125
use std::sync::Arc;
@@ -131,6 +134,7 @@ pub(crate) mod test_seam {
131134
tokio::task_local! {
132135
pub(crate) static PAUSE: Arc<Pause>;
133136
pub(crate) static RELEASE_PAUSE: Arc<Pause>;
137+
pub(crate) static CLEANUP_PAUSE: Arc<Pause>;
134138
}
135139

136140
pub(super) async fn pause_at(point: &'static tokio::task::LocalKey<Arc<Pause>>) {

‎scripts/run-tests.sh‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -344,6 +344,7 @@ run_unit_tests() {
344344
connection::terminal_close_delivery_tests::session_expiry_denial_closes_with_1008_on_real_socket
345345
connection::terminal_close_delivery_tests::server_holds_socket_for_client_close_reply
346346
connection::terminal_close_delivery_tests::server_finishes_promptly_when_client_answers_close
347+
connection::terminal_close_delivery_tests::stalled_cleanup_does_not_hold_socket_open
347348
handlers::auth::tests::b2_pre_cancelled_connection_never_becomes_authenticated
348349
handlers::auth::tests::fi_ban_check_error_emits_terminal_authorization_unavailable
349350
handlers::auth::tests::fi_root_authorization_denied_rows_emit_identical_frames

0 commit comments

Comments
 (0)