From 6a47da5b381c5fce7677bb65578953f845dd8a5d Mon Sep 17 00:00:00 2001 From: dawn <90008@klbr.net> Date: Thu, 8 Oct 2026 03:20:18 +0300 Subject: [PATCH] [ingest] cursor tracker: tickets from the reader, relay-side done, relay-only watermark the saved cursor was whatever seq a shard finished last, so it could run ahead of frames still queued in other shards (finding 1). the reader now takes a ticket per frame in read order, and a SeqTracker per source url moves the watermark only over a done prefix. the trackers live on AppState next to the cursors, so an ingestor that replaces another for the same url (restart, remove and add) queues its frames behind the old one's instead of publishing a second watermark into the same cursor. the reader marks filtered frames done, the relay shard marks the rest done after its commit. relay-only builds store the watermark as the source cursor in place of the last seq. indexer builds compute it unused, their shards still store the cursor until they mark tickets done after their own commit. --- src/ingest/cursor.rs | 182 +++++++++++++++++++++++++++++++ src/ingest/firehose.rs | 60 +++++++++- src/ingest/mod.rs | 2 + src/ingest/relay/sink/indexer.rs | 8 +- src/ingest/relay/sink/none.rs | 5 +- src/ingest/relay/sink/relay.rs | 13 +-- src/ingest/relay/worker.rs | 19 +++- src/state.rs | 15 +++ 8 files changed, 279 insertions(+), 25 deletions(-) create mode 100644 src/ingest/cursor.rs 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, -- 2.54.0