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();