From 234f6c9c54406060ee2d781af9fd727e0d790650 Mon Sep 17 00:00:00 2001 From: Alessandro Asoni Date: Fri, 10 Jul 2026 15:13:39 +0200 Subject: [PATCH 1/2] 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 | 19 ++- crates/standalone/config.toml | 8 + 3 files changed, 177 insertions(+), 27 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 2e4b1bd62ad..f3ae0ec522a 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -366,8 +366,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`. /// @@ -461,7 +463,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(); @@ -475,6 +482,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( @@ -507,6 +515,15 @@ async fn ws_client_actor_inner( log::info!("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 @@ -518,8 +535,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 @@ -631,9 +652,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(); @@ -666,13 +692,31 @@ async fn ws_main_loop( break; }, - // Exit if we haven't heard from the client for too long. - _ = &mut idle_timer => { - log::warn!("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(); + 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; }, @@ -764,7 +808,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, @@ -820,7 +864,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 { // Get the next message from `ws`, or `None` if the stream is exhausted. @@ -1113,6 +1157,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, @@ -1120,9 +1165,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 @@ -1133,6 +1179,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, @@ -1156,6 +1203,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. // @@ -1262,6 +1319,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(); + } } }, @@ -1911,7 +1974,7 @@ mod tests { use std::{ future::{poll_fn, Future}, pin::Pin, - sync::atomic::AtomicUsize, + sync::{atomic::AtomicUsize, Mutex}, task::{Context, Poll}, }; @@ -1958,6 +2021,10 @@ mod tests { ActorState::new(Identity::ZERO, dummy_client_id(), config) } + 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); @@ -1988,7 +2055,7 @@ mod tests { let input = stream::iter(vec![Ok(WsMessage::Ping(Bytes::new()))]); pin_mut!(input); - 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); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2007,7 +2074,7 @@ mod tests { ]); pin_mut!(input); - 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); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2025,7 +2092,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(); @@ -2046,7 +2113,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(); @@ -2065,7 +2132,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(); @@ -2124,6 +2191,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2148,6 +2216,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2201,6 +2270,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2250,6 +2320,7 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), + dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2292,6 +2363,20 @@ mod tests { })); 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(); @@ -2302,7 +2387,7 @@ mod tests { ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), - drop, + unordered_tx, ) .await } @@ -2318,8 +2403,48 @@ 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()); + } + + #[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] @@ -2362,7 +2487,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(); @@ -2505,7 +2632,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 85c7d97a620..891fdd854c2 100644 --- a/crates/core/src/client/client_connection.rs +++ b/crates/core/src/client/client_connection.rs @@ -445,7 +445,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(); @@ -871,13 +871,20 @@ impl ClientConnection { let _gauge_guard = module_info.metrics.connected_clients.inc_scope(); module_info.metrics.ws_clients_spawned.inc(); - scopeguard::defer! { - let database_identity = module_info.database_identity; + // Runs only if this task is aborted (e.g. by `ClientConnectionSender::send` + // kicking a client whose outgoing queue is full) or if `fut` panics. + // Normal disconnects complete `fut`, defusing this guard; + // they are logged by the websocket actor. + let aborted_metric = module_info.metrics.ws_clients_aborted.clone(); + let abort_guard = scopeguard::guard((), move |()| { log::warn!("websocket connection aborted for client identity `{client_identity}` and database identity `{database_identity}`"); - module_info.metrics.ws_clients_aborted.inc(); - }; + aborted_metric.inc(); + }); - fut.await + fut.await; + + scopeguard::ScopeGuard::into_inner(abort_guard); + log::debug!("websocket connection ended for client identity `{client_identity}` and database identity `{database_identity}`"); }) .abort_handle(); 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 c81264bf17bef62c4e726c220005f7fe70bd45ba Mon Sep 17 00:00:00 2001 From: Jeffrey Dallatezza Date: Mon, 20 Jul 2026 06:12:21 -0700 Subject: [PATCH 2/2] Change idle timeout to track last activity (#5521) # Description of Changes 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 | 163 +++++++--------------- 1 file changed, 51 insertions(+), 112 deletions(-) diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index f3ae0ec522a..e5b64eed01f 100644 --- a/crates/client-api/src/routes/subscribe.rs +++ b/crates/client-api/src/routes/subscribe.rs @@ -4,7 +4,7 @@ use std::num::NonZeroUsize; use std::panic; use std::pin::{pin, Pin}; use std::sync::atomic::{AtomicBool, Ordering}; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::task::{Context, Poll}; use std::time::Duration; @@ -41,7 +41,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}; @@ -317,6 +317,8 @@ struct ActorState { config: WebSocketOptions, closed: AtomicBool, got_pong: AtomicBool, + // used to determine if the connection is idle. + last_activity: Arc>, } impl ActorState { @@ -327,6 +329,7 @@ impl ActorState { config, closed: AtomicBool::new(false), got_pong: AtomicBool::new(true), + last_activity: Arc::new(Mutex::new(Instant::now())), } } @@ -346,8 +349,20 @@ impl ActorState { self.got_pong.swap(false, Ordering::Relaxed) } - 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 } } @@ -462,15 +477,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 @@ -482,12 +488,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(); @@ -508,7 +512,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; @@ -519,11 +523,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 @@ -634,7 +633,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), @@ -657,6 +655,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); @@ -756,34 +756,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; } } @@ -808,7 +793,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, @@ -821,7 +805,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); @@ -864,7 +848,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 { // Get the next message from `ws`, or `None` if the stream is exhausted. @@ -909,7 +892,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); @@ -1165,10 +1148,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 @@ -1179,7 +1161,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, @@ -1203,16 +1184,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. // @@ -1319,12 +1290,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(); } }, @@ -2021,10 +1987,6 @@ mod tests { ActorState::new(Identity::ZERO, dummy_client_id(), config) } - 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); @@ -2055,7 +2017,7 @@ mod tests { let input = stream::iter(vec![Ok(WsMessage::Ping(Bytes::new()))]); pin_mut!(input); - let recv_loop = ws_recv_loop(state, Arc::new(idle_tx), input); + let recv_loop = ws_recv_loop(state, input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2074,7 +2036,7 @@ mod tests { ]); pin_mut!(input); - let recv_loop = ws_recv_loop(state, Arc::new(idle_tx), input); + let recv_loop = ws_recv_loop(state, input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2092,7 +2054,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(); @@ -2113,7 +2075,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(); @@ -2125,23 +2087,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] @@ -2191,7 +2152,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2216,7 +2176,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2270,7 +2229,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2320,7 +2278,6 @@ mod tests { messages_rx, unordered_rx, BsatnRowListBuilderPool::new(), - dummy_idle_tx(), ); pin_mut!(send_loop); @@ -2338,7 +2295,6 @@ mod tests { ws_main_loop( state.clone(), future::pending, - future::pending(), tokio::spawn(sleep(Duration::from_millis(10))), tokio::spawn(future::pending()), drop, @@ -2347,7 +2303,6 @@ mod tests { ws_main_loop( state, future::pending, - future::pending(), tokio::spawn(future::pending()), tokio::spawn(sleep(Duration::from_millis(10))), drop, @@ -2361,7 +2316,6 @@ mod tests { idle_timeout: Duration::from_millis(10), ..<_>::default() })); - 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 @@ -2384,7 +2338,6 @@ mod tests { ws_main_loop( state, future::pending, - ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), unordered_tx, @@ -2396,7 +2349,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); @@ -2414,7 +2367,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. @@ -2432,7 +2384,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, @@ -2454,7 +2405,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 = { @@ -2465,7 +2415,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(); } } } @@ -2476,9 +2426,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, @@ -2509,7 +2458,6 @@ mod tests { async fn main_loop_terminates_when_module_exits() { let state = Arc::new(dummy_actor_state()); - let (_idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); let unordered_tx = { let state = state.clone(); move |m| { @@ -2529,7 +2477,6 @@ mod tests { ws_main_loop( state.clone(), hotswap, - ws_idle_timer(idle_rx), // Pretend we received a close immediately after sending one. tokio::spawn(async move { loop { @@ -2632,15 +2579,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();