[api] close streams with 1012 when hydrant stops
83 unmodified lines84858687888990919239 unmodified lines13213313413513513613713813914014014114214314414514614714860 unmodified lines83 unmodified lines
// long enough for a busy database to finish what it's writing, short enough that a// supervisor's own stop timeout (systemd's is 90s) doesn't kill it firstconst SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(30);const DRAIN_TIMEOUT: Duration = Duration::from_secs(2);
fn main() -> miette::Result<()> { // config lives in env vars, so a stray argument (a `--version` probe, a typo) would // otherwise boot a full default instance without anyone noticing39 unmodified lines let debug_fut = app .enable_debug .then(|| hydrant.serve_debug(app.debug_port).boxed()) .unwrap_or_else(|| std::future::pending().boxed()); tokio::select! { let result = tokio::select! { r = hydrant.run()? => r, r = api_fut => r, r = debug_fut => r, _ = stop.notified() => Ok(()), } }; // the select dropped the listeners. reconnects go to the next hydrant hydrant.drain_streams(DRAIN_TIMEOUT).await; result });
// dropping the runtime would wait on blocking tasks that only `Hydrant::shutdown` ends runtime.shutdown_background();60 unmodified lines22 unmodified lines23242526272728293031102 unmodified linesExpand all134135136137138139140141142135 unmodified linesExpand all278279280281282283284285286351 unmodified linesExpand all22 unmodified lines filter::{FilterHandle, new_handle as new_filter_handle}, pds_meta::{PdsMeta, PdsMetaHandle, TierPolicy, new_handle as new_pds_handle}, resolver::Resolver, types::DocIdentity, util::{stop::Stop, throttle::Throttler}, util::{drain::Drain, stop::Stop, throttle::Throttler},};
const IDENTITY_RECONCILE_COOL_OFF: Duration = Duration::from_secs(30);
102 unmodified linesExpand all pub(crate) after_job_step: parking_lot::Mutex<Option<AfterChunk>>, pub get_repo_semaphore: Semaphore, pub maintenance_timeout: Duration, pub(crate) stop: Stop, pub(crate) drain: Drain,}
impl AppState { pub fn new(config: &Config) -> Result<Self> {135 unmodified linesExpand all after_job_step: parking_lot::Mutex::new(None), get_repo_semaphore: Semaphore::new(config.get_repo_concurrency_limit), maintenance_timeout: config.maintenance_timeout, stop: Stop::default(), drain: Drain::default(), throttler: Throttler::new(config.per_pds_concurrency), pds_daily_limit, }) }351 unmodified linesExpand all61 unmodified lines626364656667686970912 unmodified linesExpand all61 unmodified linesapi-boundary:src/api/stream.rs::router nu:tests/stream_cursor_replay.nuapi-boundary:src/api/ws.rs::WebSocketUpgrade::limits nu:tests/stream_ping.nuapi-boundary:src/api/ws.rs::WebSocketUpgrade::on_upgrade nu:tests/stream_ping.nuapi-boundary:src/api/ws.rs::control_frame_only_limits nu:tests/stream_ping.nuapi-boundary:src/api/ws.rs::restarting rust:src/api/ws.rs::a_drain_closes_an_idle_stream_as_a_restartapi-boundary:src/api/ws.rs::run_socket nu:tests/stream_ping.nuapi-boundary:src/api/xrpc/com_atproto_describe_repo.rs::handle rust:src/api/xrpc/com_atproto_describe_repo.rs::describe_repo_json_projection_contractapi-boundary:src/api/xrpc/count_records.rs::handle nu:tests/api_xrpc.nuapi-boundary:src/api/xrpc/describe_repo.rs::handle nu:tests/api_xrpc.nu912 unmodified linesExpand all114 unmodified linesExpand all115116117118119120114 unmodified linesExpand all- **live stream scope**: by default, cursor replay returns both live firehose events (`live: true`) and historical backfill / sync replay events (`live: false`), interleaved in commit order, while the live tail carries only live events, even after the subscriber falls behind and catches up from the database. filter them with the `wantedEventTypes` query parameter. asking for `historical` explicitly also delivers backfills that finish while the socket is open, from the live tail.- **ordering**: events are delivered in commit order, which is also `time_us` order, so a client that reconnects with its last `time_us` as the cursor misses nothing the socket would have sent.- **replay speed**: a cursor replay runs at most `HYDRANT_STREAM_REPLAY_SPEED` times the bandwidth the live stream would take (10 by default), or `HYDRANT_STREAM_REPLAY_FLOOR_BYTES` a second (about 80 Mbit/s) when that's more. events your filters skip count a little too, since hydrant still reads them. the live tail never waits on it. a single trusted consumer that has to catch up fast (such as a service resuming from its saved cursor) can leave jetstream out of `HYDRANT_THROTTLED_STREAMS` to replay uncapped.- **falling behind**: a subscriber that falls behind the live broadcast catches up from the database at its own pace instead of being disconnected.- **restarts**: when hydrant stops, it closes the connection with close code `1012` (service restart).- **slow consumers**: if a socket's send buffer remains full for longer than the configured timeout, the server sends a `{"type":"error","error":"ConsumerTooSlow"}` message and drops the connection.129 unmodified linesExpand all13013113213313413513613713817 unmodified lines129 unmodified linesExpand all- client text and binary frames are not accepted; hydrant closes the connection when it receives one.- when `/stream` is in `HYDRANT_THROTTLED_STREAMS` (it isn't by default), a replay from a cursor runs at most `HYDRANT_STREAM_REPLAY_SPEED` times the bandwidth the live stream would take (10 by default), or `HYDRANT_STREAM_REPLAY_FLOOR_BYTES` a second (about 80 Mbit/s) when that's more. the live tail never waits on it.- a subscriber that falls behind the live broadcast catches up from the database at its own pace. `record` events are never skipped, but `identity` and `account` events it fell behind on are not replayed.- if hydrant can't read the next events from the database, it closes the connection rather than skip them. reconnect with the last `id` you received as the cursor.- when hydrant stops, it closes the connection with close code `1012` (service restart). reconnect the same way. under systemd socket activation, the reconnect waits for the next hydrant instead of being refused.- if sending blocks for `HYDRANT_STREAM_SEND_TIMEOUT` (30 seconds by default), hydrant sends an error text frame and closes the connection:
```json{17 unmodified lines1 unmodified line2345667891094 unmodified lines105106107108109110111112113114115116117118119120121122123124262 unmodified linesExpand all1 unmodified line
use axum::extract::State;use axum::http::{HeaderMap, StatusCode};use axum::response::{IntoResponse, Response};use crate::api::ws::{Message, WebSocket, WebSocketUpgrade};use crate::api::ws::{Message, WebSocket, WebSocketUpgrade, restarting};use futures::{SinkExt, StreamExt};use serde::Deserialize;use tokio::time::{MissedTickBehavior, interval, timeout};use tracing::{debug, warn};94 unmodified lines } } }
let mut draining = hydrant.state.drain.subscribe(); let mut events = hydrant.subscribe_jetstream_network(cursor, options.clone()); let mut ping_timer = interval(PING_INTERVAL); ping_timer.set_missed_tick_behavior(MissedTickBehavior::Delay); ping_timer.tick().await;
loop { tokio::select! { _ = draining.started() => { let _ = timeout(CLOSE_TIMEOUT, sink.send(restarting())).await; break; } inbound = ws_recv.next() => match inbound { Some(Ok(m)) if m.is_close() => break, Some(Ok(m)) if m.is_ping() => { if let Err(err) = sink.send(Message::pong(m.into_payload())).await {262 unmodified linesExpand all47 unmodified lines48495051525354555619 unmodified lines76777879808182838447 unmodified lines47 unmodified lines}
async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: StreamQuery) { let send_timeout = hydrant.stream_send_timeout(); let draining = hydrant.state.drain.subscribe(); let frames = hydrant.subscribe_frames(query.cursor); run_socket( socket, frames,19 unmodified lines ), }); Some(Message::text(json.to_string())) }, draining, ) .await;}
47 unmodified lines9 unmodified lines1011121314151617181920212223100 unmodified lines1241251261271281291301311321331341351363 unmodified lines1401411421431441451461471485 unmodified lines1541551561571581591601611623 unmodified lines1661671681691701711721731748 unmodified lines183184185186187188189190191192193194108 unmodified linesExpand all30330430530630730830931031131217 unmodified lines33033133233333433533633733833934010 unmodified lines35135235335435535635735835936036136236336436536636736833936937037137237337437537637752 unmodified lines4304314324334344354364374384394404414424434444454464474484494504514524534544559 unmodified linesuse bytes::Bytes;use futures::{FutureExt, Sink, SinkExt, Stream, StreamExt};use hyper_util::rt::TokioIo;use tokio::time::{MissedTickBehavior, interval, timeout};use tokio_websockets::CloseCode;pub(crate) use tokio_websockets::{Limits, Message};use tracing::{debug, warn};
use crate::util::drain::Draining;
pub(crate) type WebSocket = tokio_websockets::WebSocketStream<TokioIo<hyper::upgrade::Upgraded>>;
pub(crate) struct WebSocketUpgrade { limits: Limits,100 unmodified linespub(super) fn control_frame_only_limits() -> Limits { Limits::default().max_payload_len(Some(CONTROL_FRAME_PAYLOAD_LIMIT))}
/// 1012 is what makes a client reconnect now instead of backing offpub(super) fn restarting() -> Message { Message::close(Some(CloseCode::SERVICE_RESTART), "hydrant is restarting")}
pub(super) enum WsAction { Send(Message), Close(Option<Message>),}3 unmodified lines events: S, to_action: F, send_timeout: Duration, slow_consumer: G, draining: Draining,) where S: Stream + Unpin, F: FnMut(S::Item) -> WsAction, G: FnOnce(Duration) -> Option<Message>,5 unmodified lines events, to_action, send_timeout, slow_consumer, draining, ) .await;}
3 unmodified lines mut events: S, mut to_action: F, send_timeout: Duration, slow_consumer: G, mut draining: Draining,) where Si: Sink<Message, Error = E> + Unpin, R: Stream<Item = Result<Message, E>> + Unpin, E: Display,8 unmodified lines let mut slow_consumer = Some(slow_consumer);
loop { tokio::select! { _ = draining.started() => { let _ = timeout(CLOSE_TIMEOUT, sink.send(restarting())).await; break; } 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");108 unmodified linesExpand all use std::pin::Pin; use std::sync::{Arc, Mutex}; use std::task::{Context, Poll};
use crate::util::drain::Drain;
/// a client socket that keeps what it was sent and how many flushes wrote it, or that stops /// reading, so no flush ever finishes #[derive(Clone, Default)] struct Client {17 unmodified lines
fn start_send(self: Pin<&mut Self>, msg: Message) -> Result<(), Self::Error> { if let Some(text) = msg.as_text() { self.sent.lock().unwrap().push(text.to_owned()); } else if let Some((code, reason)) = msg.as_close() { let code = u16::from(code); self.sent.lock().unwrap().push(format!("close {code} {reason}")); } Ok(()) }
10 unmodified lines } }
async fn run(client: Client, events: Vec<WsAction>, send_timeout: Duration) { let drain = Drain::default(); serve_until(client, futures::stream::iter(events), send_timeout, &drain).await; }
async fn serve_until<S: Stream<Item = WsAction> + Unpin>( client: Client, events: S, send_timeout: Duration, drain: &Drain, ) { let quiet = futures::stream::pending::<Result<Message, std::convert::Infallible>>(); serve( client, quiet, futures::stream::iter(events), events, |action| action, send_timeout, |_| Some(Message::text("too slow")), drain.subscribe(), ) .await; }
52 unmodified lines (0..10).map(|n| n.to_string()).collect::<Vec<_>>() ); }
#[tokio::test] async fn a_drain_closes_an_idle_stream_as_a_restart() { let client = Client::default(); let drain = Drain::default(); let events = futures::stream::iter((0..3).map(text)).chain(futures::stream::pending()); let sent_all = async { while client.texts().len() < 3 { tokio::time::sleep(Duration::from_millis(5)).await; } drain.drain(Duration::from_secs(5)).await; }; let served = serve_until(client.clone(), events, Duration::from_secs(5), &drain); tokio::time::timeout(Duration::from_secs(5), futures::future::join(served, sent_all)) .await .expect("the drain must end the stream");
assert_eq!( client.texts(), ["0", "1", "2", "close 1012 hydrant is restarting"] ); }}219 unmodified linesExpand all220221222223224225226227228229230231232158 unmodified linesExpand all219 unmodified linesExpand all let state = self.state.clone(); async move { crate::api::serve_debug(state, port).await } }
/// call after the api stops accepting and before [`shutdown`](Self::shutdown). pub async fn drain_streams(&self, timeout: Duration) { self.state.drain.drain(timeout).await; }
/// stops hydrant's own threads and waits until its database is closed, which /// gives its memory back and lets the same path be opened again. /// /// the tasks [`run`](Self::run) and [`serve`](Self::serve) spawned live on your158 unmodified linesExpand all12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273use std::time::Duration;
use tokio::sync::watch;
pub struct Drain(watch::Sender<bool>);
pub struct Draining(watch::Receiver<bool>);
impl Default for Drain { fn default() -> Self { Self(watch::Sender::new(false)) }}
impl Drain { pub fn subscribe(&self) -> Draining { Draining(self.0.subscribe()) }
pub async fn drain(&self, timeout: Duration) { self.0.send_replace(true); let _ = tokio::time::timeout(timeout, self.0.closed()).await; }}
impl Draining { pub async fn started(&mut self) { // a dropped drain never starts and this pends forever instead of returning if self.0.wait_for(|draining| *draining).await.is_err() { std::future::pending::<()>().await; } }}
#[cfg(test)]mod tests { use super::*; use std::time::Instant;
#[tokio::test] async fn drain_waits_for_subscribers_to_go() { let drain = Drain::default(); let mut draining = drain.subscribe(); let subscriber = tokio::spawn(async move { draining.started().await; tokio::time::sleep(Duration::from_millis(50)).await; }); let started = Instant::now(); drain.drain(Duration::from_secs(5)).await; assert!(started.elapsed() >= Duration::from_millis(50)); assert!(started.elapsed() < Duration::from_secs(5)); subscriber.await.unwrap(); }
#[tokio::test] async fn drain_gives_up_on_a_subscriber_that_never_goes() { let drain = Drain::default(); let _stuck = drain.subscribe(); let started = Instant::now(); drain.drain(Duration::from_millis(50)).await; assert!(started.elapsed() < Duration::from_secs(5)); }
#[tokio::test] async fn a_subscriber_after_the_drain_started_sees_it_right_away() { let drain = Drain::default(); drain.drain(Duration::ZERO).await; let mut late = drain.subscribe(); tokio::time::timeout(Duration::from_secs(1), late.started()) .await .expect("a late subscriber must see the drain"); }}10 unmodified lines111213141516171819507 unmodified linesExpand all10 unmodified lines
use crate::{db::types::DidKey, types::RepoStatus};
pub mod drain;pub mod stop;pub mod throttle;
pub(crate) mod hex {507 unmodified linesExpand all27 unmodified lines28293031323334353622 unmodified lines596061626364656627 unmodified lines}
async fn handle_socket(socket: WebSocket, hydrant: Hydrant, query: SubscribeReposQuery) { let send_timeout = hydrant.stream_send_timeout(); let draining = hydrant.state.drain.subscribe(); let stream = hydrant.subscribe_repos_network(query.cursor); run_socket( socket, stream,22 unmodified lines None } } }, draining, ) .await;}