fix: close dropped subscriber sockets and ping idle ones

stream_labels only noticed a disconnect on its next send, so with few labels
closed subscribers piled up in CLOSE-WAIT until the server ran out of file
descriptors. It now watches the socket and pings every 30s.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
govandClaude Opus 5.5 committed 2026-11-06 15:18:01 -05:00
1 parent b4ac30c010
commit edf19507b4
1 file changed
+14 -1
+14 -1
View File
@@ -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,8 +139,18 @@ 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 {
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?;
@@ -147,6 +159,7 @@ async fn stream_labels(
// fell behind the channel; catch up from the database
Err(RecvError::Lagged(_)) => replay(socket, labeler, &mut last).await?,
Err(RecvError::Closed) => return Ok(()),
},
}
}
}