diff --git a/crates/client-api/src/routes/subscribe.rs b/crates/client-api/src/routes/subscribe.rs index 2e4b1bd62ad..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 } } @@ -366,8 +381,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`. /// @@ -460,10 +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. - let (idle_tx, idle_rx) = watch::channel(state.next_idle_deadline()); - 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 @@ -479,7 +492,6 @@ async fn ws_client_actor_inner( // 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(); @@ -500,13 +512,17 @@ 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; 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); + /// The main `select!` loop of the websocket client actor. /// /// > This function is defined standalone with generic parameters so that its @@ -518,8 +534,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 @@ -613,7 +633,6 @@ async fn ws_client_actor_inner( 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), @@ -631,9 +650,16 @@ 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; + + let idle_timer = state.idle_timer(); 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; }, @@ -712,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; } } @@ -764,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: watch::Sender, client_closed_metric: IntGauge, message_handler: impl Fn(DataMessage, Instant) -> MessageHandler, unordered_tx: mpsc::UnboundedSender, @@ -777,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); @@ -820,7 +848,6 @@ async fn ws_recv_task( /// state are dropped. fn ws_recv_loop( state: Arc, - idle_tx: watch::Sender, mut ws: impl Stream> + Unpin, ) -> impl Stream { // Get the next message from `ws`, or `None` if the stream is exhausted. @@ -865,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); @@ -1113,6 +1140,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, @@ -1262,6 +1290,7 @@ async fn ws_send_loop_inner( log::warn!("error sending frame: {e:#}"); break 'outer; } + state.record_activity(); } }, @@ -1911,7 +1940,7 @@ mod tests { use std::{ future::{poll_fn, Future}, pin::Pin, - sync::atomic::AtomicUsize, + sync::{atomic::AtomicUsize, Mutex}, task::{Context, Poll}, }; @@ -1988,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, idle_tx, input); + let recv_loop = ws_recv_loop(state, input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2007,7 +2036,7 @@ mod tests { ]); pin_mut!(input); - let recv_loop = ws_recv_loop(state, idle_tx, input); + let recv_loop = ws_recv_loop(state, input); pin_mut!(recv_loop); assert_matches!(recv_loop.next().await, Some(ClientMessage::Ping(_))); @@ -2025,7 +2054,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(), &mut input); pin_mut!(recv_loop); state.close(); @@ -2046,7 +2075,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(), &mut input); pin_mut!(recv_loop); state.close(); @@ -2058,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, 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] @@ -2267,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, @@ -2276,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, @@ -2290,7 +2316,20 @@ 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 + // 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({ @@ -2299,10 +2338,9 @@ mod tests { ws_main_loop( state, future::pending, - ws_idle_timer(idle_rx), tokio::spawn(future::pending()), tokio::spawn(future::pending()), - drop, + unordered_tx, ) .await } @@ -2311,15 +2349,53 @@ 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); 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() + })); + + // 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, + 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] @@ -2329,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 = { @@ -2340,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(); } } } @@ -2351,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, @@ -2362,7 +2436,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(); @@ -2382,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| { @@ -2402,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 { 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