diff --git a/src/ingest/cursor.rs b/src/ingest/cursor.rs new file mode 100644 index 0000000..d0dc87e --- /dev/null +++ b/src/ingest/cursor.rs @@ -0,0 +1,182 @@ +//! the firehose cursor as a low watermark. the reader takes a ticket for every frame in the +//! order it reads them, and the cursor only moves past a seq once every frame read before it is +//! done too, so a crash replays whatever was still queued or in flight anywhere + +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use parking_lot::Mutex; +use url::Url; + +use crate::state::AppState; + +pub(crate) struct SeqTracker { + window: Mutex, + /// whether the watermark goes into the source's cursor. indexer builds don't use it yet, + /// their shards store the cursor themselves + publishes: bool, +} + +#[derive(Default)] +struct Window { + /// the ticket id of `pending[0]` + front: u64, + /// seq and whether it's done, per frame not yet under the watermark, in read order + pending: VecDeque<(i64, bool)>, +} + +/// one frame the reader handed on. it has to reach [`Ticket::done`] once whatever it caused is +/// committed, or the cursor never moves past it again. dropping one leaves the cursor before it +/// on purpose, so messages a shutdown drops are read again next start +#[must_use = "a ticket that's never done pins the firehose cursor"] +pub(crate) struct Ticket { + tracker: Arc, + id: u64, + seq: i64, +} + +impl SeqTracker { + pub(crate) fn new(publishes: bool) -> Arc { + Arc::new(Self { + window: Mutex::default(), + publishes, + }) + } + + pub(crate) fn ticket(self: &Arc, seq: i64) -> Ticket { + let mut window = self.window.lock(); + let id = window.front + window.pending.len() as u64; + window.pending.push_back((seq, false)); + Ticket { + tracker: self.clone(), + id, + seq, + } + } + + /// frames read and not done yet, plus done ones still behind an earlier one + #[cfg(test)] + pub(crate) fn pending(&self) -> usize { + self.window.lock().pending.len() + } +} + +impl Ticket { + #[cfg(test)] + pub(crate) fn seq(&self) -> i64 { + self.seq + } + + /// marks the frame done. when that moves the watermark, `publish` gets the new one, under + /// the tracker's lock so two shards finishing at once can't store it out of order + pub(crate) fn done(self, publish: impl FnOnce(i64)) { + let mut window = self.tracker.window.lock(); + let index = (self.id - window.front) as usize; + window.pending[index].1 = true; + let mut watermark = None; + while let Some(&(seq, true)) = window.pending.front() { + window.pending.pop_front(); + window.front += 1; + watermark = Some(seq); + } + if let Some(seq) = watermark + && self.tracker.publishes + { + publish(seq); + } + } + + /// a ticket for a frame no reader took, for tests that hand messages to a shard directly + #[cfg(all(test, feature = "indexer"))] + pub(crate) fn detached(seq: i64) -> Self { + SeqTracker::new(false).ticket(seq) + } +} + +impl std::fmt::Debug for Ticket { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Ticket") + .field("id", &self.id) + .field("seq", &self.seq) + .finish() + } +} + +/// marks `ticket` done, storing a moved watermark as `firehose`'s cursor +pub(crate) fn done(state: &AppState, firehose: &Url, ticket: Ticket) { + ticket.done(|seq| { + state + .firehose_cursors + .peek_with(firehose, |_, cursor| cursor.store(seq, Ordering::SeqCst)); + }); +} + +#[cfg(test)] +mod tests { + use super::*; + use rand::seq::SliceRandom; + + fn published(tracker: &Arc, tickets: Vec) -> Vec> { + let _ = tracker; + tickets + .into_iter() + .map(|ticket| { + let mut moved = None; + ticket.done(|seq| moved = Some(seq)); + moved + }) + .collect() + } + + #[test] + fn the_watermark_waits_for_the_oldest_frame() { + let tracker = SeqTracker::new(true); + let mut tickets = [10, 11, 12].map(|seq| Some(tracker.ticket(seq))); + let take = |tickets: &mut [Option; 3], i: usize| tickets[i].take().unwrap(); + + let last = take(&mut tickets, 2); + let middle = take(&mut tickets, 1); + let first = take(&mut tickets, 0); + assert_eq!( + published(&tracker, vec![last, middle, first]), + [None, None, Some(12)] + ); + assert_eq!(tracker.pending(), 0); + } + + #[test] + fn a_tracker_that_doesnt_publish_still_drains() { + let tracker = SeqTracker::new(false); + let tickets = (0..3).map(|seq| tracker.ticket(seq)).collect(); + assert_eq!(published(&tracker, tickets), [None, None, None]); + assert_eq!(tracker.pending(), 0); + } + + /// whatever order shards finish in, the cursor never passes a frame that isn't done + #[test] + fn the_cursor_never_passes_an_unfinished_frame() { + let mut rng = rand::rng(); + for _ in 0..200 { + let tracker = SeqTracker::new(true); + let mut tickets = (0..64) + .map(|seq| tracker.ticket(seq * 2)) + .collect::>(); + tickets.shuffle(&mut rng); + let mut done = std::collections::BTreeSet::new(); + let mut cursor = -1; + for ticket in tickets { + done.insert(ticket.seq()); + ticket.done(|seq| { + assert!(seq > cursor, "the watermark went backwards"); + cursor = seq; + }); + let unfinished = (0..64).map(|seq| seq * 2).find(|seq| !done.contains(seq)); + if let Some(unfinished) = unfinished { + assert!(cursor < unfinished, "cursor {cursor} passed {unfinished}"); + } + } + assert_eq!(cursor, 126); + } + } +} diff --git a/src/ingest/firehose.rs b/src/ingest/firehose.rs index 89d7035..c05bc91 100644 --- a/src/ingest/firehose.rs +++ b/src/ingest/firehose.rs @@ -1,4 +1,5 @@ use crate::filter::{FilterHandle, FilterMode}; +use crate::ingest::cursor::{self, SeqTracker}; use crate::ingest::stream::{ FirehoseError, FirehoseStream, InfoName, SubscribeReposMessage, decode_frame, }; @@ -164,6 +165,9 @@ pub struct FirehoseIngestor { throttle: ThrottleHandle, max_failures: usize, stats: Arc, + /// shared with every ingestor this source ever gets, so frames a reconnect or a respawn + /// reads again queue up behind the ones still in flight from before + tracker: Arc, } impl FirehoseIngestor { @@ -180,6 +184,7 @@ impl FirehoseIngestor { ) -> Self { let throttle = state.throttler.get_handle(&relay_host).await; let stats = state.firehose_stats.handle(&relay_host).await; + let tracker = state.firehose_tracker(&relay_host).await; Self { state, buffer_tx, @@ -192,6 +197,7 @@ impl FirehoseIngestor { throttle, max_failures, stats, + tracker, } } @@ -589,13 +595,14 @@ impl FirehoseIngestor { } async fn handle_message(&self, msg: SubscribeReposMessage<'_>) { - let did = match &msg { - SubscribeReposMessage::Commit(commit) => &commit.repo, - SubscribeReposMessage::Identity(identity) => &identity.did, - SubscribeReposMessage::Account(account) => &account.did, - SubscribeReposMessage::Sync(sync) => &sync.did, + let (did, seq) = match &msg { + SubscribeReposMessage::Commit(commit) => (&commit.repo, commit.seq), + SubscribeReposMessage::Identity(identity) => (&identity.did, identity.seq), + SubscribeReposMessage::Account(account) => (&account.did, account.seq), + SubscribeReposMessage::Sync(sync) => (&sync.did, sync.seq), _ => return, }; + let ticket = self.tracker.ticket(seq); let should_process_started = crate::ingest::firehose_stats::StatsInstant::now(); let process = self @@ -608,6 +615,7 @@ impl FirehoseIngestor { if !process { self.stats.record_skipped(); trace!(did = %did, "skipping: not in filter"); + cursor::done(&self.state, &self.relay_host, ticket); return; } trace!(did = %did, "forwarding message to ingest buffer"); @@ -619,6 +627,7 @@ impl FirehoseIngestor { url: self.relay_host.clone(), is_pds: self.is_pds, msg: msg.into_static(), + ticket, }) .await; { @@ -785,6 +794,47 @@ mod tests { Ok(()) } + /// a source whose ingestor is replaced (restart, remove and add again) while frames the old + /// one read are still in the shards must not get a second watermark for the same cursor + #[tokio::test] + async fn a_replaced_ingestor_waits_for_the_frames_of_the_old_one() -> Result<()> { + let tmp = tempfile::tempdir().into_diagnostic()?; + let cfg = Config { + database_path: tmp.path().to_path_buf(), + ..Default::default() + }; + let state = Arc::new(AppState::new(&cfg)?); + let relay = Url::parse("wss://relay.example/").into_diagnostic()?; + let ingestor = || { + FirehoseIngestor::new( + state.clone(), + BufferTx::channel(1).0, + relay.clone(), + false, + false, + state.filter.clone(), + state.firehose_enabled.subscribe(), + false, + 1, + ) + }; + + let first = ingestor().await; + let in_flight = first.tracker.ticket(1); + drop(first); + let second = ingestor().await; + let mut moved = None; + second.tracker.ticket(2).done(|seq| moved = Some(seq)); + assert_eq!( + moved, None, + "the cursor passed a frame the first ingestor read" + ); + assert_eq!(second.tracker.pending(), 2); + in_flight.done(|_| {}); + assert_eq!(second.tracker.pending(), 0); + Ok(()) + } + #[tokio::test] async fn clear_stale_cursor_resets_db_and_memory() -> Result<()> { let tmp = tempfile::tempdir().into_diagnostic()?; diff --git a/src/ingest/mod.rs b/src/ingest/mod.rs index 76f4933..4c851c9 100644 --- a/src/ingest/mod.rs +++ b/src/ingest/mod.rs @@ -1,3 +1,4 @@ +pub(crate) mod cursor; pub mod firehose; pub mod firehose_stats; #[cfg(feature = "indexer")] @@ -19,6 +20,7 @@ pub enum IngestMessage { /// enables host authority enforcement in the worker. is_pds: bool, msg: SubscribeReposMessage<'static>, + ticket: cursor::Ticket, }, } diff --git a/src/ingest/relay/sink/indexer.rs b/src/ingest/relay/sink/indexer.rs index 5abe8ad..39a211b 100644 --- a/src/ingest/relay/sink/indexer.rs +++ b/src/ingest/relay/sink/indexer.rs @@ -11,6 +11,10 @@ use crate::ingest::stream::{Account, Commit, Identity, Sync}; use crate::state::AppState; use crate::types::{RepoState, RepoStatus}; +/// the indexer shards still store the cursor themselves, so the relay side's watermark is only +/// computed until they mark frames done after their own commit +pub(crate) const PUBLISHES_WATERMARK: bool = false; + /// per-worker seed from which each shard builds its sink. #[derive(Clone)] pub(crate) struct SinkSeed { @@ -161,10 +165,6 @@ impl EventSink { let _ = self.hook.blocking_send(msg); } } - - pub(crate) fn advance_cursor(&self, _state: &AppState, _firehose: &Url, _seq: i64) { - // in indexer mode the FirehoseWorker advances the cursor after processing - } } #[cfg(test)] diff --git a/src/ingest/relay/sink/none.rs b/src/ingest/relay/sink/none.rs index fc8e590..63c2f2e 100644 --- a/src/ingest/relay/sink/none.rs +++ b/src/ingest/relay/sink/none.rs @@ -9,6 +9,9 @@ use crate::ingest::stream::{Account, Commit, Identity, Sync}; use crate::state::AppState; use crate::types::RepoState; +/// nothing consumes the events, so nothing ever moved the cursor of a bare build +pub(crate) const PUBLISHES_WATERMARK: bool = false; + /// per-worker seed from which each shard builds its sink. #[derive(Clone, Default)] pub(crate) struct SinkSeed; @@ -91,6 +94,4 @@ impl EventSink { } pub(crate) fn flush(&mut self, _state: &AppState, _staged: Staged) {} - - pub(crate) fn advance_cursor(&self, _state: &AppState, _firehose: &Url, _seq: i64) {} } diff --git a/src/ingest/relay/sink/relay.rs b/src/ingest/relay/sink/relay.rs index dc31d9c..33f15ee 100644 --- a/src/ingest/relay/sink/relay.rs +++ b/src/ingest/relay/sink/relay.rs @@ -2,7 +2,6 @@ use fjall::OwnedWriteBatch; use miette::Result; #[cfg(feature = "jetstream")] use smol_str::ToSmolStr; -use std::sync::atomic::Ordering; use url::Url; use crate::db::outbox::{Outbox, RelayFrameRef, RelayOutbox}; @@ -16,6 +15,10 @@ use crate::types::StoredJetstreamEvent; #[cfg(feature = "jetstream")] use jacquard_common::CowStr; +/// the relay is the last stage, so a frame is done once its relay commit lands, and the +/// watermark is the source's cursor +pub(crate) const PUBLISHES_WATERMARK: bool = true; + /// per-worker seed from which each shard builds its sink. #[derive(Clone, Default)] pub(crate) struct SinkSeed; @@ -180,14 +183,6 @@ impl EventSink { /// frames are broadcast as their transaction commits. pub(crate) fn flush(&mut self, _state: &AppState, _staged: Staged) {} - - /// the relay is the terminal consumer, so it advances the firehose cursor - /// itself once a message is fully processed. - pub(crate) fn advance_cursor(&self, state: &AppState, firehose: &Url, seq: i64) { - state - .firehose_cursors - .peek_with(firehose, |_, c| c.store(seq, Ordering::SeqCst)); - } } #[cfg(feature = "jetstream")] diff --git a/src/ingest/relay/worker.rs b/src/ingest/relay/worker.rs index 69ded67..e80850d 100644 --- a/src/ingest/relay/worker.rs +++ b/src/ingest/relay/worker.rs @@ -8,6 +8,7 @@ use tracing::{debug, error, info, info_span, trace, warn}; use url::Url; use crate::db::CountDeltas; +use crate::ingest::cursor; use crate::ingest::stream::{InfoName, SubscribeReposMessage}; use crate::ingest::validation::ValidationOptions; use crate::ingest::{BufferRx, BufferTx, IngestMessage}; @@ -130,8 +131,14 @@ impl RelayWorker { break; } let message_started = StatsInstant::now(); - let IngestMessage::Firehose { url, is_pds, msg } = msg; + let IngestMessage::Firehose { + url, + is_pds, + msg, + ticket, + } = msg; if let SubscribeReposMessage::Info(inf) = msg { + cursor::done(&state, &url, ticket); shard_stats.record_info(); match inf.name { InfoName::OutdatedCursor => {} @@ -157,7 +164,7 @@ impl RelayWorker { SubscribeReposMessage::Identity(i) => (i.did.clone(), i.seq), SubscribeReposMessage::Account(a) => (a.did.clone(), a.seq), SubscribeReposMessage::Sync(s) => (s.did.clone(), s.seq), - _ => continue, + SubscribeReposMessage::Info(_) => unreachable!("info frames are handled above"), }; shard_stats.record_received(seq); @@ -183,6 +190,8 @@ impl RelayWorker { Err(e) => { shard_stats.record_commit_error(); error!(shard = id, err = %e, "relay shard: failed to commit batch"); + // a later frame's commit would have moved the cursor past this one anyway + cursor::done(&state, &firehose, ticket); continue; } }; @@ -191,10 +200,8 @@ impl RelayWorker { ctx.sink.flush(&state, staged); let broadcast = broadcast_started.elapsed(); - // advance cursor for this firehose only if we are the terminal consumer (relay mode) - // in events mode, FirehoseWorker will advance the cursor after processing let cursor_started = StatsInstant::now(); - ctx.sink.advance_cursor(&state, &firehose, seq); + cursor::done(&state, &firehose, ticket); shard_stats.record_processed(crate::ingest::firehose_stats::RelayShardTimings { process_message, stage_counts: commit_timings.stage_counts, @@ -485,6 +492,7 @@ mod tests { .send(IngestMessage::Firehose { url: Url::parse("wss://relay.example").into_diagnostic()?, is_pds: false, + ticket: cursor::Ticket::detached(1), msg: SubscribeReposMessage::Identity(Box::new(Identity { did: self.did.clone(), handle: None, @@ -533,6 +541,7 @@ mod tests { .send(IngestMessage::Firehose { url: Url::parse("wss://relay.example").into_diagnostic()?, is_pds: false, + ticket: cursor::Ticket::detached(secs), msg: SubscribeReposMessage::Account(Box::new(Account { active, did: self.did.clone(), diff --git a/src/state.rs b/src/state.rs index c292e39..0c9c4b8 100644 --- a/src/state.rs +++ b/src/state.rs @@ -12,6 +12,7 @@ use tokio::sync::Notify; use tokio::sync::{Mutex, Semaphore, watch}; use url::Url; +use crate::ingest::cursor::SeqTracker; use crate::ingest::firehose_stats::FirehoseStats; use crate::pds_daily_limit::PdsDailyLimit; use crate::{ @@ -57,6 +58,10 @@ pub struct AppState { pub(crate) pds_daily_limit: PdsDailyLimit, pub(crate) tier_policy: TierPolicy, pub firehose_cursors: scc::HashIndex, + /// one per source url for the life of the process, so an ingestor that replaces another + /// orders its frames behind the ones the old one left in flight instead of publishing a + /// second watermark into the same cursor + firehose_trackers: scc::HashMap>, // nothing submits jobs yet #[allow(dead_code)] pub(crate) lookups: crate::ingest::relay::lookup::Lookups, @@ -193,6 +198,7 @@ impl AppState { pds_status_write: std::sync::Arc::new(parking_lot::Mutex::new(())), tier_policy: config.tier_policy.clone(), firehose_cursors: relay_cursors, + firehose_trackers: scc::HashMap::new(), lookups: Default::default(), firehose_stats: FirehoseStats::default(), #[cfg(feature = "indexer")] @@ -253,6 +259,15 @@ impl AppState { }) } + pub(crate) async fn firehose_tracker(&self, url: &Url) -> std::sync::Arc { + self.firehose_trackers + .entry_async(url.clone()) + .await + .or_insert_with(|| SeqTracker::new(crate::ingest::relay::sink::PUBLISHES_WATERMARK)) + .get() + .clone() + } + pub(crate) async fn identity_reconcile_is_cooled_down( &self, did: &Did,