//! 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<Window>,
/// whether the watermark goes anywhere. a bare build has nothing downstream to move it
publishes: bool,
}
#[derive(Default)]
struct Window {
/// the ticket id of `pending[0]`
front: u64,
/// per frame not yet under the watermark, in read order
pending: VecDeque<Frame>,
}
#[derive(Clone, Copy)]
struct Frame {
seq: i64,
/// the event time the frame carries, unix ms
time_ms: i64,
done: 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<SeqTracker>,
id: u64,
seq: i64,
}
/// the newest frame every frame read before it is done with too
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct Watermark {
pub(crate) seq: i64,
pub(crate) time_ms: i64,
}
impl SeqTracker {
pub(crate) fn new(publishes: bool) -> Arc<Self> {
Arc::new(Self {
window: Mutex::default(),
publishes,
})
}
pub(crate) fn ticket(self: &Arc<Self>, seq: i64, time_ms: i64) -> Ticket {
let mut window = self.window.lock();
let id = window.front + window.pending.len() as u64;
window.pending.push_back(Frame {
seq,
time_ms,
done: false,
});
Ticket {
tracker: self.clone(),
id,
seq,
}
}
/// frames read and not done yet, plus done ones still behind an earlier one
#[cfg(any(test, feature = "indexer_stream"))]
pub(crate) fn pending(&self) -> usize {
self.window.lock().pending.len()
}
}
impl Ticket {
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(Watermark)) {
let mut window = self.tracker.window.lock();
let index = (self.id - window.front) as usize;
window.pending[index].done = true;
let mut watermark = None;
while let Some(&frame) = window.pending.front()
&& frame.done
{
window.pending.pop_front();
window.front += 1;
watermark = Some(Watermark {
seq: frame.seq,
time_ms: frame.time_ms,
});
}
if let Some(watermark) = watermark
&& self.tracker.publishes
{
publish(watermark);
}
}
/// 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, 0)
}
}
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 and as how far it got
/// in event time
pub(crate) fn done(state: &AppState, firehose: &Url, ticket: Ticket) {
ticket.done(|watermark| {
state.firehose_cursors.peek_with(firehose, |_, cursor| {
cursor.store(watermark.seq, Ordering::SeqCst)
});
state.note_dispatched(firehose, watermark.time_ms);
});
}
#[cfg(test)]
mod tests {
use super::*;
use rand::seq::SliceRandom;
fn published(tracker: &Arc<SeqTracker>, tickets: Vec<Ticket>) -> Vec<Option<i64>> {
let _ = tracker;
tickets
.into_iter()
.map(|ticket| {
let mut moved = None;
ticket.done(|watermark| moved = Some(watermark.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, 0)));
let take = |tickets: &mut [Option<Ticket>; 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, 0)).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, 0))
.collect::<Vec<_>>();
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(|watermark| {
assert!(watermark.seq > cursor, "the watermark went backwards");
cursor = watermark.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);
}
}
}