diff --git a/src/server.rs b/src/server.rs index ac5d9fe..026727f 100644 --- a/src/server.rs +++ b/src/server.rs @@ -1,4 +1,5 @@ use std::sync::Arc; +use std::time::Duration; use anyhow::Result; use axum::extract::ws::{Message, WebSocket, WebSocketUpgrade}; @@ -17,6 +18,7 @@ use crate::db::{LabelQuery, SeqLabel}; use crate::labeler::Labeler; const REPLAY_BATCH: i64 = 500; +const PING_INTERVAL: Duration = Duration::from_secs(30); fn bad_request(message: &str) -> GenericXrpcErrorResponse { GenericXrpcErrorResponse::new(StatusCode::BAD_REQUEST, "InvalidRequest", Some(message)) @@ -137,16 +139,27 @@ async fn stream_labels( replay(socket, labeler, &mut last).await?; } + // Labels are rare, so also watch the socket: otherwise a subscriber that + // hangs up is only noticed on the next send, and its connection leaks. + let mut ping = tokio::time::interval(PING_INTERVAL); loop { - match live.recv().await { - Ok(row) if row.seq > last => { - last = row.seq; - socket.send(labels_frame(row)?).await?; - } - Ok(_) => {} - // fell behind the channel; catch up from the database - Err(RecvError::Lagged(_)) => replay(socket, labeler, &mut last).await?, - Err(RecvError::Closed) => return Ok(()), + tokio::select! { + incoming = socket.recv() => match incoming { + None | Some(Err(_)) | Some(Ok(Message::Close(_))) => return Ok(()), + Some(Ok(_)) => {} + }, + // keeps idle connections from being dropped by peers and proxies + _ = ping.tick() => socket.send(Message::Ping(Default::default())).await?, + published = live.recv() => match published { + Ok(row) if row.seq > last => { + last = row.seq; + socket.send(labels_frame(row)?).await?; + } + Ok(_) => {} + // fell behind the channel; catch up from the database + Err(RecvError::Lagged(_)) => replay(socket, labeler, &mut last).await?, + Err(RecvError::Closed) => return Ok(()), + }, } } }