From d8068f5632814cdc08427bdc00074861ec2db9ab Mon Sep 17 00:00:00 2001 From: "oyster.cafe" Date: Thu, 7 May 2026 11:37:28 +0000 Subject: [PATCH] [api] websocket pong & better ping, round 2 Hydrant should ping others! :P And also, we shouldn't let others pongtimeout, we should pong. Sorry I snuck in a little flake fix too. --- flake.nix | 1 + src/api/mod.rs | 2 + src/api/stream.rs | 79 +++++++++++---------------- src/api/ws.rs | 94 +++++++++++++++++++++++++++++++++ src/api/xrpc/subscribe_repos.rs | 76 ++++++++++++-------------- src/ingest/stream.rs | 2 +- tests/authenticated_stream.nu | 2 +- tests/stream.nu | 4 +- tests/stream_ping.nu | 52 ++++++++++++++++++ tests/subscribe_repos_ping.nu | 54 +++++++++++++++++++ 10 files changed, 271 insertions(+), 95 deletions(-) create mode 100644 src/api/ws.rs create mode 100644 tests/stream_ping.nu create mode 100644 tests/subscribe_repos_ping.nu diff --git a/flake.nix b/flake.nix index 46b0f7f..8cb41e7 100644 --- a/flake.nix +++ b/flake.nix @@ -66,6 +66,7 @@ cmake websocat http-nu + nushell clang wild psmisc diff --git a/src/api/mod.rs b/src/api/mod.rs index 2f1e7e4..976328f 100644 --- a/src/api/mod.rs +++ b/src/api/mod.rs @@ -19,6 +19,8 @@ mod repos; mod stats; #[cfg(feature = "indexer_stream")] mod stream; +#[cfg(any(feature = "relay", feature = "indexer_stream"))] +mod ws; mod xrpc; pub async fn serve(hydrant: Hydrant, binds: ApiBinds) -> miette::Result<()> { diff --git a/src/api/stream.rs b/src/api/stream.rs index df7db6d..e88783c 100644 --- a/src/api/stream.rs +++ b/src/api/stream.rs @@ -6,10 +6,11 @@ use axum::{ response::IntoResponse, }; use axum_tws::{Message, WebSocket, WebSocketUpgrade}; -use futures::StreamExt; use serde::Deserialize; use tracing::error; +use super::ws::{WsAction, run_socket}; + pub fn router() -> Router { Router::new().route("/", get(handle_stream)) } @@ -27,59 +28,41 @@ pub async fn handle_stream( ws.on_upgrade(move |socket| handle_socket(socket, hydrant, query)) } -async fn handle_socket(mut socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { +async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { let send_timeout = hydrant.stream_send_timeout(); - let mut stream = hydrant.subscribe(query.cursor); - - while let Some(item) = stream.next().await { - let evt = match item { - Ok(evt) => evt, + let events = hydrant.subscribe(query.cursor); + run_socket( + socket, + events, + |item| match item { + Ok(evt) => match serde_json::to_string(&evt) { + Ok(json) => WsAction::Send(Message::text(json)), + Err(e) => { + error!(err = %e, "failed to serialize event"); + WsAction::Skip + } + }, Err(err) => { let json = serde_json::json!({ "type": "error", "error": err.code(), "message": err.to_string(), }); - let _ = tokio::time::timeout( - send_timeout, - socket.send(Message::text(json.to_string())), - ) - .await; - let _ = - tokio::time::timeout(std::time::Duration::from_secs(1), socket.close()).await; - break; - } - }; - - match serde_json::to_string(&evt) { - Ok(json) => { - match tokio::time::timeout(send_timeout, socket.send(Message::text(json))).await { - Ok(Ok(())) => {} - Ok(Err(_)) => break, - Err(_) => { - let err = serde_json::json!({ - "type": "error", - "error": "ConsumerTooSlow", - "message": format!( - "stream socket send blocked for at least {} seconds", - send_timeout.as_secs() - ), - }); - let _ = tokio::time::timeout( - std::time::Duration::from_secs(1), - socket.send(Message::text(err.to_string())), - ) - .await; - let _ = - tokio::time::timeout(std::time::Duration::from_secs(1), socket.close()) - .await; - break; - } - } - } - Err(e) => { - error!(err = %e, "failed to serialize event"); + WsAction::Close(Some(Message::text(json.to_string()))) } - } - } + }, + send_timeout, + |timeout_dur| { + let json = serde_json::json!({ + "type": "error", + "error": "ConsumerTooSlow", + "message": format!( + "stream socket send blocked for at least {} seconds", + timeout_dur.as_secs() + ), + }); + Some(Message::text(json.to_string())) + }, + ) + .await; } diff --git a/src/api/ws.rs b/src/api/ws.rs new file mode 100644 index 0000000..8cc6f84 --- /dev/null +++ b/src/api/ws.rs @@ -0,0 +1,94 @@ +use std::time::Duration; + +use axum_tws::{Message, WebSocket}; +use bytes::Bytes; +use futures::{SinkExt, Stream, StreamExt}; +use tokio::time::{MissedTickBehavior, interval, timeout}; +use tracing::{debug, warn}; + +const PING_INTERVAL: Duration = Duration::from_secs(30); +const CLOSE_TIMEOUT: Duration = Duration::from_secs(1); + +pub(super) enum WsAction { + Send(Message), + Skip, + Close(Option), +} + +pub(super) async fn run_socket( + socket: WebSocket, + mut events: S, + mut to_action: F, + send_timeout: Duration, + slow_consumer: G, +) where + S: Stream + Unpin, + F: FnMut(S::Item) -> WsAction, + G: FnOnce(Duration) -> Option, +{ + let (mut sink, mut ws_recv) = socket.split(); + + let mut ping_timer = interval(PING_INTERVAL); + ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); + ping_timer.tick().await; + + let mut slow_consumer = Some(slow_consumer); + + loop { + tokio::select! { + inbound = ws_recv.next() => match inbound { + Some(Ok(m)) if m.is_close() => break, + Some(Ok(m)) if m.is_text() || m.is_binary() => { + debug!("client sent unsolicited data frame, closing"); + break; + } + Some(Ok(m)) if m.is_ping() => { + if let Err(err) = sink.send(Message::pong(m.into_payload())).await { + warn!(err = %err, "ws pong send error"); + break; + } + } + Some(Ok(_)) => {} + Some(Err(e)) => { + warn!(err = %e, "ws recv error"); + break; + } + None => break, + }, + evt = events.next() => match evt { + Some(item) => match to_action(item) { + WsAction::Skip => {} + WsAction::Send(msg) => match timeout(send_timeout, sink.send(msg)).await { + Ok(Ok(())) => {} + Ok(Err(err)) => { + warn!(err = %err, "ws send error"); + break; + } + Err(_) => { + if let Some(final_msg) = + slow_consumer.take().and_then(|f| f(send_timeout)) + { + let _ = timeout(CLOSE_TIMEOUT, sink.send(final_msg)).await; + } + break; + } + }, + WsAction::Close(final_msg) => { + if let Some(m) = final_msg { + let _ = timeout(send_timeout, sink.send(m)).await; + } + break; + } + }, + None => break, + }, + _ = ping_timer.tick() => { + if let Err(err) = sink.send(Message::ping(Bytes::new())).await { + warn!(err = %err, "ws ping send error"); + break; + } + } + } + } + let _ = timeout(CLOSE_TIMEOUT, sink.close()).await; +} diff --git a/src/api/xrpc/subscribe_repos.rs b/src/api/xrpc/subscribe_repos.rs index c1703a3..d4692bf 100644 --- a/src/api/xrpc/subscribe_repos.rs +++ b/src/api/xrpc/subscribe_repos.rs @@ -3,10 +3,10 @@ use axum::{ response::IntoResponse, }; use axum_tws::{Message, WebSocket, WebSocketUpgrade}; -use futures::StreamExt; use serde::Deserialize; use tracing::error; +use crate::api::ws::{WsAction, run_socket}; use crate::control::{Hydrant, RelayStreamError}; use crate::ingest::stream::encode_error_frame; @@ -23,48 +23,38 @@ pub async fn handle( ws.on_upgrade(move |socket| handle_socket(socket, hydrant, query)) } -async fn handle_socket(mut socket: WebSocket, hydrant: Hydrant, query: SubscribeReposQuery) { +async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: SubscribeReposQuery) { let send_timeout = hydrant.stream_send_timeout(); - let mut stream = hydrant.subscribe_repos(query.cursor); - - while let Some(item) = stream.next().await { - let frame = match item { - Ok(frame) => frame, - Err(err) => { - send_error_frame(&mut socket, send_timeout, &err).await; - break; - } - }; - - match tokio::time::timeout(send_timeout, socket.send(Message::binary(frame))).await { - Ok(Ok(())) => {} - Ok(Err(_)) => break, - Err(_) => { - let err = RelayStreamError::ConsumerTooSlow { - reason: format!( - "relay stream socket send blocked for at least {} seconds", - send_timeout.as_secs() - ), - }; - send_error_frame(&mut socket, std::time::Duration::from_secs(1), &err).await; - break; + let stream = hydrant.subscribe_repos(query.cursor); + run_socket( + socket, + stream, + |item| match item { + Ok(frame) => WsAction::Send(Message::binary(frame)), + Err(err) => match encode_error_frame(err.code(), Some(&err.to_string())) { + Ok(frame) => WsAction::Close(Some(Message::binary(frame))), + Err(e) => { + error!(err = %e, "failed to encode relay stream error frame"); + WsAction::Close(None) + } + }, + }, + send_timeout, + |timeout_dur| { + let err = RelayStreamError::ConsumerTooSlow { + reason: format!( + "relay stream socket send blocked for at least {} seconds", + timeout_dur.as_secs() + ), + }; + match encode_error_frame(err.code(), Some(&err.to_string())) { + Ok(frame) => Some(Message::binary(frame)), + Err(e) => { + error!(err = %e, "failed to encode relay stream error frame"); + None + } } - } - } -} - -async fn send_error_frame( - socket: &mut WebSocket, - timeout: std::time::Duration, - err: &RelayStreamError, -) { - match encode_error_frame(err.code(), Some(&err.to_string())) { - Ok(frame) => { - let _ = tokio::time::timeout(timeout, socket.send(Message::binary(frame))).await; - } - Err(e) => { - error!(err = %e, "failed to encode relay stream error frame"); - } - } - let _ = tokio::time::timeout(std::time::Duration::from_secs(1), socket.close()).await; + }, + ) + .await; } diff --git a/src/ingest/stream.rs b/src/ingest/stream.rs index 29a0d0d..6ec8b91 100644 --- a/src/ingest/stream.rs +++ b/src/ingest/stream.rs @@ -110,7 +110,6 @@ impl FirehoseStream { } return Ok(bytes); } - msg if msg.is_ping() => self.ws.send(WsMsg::pong(msg.into_payload())).await?, // if ws closed treat it as an error, since why would a host close the stream?? // TODO: treat hosts that return these as offline ?????? msg if msg.is_close() => { @@ -120,6 +119,7 @@ impl FirehoseStream { ); return Err(FirehoseError::StreamClosed { code, reason }); } + msg if msg.is_ping() => self.ws.send(WsMsg::pong(msg.into_payload())).await?, msg if msg.is_pong() => continue, x => { trace!(msg = ?x, "host sent unexpected message"); diff --git a/tests/authenticated_stream.nu b/tests/authenticated_stream.nu index 7373ef7..5837435 100644 --- a/tests/authenticated_stream.nu +++ b/tests/authenticated_stream.nu @@ -26,7 +26,7 @@ def run-auth-test [did: string, password: string, pds_url: string, relays: strin let output_file = $"($db_path)/stream_output.txt" print $"starting stream listener -> ($output_file)" # use websocat to capture output. - let stream_pid = (bash -c $"websocat '($ws_url)' > '($output_file)' & echo $!" | str trim | into int) + let stream_pid = (bash -c $"websocat -n '($ws_url)' > '($output_file)' & echo $!" | str trim | into int) print $"listener pid: ($stream_pid)" # 4. add repo to hydrant (backfill trigger) diff --git a/tests/stream.nu b/tests/stream.nu index 8326ab0..d155273 100644 --- a/tests/stream.nu +++ b/tests/stream.nu @@ -25,7 +25,7 @@ def main [] { print $"starting stream listener -> ($live_output)" # start websocat in background to capture live events (no cursor = live only) - let stream_pid = (bash -c $"websocat '($ws_url)' > '($live_output)' 2>&1 & echo $!" | str trim | into int) + let stream_pid = (bash -c $"websocat -n '($ws_url)' > '($live_output)' 2>&1 & echo $!" | str trim | into int) print $"stream listener pid: ($stream_pid)" sleep 1sec @@ -77,7 +77,7 @@ def main [] { # use same approach as test 1: background process with file output # cursor=0 replays from the beginning (no cursor = live-tail only) print "starting historical stream listener..." - let history_pid = (bash -c $"websocat '($ws_url)?cursor=0' > '($history_output)' 2>&1 & echo $!" | str trim | into int) + let history_pid = (bash -c $"websocat -n '($ws_url)?cursor=0' > '($history_output)' 2>&1 & echo $!" | str trim | into int) print $"history listener pid: ($history_pid)" # wait for events to be streamed (should be fast for historical replay) diff --git a/tests/stream_ping.nu b/tests/stream_ping.nu new file mode 100644 index 0000000..2c92fd3 --- /dev/null +++ b/tests/stream_ping.nu @@ -0,0 +1,52 @@ +#!/usr/bin/env nu +use common.nu * + +def main [] { + let port = resolve-test-port 3010 + let url = $"http://localhost:($port)" + let ws_url = $"ws://localhost:($port)/stream" + let db_path = (mktemp -d -t hydrant_stream_ping_test.XXXXXX) + + print "testing ping/pong handling on /stream..." + print $"database path: ($db_path)" + + let binary = build-hydrant + let instance = start-hydrant $binary $db_path $port + + mut passed = false + + if (wait-for-api $url) { + let log_file = $"($db_path)/ws.log" + let pid_file = $"($db_path)/ws.pid" + + bash -c $"websocat -n --ping-interval 1 --ping-timeout 4 '($ws_url)' > '($log_file)' 2>&1 & echo $! > '($pid_file)'" + sleep 200ms + let ws_pid = (open $pid_file | str trim | into int) + print $"websocat pid: ($ws_pid)" + + sleep 6sec + + let alive = (do { ^kill -0 $ws_pid } | complete | get exit_code) == 0 + if $alive { + print "ping/pong test PASSED: connection alive after 6s of pings" + $passed = true + try { kill $ws_pid } + } else { + print "ping/pong test FAILED: websocat exited, pong likely not received in time" + try { open $log_file | print } + } + } else { + print "api failed to start." + } + + let hydrant_pid = $instance.pid + print $"stopping hydrant - pid: ($hydrant_pid)..." + try { kill $hydrant_pid } + + if $passed { + print "=== TEST PASSED ===" + } else { + print "=== TEST FAILED ===" + exit 1 + } +} diff --git a/tests/subscribe_repos_ping.nu b/tests/subscribe_repos_ping.nu new file mode 100644 index 0000000..2d6dace --- /dev/null +++ b/tests/subscribe_repos_ping.nu @@ -0,0 +1,54 @@ +#!/usr/bin/env nu +use common.nu * + +def main [] { + let port = resolve-test-port 3011 + let url = $"http://localhost:($port)" + let ws_url = $"ws://localhost:($port)/xrpc/com.atproto.sync.subscribeRepos" + let db_path = (mktemp -d -t hydrant_subscribe_repos_ping_test.XXXXXX) + + print "testing ping/pong handling on /xrpc/com.atproto.sync.subscribeRepos..." + print $"database path: ($db_path)" + + let binary = build-hydrant-relay + let instance = (with-env { HYDRANT_RELAY: "true" } { + start-hydrant $binary $db_path $port + }) + + mut passed = false + + if (wait-for-api $url) { + let log_file = $"($db_path)/ws.log" + let pid_file = $"($db_path)/ws.pid" + + bash -c $"websocat -n --ping-interval 1 --ping-timeout 4 '($ws_url)' > '($log_file)' 2>&1 & echo $! > '($pid_file)'" + sleep 200ms + let ws_pid = (open $pid_file | str trim | into int) + print $"websocat pid: ($ws_pid)" + + sleep 6sec + + let alive = (do { ^kill -0 $ws_pid } | complete | get exit_code) == 0 + if $alive { + print "ping/pong test PASSED: connection alive after 6s of pings" + $passed = true + try { kill $ws_pid } + } else { + print "ping/pong test FAILED: websocat exited, pong likely not received in time" + try { open $log_file | print } + } + } else { + print "api failed to start." + } + + let hydrant_pid = $instance.pid + print $"stopping hydrant - pid: ($hydrant_pid)..." + try { kill $hydrant_pid } + + if $passed { + print "=== TEST PASSED ===" + } else { + print "=== TEST FAILED ===" + exit 1 + } +} -- 2.51.2