From e06d1124c9ffaa13aff4c56c510ad7f7bd6a19e1 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Fri, 10 Jul 2026 15:13:39 +0200 Subject: [PATCH 1/3] Improve websocket liveness checks and idle timeout handling - Idle timeout no longer requires receive activity alone: the send loop extends the idle deadline whenever it makes write progress (rate-limited to once per second). - On idle timeout, initiate a proper close handshake carrying an "idle timeout" reason instead of abruptly tearing down the connection, so well-behaved clients can tell why they were disconnected. The connection is torn down if the handshake does not complete within a 10s grace period. - The "websocket connection aborted" warning and ws_clients_aborted metric now fire only when the client actor task is actually aborted or panics, not on every normal disconnect. - Fix the kick log message to report the channel's configured capacity instead of its remaining capacity (always 0 at that point). --- crates/client-api/src/routes/subscribe.rs | 177 +++++++++++++++++--- crates/core/src/client/client_connection.rs | 2 +- crates/standalone/config.toml | 8 + 3 files changed, 165 insertions(+), 22 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 165216202e8..91bc0e5ec33 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -427,8 +427,10 @@ pub struct WebSocketOptions { pub ping_interval: Duration, /// Amount of time after which an idle connection is closed. /// - /// A connection is considered idle if no data is received nor sent. - /// This includes `Ping`/`Pong` frames used for keep-alive. + /// A connection is considered idle if no data is received from the client + /// (including `Pong` frames answering our keep-alive `Ping`s) *and* no + /// send progress is made towards it. A slow client that keeps accepting + /// data is not idle, no matter how long it takes to drain a large message. /// /// Value must be greater than `ping_interval`. /// @@ -527,7 +529,12 @@ async fn ws_client_actor_inner( let (ws_send, ws_recv) = ws.split(); // Set up the idle timer. + // + // The deadline is extended by the receive task whenever data arrives from + // the client, and by the send task whenever it makes progress writing to + // the socket (a slow client that keeps accepting data is not idle). let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); + let idle_tx = Arc::new(idle_tx); let idle_timer = ws_idle_timer(idle_rx); let bsatn_rlb_pool = client.module().subscriptions().bsatn_rlb_pool.clone(); @@ -541,6 +548,7 @@ async fn ws_client_actor_inner( sendrx, unordered_rx, bsatn_rlb_pool, + idle_tx.clone(), )); // Spawn a task to handle incoming messages. let recv_task = tokio::spawn(ws_recv_task( @@ -573,6 +581,15 @@ async fn ws_client_actor_inner( log::trace!("Client connection ended: {client_id}"); } +/// How long to wait for the close handshake to complete after the server +/// initiated a close due to idle timeout, before tearing down the connection. +const SERVER_CLOSE_GRACE: Duration = Duration::from_secs(10); + +/// How often, at most, the send loop extends the idle deadline when it makes +/// write progress. Purely to avoid hammering the idle timer's watch channel +/// once per frame on fast connections. +const WRITE_PROGRESS_INTERVAL: Duration = Duration::from_secs(1); + /// The main `select!` loop of the websocket client actor. /// /// > This function is defined standalone with generic parameters so that its @@ -584,8 +601,12 @@ async fn ws_client_actor_inner( /// - Drive the tasks handling the send and receive ends of the websockets to /// completion, terminating when either of them completes. /// -/// - Terminating if the connection is idle for longer than [`ActorConfig::idle_timeout`]. -/// The connection becomes idle if nothing is received from the socket. +/// - Initiating a close handshake if the connection is idle for longer than +/// [`ActorConfig::idle_timeout`]. The connection becomes idle if nothing is +/// received from the socket and no send progress is made. The close carries +/// an "idle timeout" reason so that clients can tell why they were +/// disconnected; if the handshake does not complete within +/// [`SERVER_CLOSE_GRACE`], the connection is torn down. /// /// - Periodically sending `Ping` frames to prevent the connection from becoming /// idle (the client is supposed to respond with `Pong`, which resets the @@ -697,9 +718,14 @@ async fn ws_main_loop( let mut ping_interval = tokio::time::interval(state.config.ping_interval); // Arm the first hotswap watcher. let watch_hotswap = hotswap(); + // Deadline for the close handshake to complete after we initiated a close + // due to idle timeout. Armed if and when the idle timer fires. + let close_grace = sleep_until(Instant::now()); + let mut timed_out = false; pin_mut!(watch_hotswap); pin_mut!(idle_timer); + pin_mut!(close_grace); loop { let closed = state.closed(); @@ -732,14 +758,32 @@ async fn ws_main_loop( break; }, - // Exit if we haven't heard from the client for too long. - _ = &mut idle_timer => { - log::debug!("Client {} timed out", state.client_id); + // If we haven't heard from the client for too long, initiate a + // close handshake carrying the reason, so well-behaved clients can + // report why they were disconnected. Give the handshake a grace + // period to complete before tearing the connection down. + _ = &mut idle_timer, if !timed_out => { + log::warn!("Client {} timed out, closing", state.client_id); WORKER_METRICS .ws_clients_idle_timed_out .with_label_values(&state.database) .inc(); state.record_disconnect(ClientDisconnectCause::IdleTimeout); + unordered_tx(UnorderedWsMessage::Close(CloseFrame { + code: CloseCode::Away, + reason: "idle timeout".into(), + })); + timed_out = true; + close_grace.as_mut().reset(Instant::now() + SERVER_CLOSE_GRACE); + }, + + // The close handshake initiated on idle timeout did not complete + // in time; tear the connection down. + _ = &mut close_grace, if timed_out => { + log::warn!( + "Client {} did not complete close handshake after idle timeout, aborting", + state.client_id + ); break; }, @@ -832,7 +876,7 @@ async fn ws_idle_timer(mut activity: watch::Receiver) { /// such that we wouldn't be able to receive any more messages anyway. async fn ws_recv_task( state: Arc, - idle_tx: watch::Sender, + idle_tx: Arc>, client_closed_metric: IntGauge, message_handler: impl Fn(DataMessage, Instant) -> MessageHandler, unordered_tx: mpsc::UnboundedSender, @@ -889,7 +933,7 @@ async fn ws_recv_task( /// state are dropped. fn ws_recv_loop( state: Arc, - idle_tx: watch::Sender, + idle_tx: Arc>, mut ws: impl Stream> + Unpin, ) -> impl Stream { fn receive_error_cause(error: &WsError) -> ClientDisconnectCause { @@ -1190,6 +1234,7 @@ impl Receiver for mpsc::Receiver { /// This is so `ws_client_actor_inner` keeps polling the receive end of the /// socket until the close handshake completes -- it would otherwise exit early /// when sending to `unordered` fails. +#[allow(clippy::too_many_arguments)] async fn ws_send_loop( state: Arc, config: ClientConfig, @@ -1197,9 +1242,10 @@ async fn ws_send_loop( messages: impl Receiver, unordered: mpsc::UnboundedReceiver, bsatn_rlb_pool: BsatnRowListBuilderPool, + idle_tx: Arc>, ) { let metrics = SendMetrics::new(state.database); - ws_send_loop_inner(state, ws, messages, unordered, move |encode_rx, frames_tx| { + ws_send_loop_inner(state, ws, messages, unordered, idle_tx, move |encode_rx, frames_tx| { ws_encode_task(metrics, config, encode_rx, frames_tx, bsatn_rlb_pool) }) .await @@ -1210,6 +1256,7 @@ async fn ws_send_loop_inner( mut ws: impl Sink + Unpin, mut messages: impl Receiver, mut unordered: mpsc::UnboundedReceiver, + idle_tx: Arc>, encoder: impl FnOnce(mpsc::UnboundedReceiver, mpsc::UnboundedSender) -> Encoder, ) where T: Into, @@ -1233,6 +1280,16 @@ async fn ws_send_loop_inner( let mut message_batch = Vec::new(); let (frames_tx, mut frames_rx) = mpsc::unbounded_channel(); + // When we last extended the idle deadline due to write progress. + // + // The socket accepting bytes means the client's TCP stack has been + // acknowledging previously sent data: the peer is alive, just possibly + // slow. Counting this as activity prevents the idle timer from + // disconnecting clients that are actively (if slowly) downloading a large + // message — such clients may not see our `Ping` for a long time, as it is + // queued in the TCP stream behind the message data. + let mut last_write_progress = Instant::now(); + let (encode_tx, encode_rx) = mpsc::unbounded_channel(); // Spawn the encode task. // @@ -1344,6 +1401,12 @@ async fn ws_send_loop_inner( log::warn!("error sending frame: {e:#}"); break 'outer; } + // Writing succeeded, so the client is making progress: + // extend the idle deadline (rate-limited). + if last_write_progress.elapsed() >= WRITE_PROGRESS_INTERVAL { + last_write_progress = Instant::now(); + idle_tx.send(state.next_idle_deadline()).ok(); + } } }, @@ -1994,7 +2057,7 @@ mod tests { use std::{ future::{poll_fn, Future}, pin::Pin, - sync::atomic::AtomicUsize, + sync::{atomic::AtomicUsize, Mutex}, task::{Context, Poll}, }; @@ -2090,6 +2153,10 @@ mod tests { } } + fn dummy_idle_tx() -> Arc> { + Arc::new(watch::channel(Instant::now()).0) + } + #[tokio::test(start_paused = true)] // see [NOTE: start_paused] async fn idle_timer_extends_sleep() { let timeout = Duration::from_millis(10); @@ -2121,7 +2188,7 @@ mod tests { let input = stream::iter(vec![Ok(WsMessage::Ping(Bytes::new()))]); pin_mut!(input); - let recv_loop = ws_recv_loop(state.clone(), idle_tx, input); + let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2143,7 +2210,7 @@ mod tests { ]); pin_mut!(input); - let recv_loop = ws_recv_loop(state.clone(), idle_tx, input); + let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2162,7 +2229,7 @@ mod tests { ]); pin_mut!(input); { - let recv_loop = ws_recv_loop(state.clone(), idle_tx, &mut input); + let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), &mut input); pin_mut!(recv_loop); state.close(); @@ -2183,7 +2250,7 @@ mod tests { ]); pin_mut!(input); { - let recv_loop = ws_recv_loop(state.clone(), idle_tx, &mut input); + let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), &mut input); pin_mut!(recv_loop); state.close(); @@ -2202,7 +2269,7 @@ mod tests { Ok(WsMessage::Ping(Bytes::new())), Ok(WsMessage::Pong(Bytes::new())), ]); - let recv_loop = ws_recv_loop(state, idle_tx, input); + let recv_loop = ws_recv_loop(state, Arc::new(idle_tx), input); pin_mut!(recv_loop); let mut new_idle_deadline = *idle_rx.borrow(); @@ -2287,6 +2354,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2311,6 +2379,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2365,6 +2434,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2416,6 +2486,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2463,6 +2534,20 @@ mod tests { let before = disconnect_count(state.database, ClientDisconnectCause::IdleTimeout); let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); + // Record the `Close` frame the main loop sends when the idle timer + // fires. Since we never complete the close handshake (both tasks are + // pending forever), the loop should tear down the connection after + // `SERVER_CLOSE_GRACE`. + let close_sent = Arc::new(Mutex::new(None)); + let unordered_tx = { + let close_sent = close_sent.clone(); + move |m| { + if let UnorderedWsMessage::Close(frame) = m { + *close_sent.lock().unwrap() = Some(frame); + } + } + }; + let start = Instant::now(); let mut t = tokio::spawn({ let state = state.clone(); @@ -2473,7 +2558,7 @@ mod tests { ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), - drop, + unordered_tx, ) .await } @@ -2489,11 +2574,51 @@ mod tests { t.await.unwrap(); let elapsed = start.elapsed(); - assert!(elapsed >= timeout); - assert!(elapsed < timeout + Duration::from_millis(10)); + assert!(elapsed >= timeout + SERVER_CLOSE_GRACE); + assert!(elapsed < timeout + SERVER_CLOSE_GRACE + Duration::from_millis(10)); + assert!(close_sent.lock().unwrap().is_some()); assert_disconnect_count_incremented(state.database, ClientDisconnectCause::IdleTimeout, before); } + #[tokio::test(start_paused = true)] // see [NOTE: start_paused] + async fn main_loop_exits_promptly_when_close_handshake_completes_after_idle_timeout() { + let state = Arc::new(dummy_actor_state_with_config(WebSocketOptions { + idle_timeout: Duration::from_millis(10), + ..<_>::default() + })); + let (_idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); + + // Pretend the client acknowledges the close immediately: + // the recv task terminates as soon as the `Close` frame is sent. + let notify = Arc::new(tokio::sync::Notify::new()); + let unordered_tx = { + let notify = notify.clone(); + move |m| { + if let UnorderedWsMessage::Close(_) = m { + notify.notify_one(); + } + } + }; + + let start = Instant::now(); + ws_main_loop( + state.clone(), + future::pending, + ws_idle_timer(idle_rx), + tokio::spawn(future::pending()), + tokio::spawn(async move { notify.notified().await }), + unordered_tx, + ) + .await; + + let elapsed = start.elapsed(); + assert!(elapsed >= Duration::from_millis(10)); + assert!( + elapsed < SERVER_CLOSE_GRACE, + "should not have waited for the close grace period: {elapsed:?}" + ); + } + #[tokio::test(start_paused = true)] // see [NOTE: start_paused] async fn main_loop_keepalive_keeps_alive() { let state = Arc::new(dummy_actor_state_with_config(WebSocketOptions { @@ -2534,7 +2659,9 @@ mod tests { } }); - let expected_timeout = (5 * state.config.ping_interval) + state.config.idle_timeout; + // After the pongs stop, the loop initiates a close handshake, which + // never completes here, so it exits after `SERVER_CLOSE_GRACE`. + let expected_timeout = (5 * state.config.ping_interval) + state.config.idle_timeout + SERVER_CLOSE_GRACE; let res = timeout(expected_timeout, t).await; let elapsed = start.elapsed(); @@ -2685,7 +2812,15 @@ mod tests { const NUM_CONTROL_FRAMES: usize = 2; let send_loop = tokio::spawn(async move { - ws_send_loop_inner(state, &mut received, messages_rx, unordered_rx, encoder).await; + ws_send_loop_inner( + state, + &mut received, + messages_rx, + unordered_rx, + dummy_idle_tx(), + encoder, + ) + .await; received }); messages_tx.send(Bytes::from_static(&[1; MESSAGE_SIZE])).await.unwrap(); diff --git a/crates/core/src/client/client_connection.rs b/crates/core/src/client/client_connection.rs index ed33e29b533..23e801e4b30 100644 --- a/crates/core/src/client/client_connection.rs +++ b/crates/core/src/client/client_connection.rs @@ -451,7 +451,7 @@ impl ClientConnectionSender { log::warn!( "Client {:?} exceeded channel capacity of {}, kicking", self.id, - self.sendtx.capacity(), + self.sendtx.max_capacity(), ); if let Some(metrics) = &self.metrics { metrics.outgoing_queue_disconnects.inc(); diff --git a/crates/standalone/config.toml b/crates/standalone/config.toml index 9eeef5d3535..8ed353c7661 100644 --- a/crates/standalone/config.toml +++ b/crates/standalone/config.toml @@ -44,6 +44,14 @@ directives = [ # Apply a V8 heap limit in MiB. Set to 0 to use V8's default limit. # heap-limit-mb = 0 +# [websocket] +# Interval at which keep-alive Ping frames are sent to clients. +# ping-interval = "15s" +# Close connections from which nothing was received, and to which no send +# progress was made, for this long. Slow clients that keep accepting data are +# not considered idle. Must be greater than ping-interval. +# idle-timeout = "30s" + [commitlog] # The maximum supported commitlog format version, also used for writing. # log-format-version = 1 From 86fbda48a3e55e0dcadd06159cd0fdec97d6d57f Mon Sep 17 00:00:00 2001 From: Jeffrey Dallatezza Date: Mon, 20 Jul 2026 06:12:21 -0700 Subject: [PATCH 2/3] Change idle timeout to track last activity (#5521) Instead of passing around part of a watch to extend the timeout, this change adds a `last_activity` field to `ActorState`. It can be updated by calling `record_activity()` on the actor state, and you can get an idle timeout future by calling `idle_timer` on the actor state. IMO this is a bit more straightforward, and it reduces how many arguments we are passing around. --- crates/client-api/src/routes/subscribe.rs | 173 +++++++--------------- 1 file changed, 53 insertions(+), 120 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 91bc0e5ec33..eef3f049855 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -43,7 +43,7 @@ use spacetimedb_client_api_messages::websocket::v2 as ws_v2; use spacetimedb_client_api_messages::websocket::v3 as ws_v3; use spacetimedb_lib::bsatn; use spacetimedb_lib::connection_id::{ConnectionId, ConnectionIdForUrl}; -use tokio::sync::{mpsc, watch}; +use tokio::sync::mpsc; use tokio::task::JoinHandle; use tokio::time::error::Elapsed; use tokio::time::{sleep_until, timeout, Instant}; @@ -346,6 +346,8 @@ struct ActorState { /// When the last `Ping` frame was written to the socket. /// Taken when the corresponding `Pong` arrives, to observe the roundtrip time. last_ping_sent: Mutex>, + // used to determine if the connection is idle. + last_activity: Arc>, } impl ActorState { @@ -363,6 +365,7 @@ impl ActorState { closed: AtomicBool::new(false), got_pong: AtomicBool::new(true), last_ping_sent: Mutex::new(None), + last_activity: Arc::new(Mutex::new(Instant::now())), } } @@ -401,8 +404,20 @@ impl ActorState { } } - pub fn next_idle_deadline(&self) -> Instant { - Instant::now() + self.config.idle_timeout + // Update the `last_activity watermark` to indicate that the connection is still active. + pub fn record_activity(&self) { + let mut last_activity = self.last_activity.lock().unwrap(); + *last_activity = Instant::now(); + } + + // This future completes if `self.config.idle_timeout` has elapsed since `self.record_activity()` was last called. + pub fn idle_timer(&self) -> impl Future + use<> { + ws_idle_timer(self.last_activity.clone(), self.config.idle_timeout) + } + + pub fn get_last_activity(&self) -> Instant { + let last_activity = self.last_activity.lock().unwrap(); + *last_activity } pub fn record_disconnect(&self, cause: ClientDisconnectCause) -> bool { @@ -528,15 +543,6 @@ async fn ws_client_actor_inner( // Split websocket into send and receive halves. let (ws_send, ws_recv) = ws.split(); - // Set up the idle timer. - // - // The deadline is extended by the receive task whenever data arrives from - // the client, and by the send task whenever it makes progress writing to - // the socket (a slow client that keeps accepting data is not idle). - let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); - let idle_tx = Arc::new(idle_tx); - let idle_timer = ws_idle_timer(idle_rx); - let bsatn_rlb_pool = client.module().subscriptions().bsatn_rlb_pool.clone(); // Spawn a task to send outgoing messages @@ -548,12 +554,10 @@ async fn ws_client_actor_inner( sendrx, unordered_rx, bsatn_rlb_pool, - idle_tx.clone(), )); // Spawn a task to handle incoming messages. let recv_task = tokio::spawn(ws_recv_task( state.clone(), - idle_tx, client_closed_metric, { let client = client.clone(); @@ -574,7 +578,7 @@ async fn ws_client_actor_inner( } }; - ws_main_loop(state, hotswap, idle_timer, send_task, recv_task, move |msg| { + ws_main_loop(state, hotswap, send_task, recv_task, move |msg| { let _ = unordered_tx.send(msg); }) .await; @@ -585,11 +589,6 @@ async fn ws_client_actor_inner( /// initiated a close due to idle timeout, before tearing down the connection. const SERVER_CLOSE_GRACE: Duration = Duration::from_secs(10); -/// How often, at most, the send loop extends the idle deadline when it makes -/// write progress. Purely to avoid hammering the idle timer's watch channel -/// once per frame on fast connections. -const WRITE_PROGRESS_INTERVAL: Duration = Duration::from_secs(1); - /// The main `select!` loop of the websocket client actor. /// /// > This function is defined standalone with generic parameters so that its @@ -700,7 +699,6 @@ const WRITE_PROGRESS_INTERVAL: Duration = Duration::from_secs(1); async fn ws_main_loop( state: Arc, hotswap: impl Fn() -> HotswapWatcher, - idle_timer: impl Future, mut send_task: JoinHandle<()>, mut recv_task: JoinHandle<()>, unordered_tx: impl Fn(UnorderedWsMessage), @@ -723,6 +721,8 @@ async fn ws_main_loop( let close_grace = sleep_until(Instant::now()); let mut timed_out = false; + let idle_timer = state.idle_timer(); + pin_mut!(watch_hotswap); pin_mut!(idle_timer); pin_mut!(close_grace); @@ -824,34 +824,19 @@ async fn ws_main_loop( } } -/// A sleep that can be extended by sending it new deadlines. -/// -/// Sleeps until the deadline appearing on the `activity` channel, -/// i.e. if a new deadline appears before the sleep finishes, -/// the sleep is reset to the new deadline. -/// -/// The `activity` should be updated whenever a new message is received. -async fn ws_idle_timer(mut activity: watch::Receiver) { - let mut deadline = *activity.borrow(); - let sleep = sleep_until(deadline); - pin_mut!(sleep); - +// Sleeps until the last_activity + idle_timeout is reached. The `last_activity` should be updated whenever +// the client proves it is not idle. +async fn ws_idle_timer(last_activity: Arc>, idle_timeout: Duration) { loop { - tokio::select! { - biased; - - Ok(()) = activity.changed() => { - let new_deadline = *activity.borrow_and_update(); - if new_deadline != deadline { - deadline = new_deadline; - sleep.as_mut().reset(deadline); - } - }, - - () = &mut sleep => { - break; - }, + let deadline = { + let last_activity = last_activity.lock().unwrap(); + *last_activity + idle_timeout + }; + let now = Instant::now(); + if deadline <= now { + break; } + sleep_until(deadline).await; } } @@ -876,7 +861,6 @@ async fn ws_idle_timer(mut activity: watch::Receiver) { /// such that we wouldn't be able to receive any more messages anyway. async fn ws_recv_task( state: Arc, - idle_tx: Arc>, client_closed_metric: IntGauge, message_handler: impl Fn(DataMessage, Instant) -> MessageHandler, unordered_tx: mpsc::UnboundedSender, @@ -889,7 +873,7 @@ async fn ws_recv_task( .total_incoming_queue_length .with_label_values(&state.database); let recv_queue = ws_recv_queue(state.clone(), unordered_tx.clone(), recv_queue_gauge, ws); - let recv_loop = pin!(ws_recv_loop(state.clone(), idle_tx, recv_queue)); + let recv_loop = pin!(ws_recv_loop(state.clone(), recv_queue)); let recv_handler = ws_client_message_handler(state.clone(), client_closed_metric, recv_loop); pin_mut!(recv_handler); @@ -933,7 +917,6 @@ async fn ws_recv_task( /// state are dropped. fn ws_recv_loop( state: Arc, - idle_tx: Arc>, mut ws: impl Stream> + Unpin, ) -> impl Stream { fn receive_error_cause(error: &WsError) -> ClientDisconnectCause { @@ -996,7 +979,7 @@ fn ws_recv_loop( }; match res { Ok(m) => { - idle_tx.send(state.next_idle_deadline()).ok(); + state.record_activity(); if !state.closed() { yield ClientMessage::from_message(m); @@ -1242,10 +1225,9 @@ async fn ws_send_loop( messages: impl Receiver, unordered: mpsc::UnboundedReceiver, bsatn_rlb_pool: BsatnRowListBuilderPool, - idle_tx: Arc>, ) { let metrics = SendMetrics::new(state.database); - ws_send_loop_inner(state, ws, messages, unordered, idle_tx, move |encode_rx, frames_tx| { + ws_send_loop_inner(state, ws, messages, unordered, move |encode_rx, frames_tx| { ws_encode_task(metrics, config, encode_rx, frames_tx, bsatn_rlb_pool) }) .await @@ -1256,7 +1238,6 @@ async fn ws_send_loop_inner( mut ws: impl Sink + Unpin, mut messages: impl Receiver, mut unordered: mpsc::UnboundedReceiver, - idle_tx: Arc>, encoder: impl FnOnce(mpsc::UnboundedReceiver, mpsc::UnboundedSender) -> Encoder, ) where T: Into, @@ -1280,16 +1261,6 @@ async fn ws_send_loop_inner( let mut message_batch = Vec::new(); let (frames_tx, mut frames_rx) = mpsc::unbounded_channel(); - // When we last extended the idle deadline due to write progress. - // - // The socket accepting bytes means the client's TCP stack has been - // acknowledging previously sent data: the peer is alive, just possibly - // slow. Counting this as activity prevents the idle timer from - // disconnecting clients that are actively (if slowly) downloading a large - // message — such clients may not see our `Ping` for a long time, as it is - // queued in the TCP stream behind the message data. - let mut last_write_progress = Instant::now(); - let (encode_tx, encode_rx) = mpsc::unbounded_channel(); // Spawn the encode task. // @@ -1401,12 +1372,7 @@ async fn ws_send_loop_inner( log::warn!("error sending frame: {e:#}"); break 'outer; } - // Writing succeeded, so the client is making progress: - // extend the idle deadline (rate-limited). - if last_write_progress.elapsed() >= WRITE_PROGRESS_INTERVAL { - last_write_progress = Instant::now(); - idle_tx.send(state.next_idle_deadline()).ok(); - } + state.record_activity(); } }, @@ -2153,20 +2119,16 @@ mod tests { } } - fn dummy_idle_tx() -> Arc> { - Arc::new(watch::channel(Instant::now()).0) - } - #[tokio::test(start_paused = true)] // see [NOTE: start_paused] async fn idle_timer_extends_sleep() { let timeout = Duration::from_millis(10); let start = Instant::now(); - let (tx, rx) = watch::channel(start + timeout); - tokio::join!(ws_idle_timer(rx), async { + let last_activity = Arc::new(Mutex::new(start)); + tokio::join!(ws_idle_timer(last_activity.clone(), timeout), async { for _ in 0..5 { sleep(Duration::from_millis(1)).await; - tx.send(Instant::now() + timeout).unwrap(); + *last_activity.lock().unwrap() = Instant::now(); } }); let elapsed = start.elapsed(); @@ -2183,12 +2145,11 @@ mod tests { async fn recv_loop_terminates_when_input_exhausted() { let state = Arc::new(actor_state_with_disconnect_recorder(1, <_>::default())); let before = disconnect_count(state.database, ClientDisconnectCause::WebsocketStreamEnded); - let (idle_tx, _idle_rx) = watch::channel(Instant::now() + state.config.idle_timeout); let input = stream::iter(vec![Ok(WsMessage::Ping(Bytes::new()))]); pin_mut!(input); - let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), input); + let recv_loop = ws_recv_loop(state.clone(), input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2201,7 +2162,6 @@ mod tests { let state = Arc::new(actor_state_with_disconnect_recorder(2, <_>::default())); let cause = ClientDisconnectCause::WebsocketReceiveConnectionClosed; let before = disconnect_count(state.database, cause); - let (idle_tx, _idle_rx) = watch::channel(Instant::now() + state.config.idle_timeout); let input = stream::iter(vec![ Ok(WsMessage::Ping(Bytes::new())), @@ -2210,7 +2170,7 @@ mod tests { ]); pin_mut!(input); - let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), input); + let recv_loop = ws_recv_loop(state.clone(), input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2221,7 +2181,6 @@ mod tests { #[tokio::test] async fn recv_loop_drains_remaining_messages_when_closed() { let state = Arc::new(dummy_actor_state()); - let (idle_tx, _idle_rx) = watch::channel(Instant::now() + state.config.idle_timeout); let input = stream::iter(vec![ Ok(WsMessage::Ping(Bytes::new())), @@ -2229,7 +2188,7 @@ mod tests { ]); pin_mut!(input); { - let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), &mut input); + let recv_loop = ws_recv_loop(state.clone(), &mut input); pin_mut!(recv_loop); state.close(); @@ -2241,7 +2200,6 @@ mod tests { #[tokio::test] async fn recv_loop_stops_at_error_while_draining() { let state = Arc::new(dummy_actor_state()); - let (idle_tx, _idle_rx) = watch::channel(Instant::now() + state.config.idle_timeout); let input = stream::iter(vec![ Ok(WsMessage::Ping(Bytes::new())), @@ -2250,7 +2208,7 @@ mod tests { ]); pin_mut!(input); { - let recv_loop = ws_recv_loop(state.clone(), Arc::new(idle_tx), &mut input); + let recv_loop = ws_recv_loop(state.clone(), &mut input); pin_mut!(recv_loop); state.close(); @@ -2262,23 +2220,22 @@ mod tests { #[tokio::test] async fn recv_loop_updates_idle_channel() { let state = Arc::new(dummy_actor_state()); - let idle_deadline = Instant::now() + state.config.idle_timeout; - let (idle_tx, mut idle_rx) = watch::channel(idle_deadline); + let mut prev_activity = state.get_last_activity(); + tokio::time::advance(Duration::from_millis(1)); let input = stream::iter(vec![ Ok(WsMessage::Ping(Bytes::new())), Ok(WsMessage::Pong(Bytes::new())), ]); - let recv_loop = ws_recv_loop(state, Arc::new(idle_tx), input); + let recv_loop = ws_recv_loop(state.clone(), input); pin_mut!(recv_loop); - let mut new_idle_deadline = *idle_rx.borrow(); while let Some(message) = recv_loop.next().await { + let last_activity = state.get_last_activity(); drop(message); - assert!(idle_rx.has_changed().unwrap()); - new_idle_deadline = *idle_rx.borrow_and_update(); + tokio::time::advance(Duration::from_millis(1)); + assert!(last_activity > prev_activity); } - assert!(new_idle_deadline > idle_deadline); } #[tokio::test] @@ -2321,14 +2278,12 @@ mod tests { async fn recv_task_records_client_message_error_disconnect() { let state = Arc::new(actor_state_with_disconnect_recorder(7, <_>::default())); let before = disconnect_count(state.database, ClientDisconnectCause::ClientMessageError); - let (idle_tx, _idle_rx) = watch::channel(state.next_idle_deadline()); let metric = IntGauge::new("bleep", "unhelpful").unwrap(); let (unordered_tx, mut unordered_rx) = mpsc::unbounded_channel(); let input = stream::iter([Ok(WsMessage::text("not useful"))]); ws_recv_task( state.clone(), - idle_tx, metric, |_data, _timer| future::ready(Err(MessageHandleError::UnsupportedVersion("test"))), unordered_tx, @@ -2354,7 +2309,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2379,7 +2333,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2434,7 +2387,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2486,7 +2438,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2505,7 +2456,6 @@ mod tests { ws_main_loop( state.clone(), future::pending, - future::pending(), tokio::spawn(sleep(Duration::from_millis(10))), tokio::spawn(future::pending()), drop, @@ -2514,7 +2464,6 @@ mod tests { ws_main_loop( state, future::pending, - future::pending(), tokio::spawn(future::pending()), tokio::spawn(sleep(Duration::from_millis(10))), drop, @@ -2532,7 +2481,6 @@ mod tests { }, )); let before = disconnect_count(state.database, ClientDisconnectCause::IdleTimeout); - let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); // Record the `Close` frame the main loop sends when the idle timer // fires. Since we never complete the close handshake (both tasks are @@ -2555,7 +2503,6 @@ mod tests { ws_main_loop( state, future::pending, - ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), unordered_tx, @@ -2567,7 +2514,7 @@ mod tests { let loop_start = Instant::now(); for _ in 0..5 { sleep(Duration::from_millis(5)).await; - idle_tx.send(state.next_idle_deadline()).unwrap(); + state.record_activity(); assert!(is_pending(&mut t).await); } let timeout = loop_start.elapsed() + Duration::from_millis(10); @@ -2586,7 +2533,6 @@ mod tests { idle_timeout: Duration::from_millis(10), ..<_>::default() })); - let (_idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); // Pretend the client acknowledges the close immediately: // the recv task terminates as soon as the `Close` frame is sent. @@ -2604,7 +2550,6 @@ mod tests { ws_main_loop( state.clone(), future::pending, - ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(async move { notify.notified().await }), unordered_tx, @@ -2626,7 +2571,6 @@ mod tests { idle_timeout: Duration::from_millis(10), ..<_>::default() })); - let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); // Pretend we received a pong immediately after sending a ping, // but only five times. let unordered_tx = { @@ -2637,7 +2581,7 @@ mod tests { let n = pings.fetch_add(1, Ordering::Relaxed); if n < 5 { state.set_ponged(); - idle_tx.send(state.next_idle_deadline()).ok(); + state.record_activity(); } } } @@ -2648,9 +2592,8 @@ mod tests { let state = state.clone(); async move { ws_main_loop( - state, + state.clone(), future::pending, - ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), unordered_tx, @@ -2682,7 +2625,6 @@ mod tests { let state = Arc::new(actor_state_with_disconnect_recorder(5, <_>::default())); let before = disconnect_count(state.database, ClientDisconnectCause::ModuleExited); - let (_idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); let unordered_tx = { let state = state.clone(); move |m| { @@ -2703,7 +2645,6 @@ mod tests { ws_main_loop( state_for_loop.clone(), hotswap, - ws_idle_timer(idle_rx), // Pretend we received a close immediately after sending one. tokio::spawn(async move { loop { @@ -2812,15 +2753,7 @@ mod tests { const NUM_CONTROL_FRAMES: usize = 2; let send_loop = tokio::spawn(async move { - ws_send_loop_inner( - state, - &mut received, - messages_rx, - unordered_rx, - dummy_idle_tx(), - encoder, - ) - .await; + ws_send_loop_inner(state, &mut received, messages_rx, unordered_rx, encoder).await; received }); messages_tx.send(Bytes::from_static(&[1; MESSAGE_SIZE])).await.unwrap(); From e9ee86e1714e3585f9fe04418712d6f5062d6359 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Wed, 12 Aug 2026 09:07:21 +0200 Subject: [PATCH 3/3] Don't count send progress as activity for the idle timeout Per review feedback: completing a ws.feed only means the socket buffer (including any LB/proxy buffers) isn't full, not that the client is making progress, so a dead client could be kept alive for a long time by small periodic messages. Testing also showed this did not help the disconnection issue being investigated. Only data received from the client now counts as activity. Also update docs that still described the old behavior, and fix the unawaited tokio::time::advance calls in the last-activity test. --- crates/client-api/src/routes/subscribe.rs | 34 +++++++++-------------- 1 file changed, 13 insertions(+), 21 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index eef3f049855..b7c47f599e3 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -415,6 +415,7 @@ impl ActorState { ws_idle_timer(self.last_activity.clone(), self.config.idle_timeout) } + #[cfg(test)] pub fn get_last_activity(&self) -> Instant { let last_activity = self.last_activity.lock().unwrap(); *last_activity @@ -442,10 +443,8 @@ pub struct WebSocketOptions { pub ping_interval: Duration, /// Amount of time after which an idle connection is closed. /// - /// A connection is considered idle if no data is received from the client - /// (including `Pong` frames answering our keep-alive `Ping`s) *and* no - /// send progress is made towards it. A slow client that keeps accepting - /// data is not idle, no matter how long it takes to drain a large message. + /// A connection is considered idle if no data is received from the client, + /// including `Pong` frames answering our keep-alive `Ping`s. /// /// Value must be greater than `ping_interval`. /// @@ -602,10 +601,9 @@ const SERVER_CLOSE_GRACE: Duration = Duration::from_secs(10); /// /// - Initiating a close handshake if the connection is idle for longer than /// [`ActorConfig::idle_timeout`]. The connection becomes idle if nothing is -/// received from the socket and no send progress is made. The close carries -/// an "idle timeout" reason so that clients can tell why they were -/// disconnected; if the handshake does not complete within -/// [`SERVER_CLOSE_GRACE`], the connection is torn down. +/// received from the socket. The close carries an "idle timeout" reason so +/// that clients can tell why they were disconnected; if the handshake does +/// not complete within [`SERVER_CLOSE_GRACE`], the connection is torn down. /// /// - Periodically sending `Ping` frames to prevent the connection from becoming /// idle (the client is supposed to respond with `Pong`, which resets the @@ -657,12 +655,6 @@ const SERVER_CLOSE_GRACE: Duration = Duration::from_secs(10); /// is `Err(NoSuchModule)`, the database was shut down and existing clients /// must be disconnected. /// -/// * **idle_timer**: -/// Abstraction for [`ws_idle_timer`]: if and when the future completes, the -/// connection is considered unresponsive, and the connection is closed. -/// -/// The idle timer should be reset whenever data is received from the websocket. -/// /// * **send_task**: /// Task handling outgoing messages. Holds the receive end of `unordered_tx`. /// @@ -843,8 +835,8 @@ async fn ws_idle_timer(last_activity: Arc>, idle_timeout: Duratio /// Consumes `ws` by composing [`ws_recv_queue`], [`ws_recv_loop`], /// [`ws_client_message_handler`] and `message_handler`. /// -/// `idle_tx` is the sending end of a [`ws_idle_timer`]. The [`ws_recv_loop`] -/// sends a new, extended deadline whenever it receives a message. +/// The [`ws_recv_loop`] records activity on the shared [`ActorState`] whenever +/// it receives a message, extending the idle deadline. /// /// `unordered_tx` is used to send message execution errors /// or to initiate a close handshake. @@ -1372,7 +1364,6 @@ async fn ws_send_loop_inner( log::warn!("error sending frame: {e:#}"); break 'outer; } - state.record_activity(); } }, @@ -2217,11 +2208,11 @@ mod tests { assert_matches!(input.next().await, Some(Ok(WsMessage::Pong(_)))); } - #[tokio::test] - async fn recv_loop_updates_idle_channel() { + #[tokio::test(start_paused = true)] // see [NOTE: start_paused] + async fn recv_loop_updates_last_activity() { let state = Arc::new(dummy_actor_state()); let mut prev_activity = state.get_last_activity(); - tokio::time::advance(Duration::from_millis(1)); + tokio::time::advance(Duration::from_millis(1)).await; let input = stream::iter(vec![ Ok(WsMessage::Ping(Bytes::new())), @@ -2233,8 +2224,9 @@ mod tests { while let Some(message) = recv_loop.next().await { let last_activity = state.get_last_activity(); drop(message); - tokio::time::advance(Duration::from_millis(1)); assert!(last_activity > prev_activity); + prev_activity = last_activity; + tokio::time::advance(Duration::from_millis(1)).await; } }