diff --git a/Cargo.lock b/Cargo.lock index 286a13f..c360f47 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1591,6 +1591,7 @@ dependencies = [ "serde", "thiserror 2.0.18", "tokio", + "tokio-stream", "tokio-util", "tracing", ] @@ -3778,6 +3779,17 @@ dependencies = [ "tokio", ] +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", +] + [[package]] name = "tokio-tungstenite" version = "0.24.0" diff --git a/Cargo.toml b/Cargo.toml index 01df49a..9dd0ca5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,7 @@ star-lite = "0.2.0" star-repo = "0.2.0" thiserror = "2.0.18" tokio = "1.51.0" +tokio-stream = "0.1" tokio-util = { version = "0.7", features = ["io"] } tracing = "0.1.44" tracing-subscriber = { version = "0.3.23", features = ["env-filter"] } diff --git a/hubble-sync/Cargo.toml b/hubble-sync/Cargo.toml index e81298f..d030fca 100644 --- a/hubble-sync/Cargo.toml +++ b/hubble-sync/Cargo.toml @@ -13,5 +13,6 @@ serde = { workspace = true } reqwest = { workspace = true } thiserror = { workspace = true } tokio = { workspace = true, features = ["macros", "io-util", "rt", "time"] } +tokio-stream = { workspace = true } tokio-util = { workspace = true } tracing = { workspace = true } diff --git a/hubble-sync/src/cancel.rs b/hubble-sync/src/cancel.rs new file mode 100644 index 0000000..f68a40c --- /dev/null +++ b/hubble-sync/src/cancel.rs @@ -0,0 +1,38 @@ +//! cancellation token nice-extensions trait + +use tokio::time::{Duration, error::Elapsed, sleep, timeout}; +use tokio_util::sync::CancellationToken; + +/// nice helpers on `tokio_util::sync::CancellationToken`s +pub trait CancelExt { + /// alias for CancellationToken::run_until_cancelled + /// + /// returns None if the token cancels before the future completes + fn run(&self, fut: F) -> impl Future>; + /// sleep that ends early if if cancelled + /// + /// returns `true` when completed after the full sleep completed + fn sleep(&self, d: Duration) -> impl Future; + /// runs the future wrapped in a timeout wrapped in a token canceller + /// + /// returns None if the token cancels before the timeout or future completes + /// returns Some(Err(elapsed)) if the timeout completes before the future + /// otherwise returns Some(Ok(F::Output)) if the future gets to finish + fn timeout( + &self, + d: Duration, + fut: F, + ) -> impl Future>>; +} + +impl CancelExt for CancellationToken { + async fn run(&self, f: F) -> Option { + self.run_until_cancelled(f).await + } + async fn sleep(&self, d: Duration) -> bool { + self.run_until_cancelled(sleep(d)).await.is_some() + } + async fn timeout(&self, d: Duration, fut: F) -> Option> { + self.run_until_cancelled(timeout(d, fut)).await + } +} diff --git a/hubble-sync/src/firehose/event.rs b/hubble-sync/src/firehose/event.rs index 25048c8..e70401b 100644 --- a/hubble-sync/src/firehose/event.rs +++ b/hubble-sync/src/firehose/event.rs @@ -2,9 +2,10 @@ use std::sync::Arc; -use jacquard_api::com_atproto::sync::subscribe_repos::{Account, Commit, Identity, Info, Sync}; +use jacquard_api::com_atproto::sync::subscribe_repos::{Account, Commit, Identity, Sync}; use tokio::sync::mpsc; +use super::Seq; use crate::Host; /// A DID-bound event from the firehose @@ -15,7 +16,7 @@ use crate::Host; #[derive(Debug)] pub struct FirehoseEvent { pub host: Arc, - pub seq: i64, + pub seq: Seq, pub payload: FirehosePayload, pub ack: FirehoseAck, } @@ -33,12 +34,12 @@ pub enum FirehosePayload { /// this lets the persisted cursor update to the next processed seq #[derive(Debug)] pub struct FirehoseAck { - seq: i64, - sender: Option>, + seq: Seq, + sender: Option>, } impl FirehoseAck { - pub fn new(seq: i64, sender: mpsc::UnboundedSender) -> Self { + pub fn new(seq: Seq, sender: mpsc::UnboundedSender) -> Self { Self { seq, sender: Some(sender), diff --git a/hubble-sync/src/firehose/mod.rs b/hubble-sync/src/firehose/mod.rs index 95eeb33..479e260 100644 --- a/hubble-sync/src/firehose/mod.rs +++ b/hubble-sync/src/firehose/mod.rs @@ -3,3 +3,16 @@ mod subscribe_repos; pub use event::{FirehoseAck, FirehoseEvent, FirehosePayload}; pub use subscribe_repos::{FirehoseConfig, FirehoseError, FirehoseSubscriber}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub struct Seq(pub(crate) u64); + +impl Seq { + /// return a sequence just-before the current one + fn prev(&self) -> Seq { + Seq(self.0.saturating_sub(1)) + } + fn is_non_zero(&self) -> bool { + self.0 > 0 + } +} diff --git a/hubble-sync/src/firehose/subscribe_repos.rs b/hubble-sync/src/firehose/subscribe_repos.rs index c806466..f1fecca 100644 --- a/hubble-sync/src/firehose/subscribe_repos.rs +++ b/hubble-sync/src/firehose/subscribe_repos.rs @@ -1,16 +1,26 @@ //! a `com.atproto.sync.subscribeRepos` websocket -> event dispatcher +use std::collections::BTreeMap; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; -use jacquard_api::com_atproto::sync::subscribe_repos::SubscribeReposMessage; +use jacquard_api::com_atproto::sync::subscribe_repos::{ + SubscribeRepos, SubscribeReposMessage, SubscribeReposStream, +}; +use jacquard_common::stream::StreamErrorKind; +use jacquard_common::websocket::tungstenite_client::TungsteniteClient; +use jacquard_common::xrpc::{SubscriptionExt, SubscriptionStream}; use tokio::sync::mpsc; +use tokio_stream::StreamExt; use tokio_util::sync::CancellationToken; -use tracing::{info, trace}; +use tracing::{info, trace, warn}; -use super::{FirehoseAck, FirehoseEvent, FirehosePayload}; +use super::{FirehoseAck, FirehoseEvent, FirehosePayload, Seq}; use crate::repo_actor::{RepoSendError, RepoTask}; -use crate::{Did, Host, RepoRegistry, StorageEngine, StorageError, SyncConsumer}; +use crate::storage::engine::{StorageBatch, StorageError}; +use crate::storage::firehose_cursor::CursorState; +use crate::storage::{LoadError, StorageEngine}; +use crate::{CancelExt, Did, Host, RepoRegistry, SyncConsumer}; #[derive(Debug, Clone)] pub struct FirehoseConfig { @@ -77,8 +87,14 @@ impl Default for FirehoseConfig { #[derive(Debug, thiserror::Error)] pub enum FirehoseError { + #[error("storage load: {0}")] + StorageLoad(#[source] LoadError), #[error("storage: {0}")] Storage(#[source] E), + #[error("invalid sequence number: {0}")] + InvalidSeq(i64), + #[error("firehose setup error: {0}")] + StreamSetupError(String), #[error("max repo backpressure retry attempts reached for repo: {0:?}")] RepoBackpressureExhausted(Did), #[error("repo actor registry full and couldn't evict any")] @@ -115,15 +131,218 @@ impl> FirehoseSubscriber { } pub async fn run(&mut self, cancel: CancellationToken) -> Result<(), FirehoseError> { - todo!(); + let (ack_tx, mut ack_rx) = mpsc::unbounded_channel::(); + let mut outstanding: BTreeMap = BTreeMap::new(); + let mut last_seq = Seq(0); + let mut cursor_tick = Instant::now(); + let mut reconnect_backoff = Duration::from_secs(1); + + 'reconnect: loop { + if cancel.is_cancelled() { + return Ok(()); + } + let cursor = self.load_cursor().await?; + if let Some(cs) = cursor.as_ref() { + last_seq = last_seq.max(cs.seq); + } + info!(host = %self.host.name(), cursor = ?cursor, "connecting to firehose"); + + let stream = match cancel.run(self.connect(cursor.map(|c| c.seq))).await { + None => return Ok(()), + Some(Ok(s)) => { + reconnect_backoff = Duration::from_secs(1); + s + } + Some(Err(e)) => { + warn!(error = %e, host = %self.host.name(), + backoff_ms = reconnect_backoff.as_millis(), + "firehose connect failed, will retry"); + if !cancel.sleep(reconnect_backoff).await { + return Ok(()); + } + reconnect_backoff = + (reconnect_backoff * 2).min(self.config.reconnect_max_backoff); + continue 'reconnect; + } + }; + + info!(host = %self.host.name(), "firehose connected"); + + match self + .consume( + stream, + &mut outstanding, + &mut last_seq, + &mut cursor_tick, + &mut reconnect_backoff, + &ack_tx, + &mut ack_rx, + &cancel, + ) + .await? + { + ConnExit::Cancelled => return Ok(()), + ConnExit::Reconnect => { + if !cancel.sleep(reconnect_backoff).await { + return Ok(()); + } + reconnect_backoff = + (reconnect_backoff * 2).min(self.config.repo_backpressure_retry_sleep); + continue 'reconnect; + } + ConnExit::ReconnectNow => continue 'reconnect, + } + } + } + + async fn consume( + &self, + stream: SubscriptionStream, + outstanding: &mut BTreeMap, + last_seq: &mut Seq, + cursor_tick: &mut Instant, + backoff: &mut Duration, + ack_tx: &mpsc::UnboundedSender, + ack_rx: &mut mpsc::UnboundedReceiver, + cancel: &CancellationToken, + ) -> Result> { + let (_sink, mut messages) = stream.into_stream(); + let mut last_msg = Instant::now(); + loop { + let can_accept = outstanding.len() < self.config.max_outstanding; + let idle_deadline = self.config.idle_timeout.saturating_sub(last_msg.elapsed()); + tokio::select! { + biased; + // first priority: don't keep working if we're cancelled + slept = cancel.sleep(idle_deadline) => if slept { + warn!(host = %self.host.name(), + timeout_secs = self.config.idle_timeout.as_secs(), + "firehose idle timeout, reconnecting"); + return Ok(ConnExit::ReconnectNow); + } else { + return Ok(ConnExit::Cancelled); + }, + // second: clear outstanding acks before creating new ones + Some(seq) = ack_rx.recv() => { + outstanding.remove(&seq); + }, + // finally (if possible): process new messages + msg = messages.next(), if can_accept => { + let Some(m) = msg else { + info!(host = %self.host.name(), "firehose stream ended, reconnecting"); + return Ok(ConnExit::ReconnectNow); + }; + match m { + Ok(m) => { + // successful message receipt: update normal-operation states + last_msg = Instant::now(); + *backoff = Duration::from_secs(1); + if let Some(seq) = self.dispatch(m, ack_tx).await? { + outstanding.insert(seq, Instant::now()); + *last_seq = seq; + } + }, + Err(e) => match e.kind() { + StreamErrorKind::Decode | StreamErrorKind::WrongMessageFormat => { + // todo: metric + warn!(error = %e, host = %self.host.name(), "firehose decode error, skipping"); + } + StreamErrorKind::Closed => { + info!(host = %self.host.name(), "firehose stream closed, reconnecting"); + return Ok(ConnExit::ReconnectNow); + } + _ => { + warn!(error = %e, host = %self.host.name(), + backoff_ms = backoff.as_millis(), + "firehose stream transport error, reconnecting"); + return Ok(ConnExit::Reconnect); + } + } + } + if cursor_tick.elapsed() >= self.config.cursor_persist_interval { + self.evict_stalled(outstanding); + let watermark = outstanding + .first_entry() + .map(|e| e.key().prev()) + .unwrap_or(*last_seq); + if watermark.is_non_zero() { + self.persist_cursor(watermark).await?; + } + *cursor_tick = Instant::now(); + } + } + } + } + } + + async fn connect( + &self, + cursor: Option, + ) -> Result, FirehoseError> { + let base = self + .host + .jacquard_uri_base("wss") + .map_err(|e| FirehoseError::StreamSetupError(format!("invalid uri: {e}")))?; + let jacquard_cursor = cursor.map(|c| c.0 as i64); + let params = SubscribeRepos { + cursor: jacquard_cursor, + }; + let client = TungsteniteClient::new(); + let stream = client + .subscription(base.to_owned()) + .subscribe(¶ms) + .await + .map_err(|e| FirehoseError::StreamSetupError(format!("subscribe: {e:?}")))?; + Ok(stream) + } + + async fn load_cursor(&self) -> Result, FirehoseError> { + let storage = self.storage.clone(); + let host = self.host.clone(); + tokio::task::spawn_blocking(move || CursorState::load(&*storage, &host)) + .await + .expect("db read not to panic") + .map_err(FirehoseError::StorageLoad) + } + + async fn persist_cursor(&self, seq: Seq) -> Result<(), FirehoseError> { + let storage = self.storage.clone(); + let host = self.host.clone(); + tokio::task::spawn_blocking(move || { + let cs = CursorState { + seq, + persisted_at: std::time::SystemTime::now(), + }; + let mut batch = storage.batch(); + cs.store(&host, &mut batch); + batch.commit() + }) + .await + .expect("db cursor persist not to panic") + .map_err(FirehoseError::Storage) + } + + fn evict_stalled(&self, outstanding: &mut BTreeMap) { + let threshold = Instant::now() - self.config.cursor_unack_hold_limit; + let mut num_evicted = 0; + outstanding.retain(|&seq, when| { + if *when < threshold { + warn!(seq = ?seq, host = %self.host.name(), "evicting stalled seq"); + num_evicted += 1; + false + } else { + true + } + }); + let _ = num_evicted; // TODO this is a metric } /// from a firehose event -> the right repo actor pub(super) async fn dispatch( &self, msg: SubscribeReposMessage, - ack_tx: &mpsc::UnboundedSender, - ) -> Result, FirehoseError> { + ack_tx: &mpsc::UnboundedSender, + ) -> Result, FirehoseError> { let Some((did, seq, payload)) = convert_message(msg) else { return Ok(None); }; @@ -163,24 +382,41 @@ impl> FirehoseSubscriber { } } -fn convert_message(m: SubscribeReposMessage) -> Option<(Did, i64, FirehosePayload)> { +enum ConnExit { + Cancelled, + Reconnect, + ReconnectNow, +} + +fn convert_message(m: SubscribeReposMessage) -> Option<(Did, Seq, FirehosePayload)> { use SubscribeReposMessage as M; + let from_i = |s: i64| { + let u = s + .try_into() + .inspect_err(|_| info!(seq = %s, "event with dropping invalid seq")) + .ok()?; + Some(Seq(u)) + }; match m { M::Commit(c) => { - trace!(did = %c.repo, rev = %c.rev, seq = c.seq, "firehose #commit"); - Some(((&c.repo).into(), c.seq, FirehosePayload::Commit(c))) + let seq = from_i(c.seq)?; + trace!(did = %c.repo, rev = %c.rev, seq = ?seq, "firehose #commit"); + Some(((&c.repo).into(), seq, FirehosePayload::Commit(c))) } M::Sync(s) => { - trace!(did = %s.did, rev = %s.rev, seq = s.seq, "firehose #sync"); - Some(((&s.did).into(), s.seq, FirehosePayload::Sync(s))) + let seq = from_i(s.seq)?; + trace!(did = %s.did, rev = %s.rev, seq = ?seq, "firehose #sync"); + Some(((&s.did).into(), seq, FirehosePayload::Sync(s))) } M::Account(a) => { - trace!(did = %a.did, active = a.active, seq = a.seq, "firehose #account"); - Some(((&a.did).into(), a.seq, FirehosePayload::Account(a))) + let seq = from_i(a.seq)?; + trace!(did = %a.did, active = a.active, seq = ?seq, "firehose #account"); + Some(((&a.did).into(), seq, FirehosePayload::Account(a))) } M::Identity(i) => { - trace!(did = %i.did, seq = i.seq, "firehose #identity"); - Some(((&i.did).into(), i.seq, FirehosePayload::Identity(i))) + let seq = from_i(i.seq)?; + trace!(did = %i.did, seq = ?seq, "firehose #identity"); + Some(((&i.did).into(), seq, FirehosePayload::Identity(i))) } M::Info(info) => { info!(name = %info.name, message = ?info.message, "firehose #info"); diff --git a/hubble-sync/src/host.rs b/hubble-sync/src/host.rs index a4feb07..d64180b 100644 --- a/hubble-sync/src/host.rs +++ b/hubble-sync/src/host.rs @@ -153,6 +153,17 @@ impl Host { &self.name } + pub fn jacquard_uri_base( + &self, + scheme: &str, + ) -> Result, String> { + // ...could maybe use the builder stuff, but it's pretty verbose + let base = format!("{scheme}://{}", self.name().as_str()); + let uri = jacquard_common::deps::fluent_uri::Uri::parse(base) + .map_err(|e| format!("failed to parse uri: {e:?}"))?; + Ok(uri) + } + pub fn sync_compliance( &self, engine: &S, diff --git a/hubble-sync/src/lib.rs b/hubble-sync/src/lib.rs index c7770d4..4461481 100644 --- a/hubble-sync/src/lib.rs +++ b/hubble-sync/src/lib.rs @@ -1,3 +1,4 @@ +mod cancel; mod config; mod did; mod firehose; @@ -8,6 +9,7 @@ mod resync_scheduler; mod storage; mod sync_consumer; +pub use cancel::CancelExt; pub use config::SyncConfig; pub use did::Did; pub use firehose::{ diff --git a/hubble-sync/src/storage/firehose_cursor.rs b/hubble-sync/src/storage/firehose_cursor.rs index a61445c..fc6bfa8 100644 --- a/hubble-sync/src/storage/firehose_cursor.rs +++ b/hubble-sync/src/storage/firehose_cursor.rs @@ -7,13 +7,23 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use crate::firehose::Seq; use crate::{Host, StorageBatch, StorageEngine, StorageError}; use super::{DecodeError, LoadError, PREFIX_FIREHOSE_CURSOR}; +impl Seq { + pub fn from_bytes(bytes: &[u8; 8]) -> Self { + Self(u64::from_be_bytes(*bytes)) + } + pub fn to_bytes(self) -> [u8; 8] { + self.0.to_be_bytes() + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct CursorState { - pub seq: u64, + pub seq: Seq, pub persisted_at: SystemTime, } @@ -33,7 +43,7 @@ impl CursorState { .expect("to be persisting cursors after the unix epoch") .as_millis() as u64; let mut out = Vec::with_capacity(8 + 8); - out.extend_from_slice(&self.seq.to_be_bytes()); + out.extend_from_slice(&self.seq.to_bytes()); out.extend_from_slice(&ms.to_be_bytes()); out } @@ -50,7 +60,7 @@ impl CursorState { } let persiste_at_ms = u64::from_be_bytes(*persisted_at_bytes); Ok(Self { - seq: u64::from_be_bytes(*seq_bytes), + seq: Seq::from_bytes(seq_bytes), persisted_at: UNIX_EPOCH + Duration::from_millis(persiste_at_ms), }) } @@ -86,7 +96,7 @@ mod tests { } fn cs(seq: u64, persisted_at_ms: u64) -> CursorState { CursorState { - seq, + seq: Seq(seq), persisted_at: t(persisted_at_ms), } } @@ -120,8 +130,8 @@ mod tests { cs(100, 1).store(&r1, &mut b); cs(200, 2).store(&r2, &mut b); b.commit().unwrap(); - assert_eq!(CursorState::load(&eng, &r1).unwrap().unwrap().seq, 100); - assert_eq!(CursorState::load(&eng, &r2).unwrap().unwrap().seq, 200); + assert_eq!(CursorState::load(&eng, &r1).unwrap().unwrap().seq, Seq(100)); + assert_eq!(CursorState::load(&eng, &r2).unwrap().unwrap().seq, Seq(200)); } #[test]