use std::{
collections::BTreeMap,
mem::{replace, size_of},
sync::{
Arc,
atomic::{AtomicBool, AtomicU64, Ordering},
},
};
use reifydb_core::{
common::CommitVersion,
interface::cdc::{Cdc, CdcChange},
metrics::{collect::MetricsCollector, sample::MetricsSample},
};
use reifydb_runtime::sync::rwlock::RwLock;
use reifydb_value::{byte_size::ByteSize, reifydb_assertions};
pub enum BacklogPull {
Hit {
items: Vec<Arc<Cdc>>,
advance_to: CommitVersion,
more: bool,
},
Behind,
}
struct BacklogInner {
entries: BTreeMap<CommitVersion, (u64, Arc<Cdc>)>,
bytes: u64,
cover_from: Option<CommitVersion>,
}
type Waker = Box<dyn Fn() + Send + Sync>;
struct BacklogShared {
inner: RwLock<BacklogInner>,
limit: u64,
anchor: AtomicU64,
waker: RwLock<Option<Waker>>,
armed: AtomicBool,
published_entries: AtomicU64,
pull_hits: AtomicU64,
pull_behinds: AtomicU64,
evicted_floor: AtomicU64,
evicted_ceiling: AtomicU64,
}
#[derive(Clone)]
pub struct FlowBacklog {
shared: Arc<BacklogShared>,
}
impl FlowBacklog {
pub fn new(limit: ByteSize) -> Self {
Self {
shared: Arc::new(BacklogShared {
inner: RwLock::new(BacklogInner {
entries: BTreeMap::new(),
bytes: 0,
cover_from: None,
}),
limit: limit.as_bytes().max(1),
anchor: AtomicU64::new(0),
waker: RwLock::new(None),
armed: AtomicBool::new(false),
published_entries: AtomicU64::new(0),
pull_hits: AtomicU64::new(0),
pull_behinds: AtomicU64::new(0),
evicted_floor: AtomicU64::new(0),
evicted_ceiling: AtomicU64::new(0),
}),
}
}
pub fn limit(&self) -> ByteSize {
ByteSize::from_bytes(self.shared.limit)
}
pub fn set_waker(&self, waker: impl Fn() + Send + Sync + 'static) {
*self.shared.waker.write() = Some(Box::new(waker));
}
pub fn notify(&self) {
if self.shared.armed.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire).is_ok()
&& let Some(waker) = self.shared.waker.read().as_ref()
{
waker();
}
}
pub fn disarm(&self) {
self.shared.armed.store(false, Ordering::Release);
}
pub fn publish(&self, version: CommitVersion, cdc: Option<Arc<Cdc>>) {
let mut inner = self.shared.inner.write();
if inner.cover_from.is_none() {
inner.cover_from = Some(CommitVersion(version.0.saturating_sub(1)));
}
let Some(cdc) = cdc else {
return;
};
if version <= inner.cover_from.expect("cover_from set above") {
return;
}
let bytes = cdc_bytes(&cdc);
if let Some((replaced, _)) = inner.entries.insert(version, (bytes, cdc)) {
inner.bytes -= replaced;
}
inner.bytes += bytes;
self.shared.published_entries.fetch_add(1, Ordering::Relaxed);
self.evict_over_ceiling(&mut inner);
}
fn evict_over_ceiling(&self, inner: &mut BacklogInner) {
let anchor = CommitVersion(self.shared.anchor.load(Ordering::Acquire));
while inner.bytes > self.shared.limit {
let Some(lowest) = inner.entries.keys().next().copied() else {
break;
};
if lowest > anchor {
break;
}
if let Some((evicted, _)) = inner.entries.remove(&lowest) {
inner.bytes -= evicted;
}
inner.cover_from = Some(inner.cover_from.map_or(lowest, |c| c.max(lowest)));
self.shared.evicted_ceiling.fetch_add(1, Ordering::Relaxed);
}
}
pub fn pull(&self, cursor: CommitVersion, up_to: CommitVersion, budget: ByteSize) -> BacklogPull {
if up_to <= cursor {
return BacklogPull::Hit {
items: Vec::new(),
advance_to: cursor,
more: false,
};
}
let inner = self.shared.inner.read();
let Some(cover_from) = inner.cover_from else {
self.shared.pull_behinds.fetch_add(1, Ordering::Relaxed);
return BacklogPull::Behind;
};
if cursor < cover_from {
self.shared.pull_behinds.fetch_add(1, Ordering::Relaxed);
return BacklogPull::Behind;
}
let budget = budget.as_bytes().max(1);
let mut items: Vec<Arc<Cdc>> = Vec::new();
let mut taken = 0u64;
let mut truncated_at: Option<CommitVersion> = None;
for (version, (bytes, cdc)) in inner.entries.range(next_version(cursor)..=up_to) {
if !items.is_empty() && taken + bytes > budget {
truncated_at = Some(*version);
break;
}
taken += bytes;
items.push(cdc.clone());
}
self.shared.pull_hits.fetch_add(1, Ordering::Relaxed);
match truncated_at {
Some(_) => BacklogPull::Hit {
advance_to: items.last().expect("truncation implies at least one item").version,
items,
more: true,
},
None => BacklogPull::Hit {
items,
advance_to: up_to,
more: false,
},
}
}
pub fn evict_below(&self, version: CommitVersion) {
let mut inner = self.shared.inner.write();
if inner.cover_from.is_none() {
return;
}
let retained = inner.entries.split_off(&next_version(version));
let evicted = replace(&mut inner.entries, retained);
let count = evicted.len() as u64;
for (bytes, _) in evicted.into_values() {
inner.bytes -= bytes;
}
inner.cover_from = Some(inner.cover_from.map_or(version, |c| c.max(version)));
self.shared.evicted_floor.fetch_add(count, Ordering::Relaxed);
}
pub fn set_anchor(&self, version: CommitVersion) {
reifydb_assertions! {
let prev = self.shared.anchor.load(Ordering::Acquire);
assert!(
version.0 >= prev,
"the backlog scan anchor moved backwards ({} -> {}), so ceiling eviction could remove \
entries the supervisor has not scanned for DDL yet",
prev,
version.0
);
}
self.shared.anchor.store(version.0, Ordering::Release);
}
}
#[inline]
fn next_version(v: CommitVersion) -> CommitVersion {
CommitVersion(v.0.saturating_add(1))
}
impl MetricsCollector for FlowBacklog {
fn collect(&self, out: &mut Vec<MetricsSample>) {
let (bytes, count, cover_from) = {
let inner = self.shared.inner.read();
(inner.bytes, inner.entries.len() as u64, inner.cover_from.map(|c| c.0).unwrap_or(0))
};
out.push(MetricsSample::heap("flow_backlog", "bytes", ByteSize::from_bytes(bytes)));
out.push(MetricsSample::count("flow_backlog", "entries", count));
out.push(MetricsSample::count("flow_backlog", "cover_from", cover_from));
out.push(MetricsSample::counter(
"flow_backlog",
"published_entries",
self.shared.published_entries.load(Ordering::Relaxed),
));
out.push(MetricsSample::counter(
"flow_backlog",
"pull_hits",
self.shared.pull_hits.load(Ordering::Relaxed),
));
out.push(MetricsSample::counter(
"flow_backlog",
"pull_behinds",
self.shared.pull_behinds.load(Ordering::Relaxed),
));
out.push(MetricsSample::counter(
"flow_backlog",
"evicted_floor",
self.shared.evicted_floor.load(Ordering::Relaxed),
));
out.push(MetricsSample::counter(
"flow_backlog",
"evicted_ceiling",
self.shared.evicted_ceiling.load(Ordering::Relaxed),
));
}
}
pub fn cdc_bytes(cdc: &Cdc) -> u64 {
let system: usize = cdc
.changes
.iter()
.map(|change| size_of::<CdcChange>() + change.key().len() + change.value_bytes())
.sum();
(size_of::<Cdc>() + system) as u64
}
#[cfg(test)]
mod tests {
use std::sync::atomic::AtomicUsize;
use reifydb_codec::{key::encoded::EncodedKey, row::bytes::EncodedBytes};
use reifydb_value::{util::cowvec::CowVec, value::datetime::DateTime};
use super::*;
fn cv(n: u64) -> CommitVersion {
CommitVersion(n)
}
fn cdc_with_payload(version: u64, payload: usize) -> Arc<Cdc> {
Arc::new(Cdc::new(
cv(version),
DateTime::default(),
vec![CdcChange::Insert {
key: EncodedKey::new(vec![0xAB; 4]),
post: EncodedBytes(CowVec::new(vec![0u8; payload])),
}],
))
}
fn backlog(limit_bytes: u64) -> FlowBacklog {
let b = FlowBacklog::new(ByteSize::from_bytes(limit_bytes));
b.set_anchor(cv(u64::MAX));
b
}
fn entry_bytes() -> u64 {
cdc_bytes(&cdc_with_payload(1, 100))
}
#[test]
fn pull_before_any_publish_is_behind() {
let b = backlog(u64::MAX);
assert!(matches!(b.pull(cv(0), cv(10), ByteSize::from_mib(1)), BacklogPull::Behind));
}
#[test]
fn coverage_starts_just_below_the_first_published_version() {
let b = backlog(u64::MAX);
b.publish(cv(100), Some(cdc_with_payload(100, 10)));
match b.pull(cv(99), cv(100), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
advance_to,
more,
} => {
assert_eq!(items.len(), 1);
assert_eq!(advance_to, cv(100));
assert!(!more);
}
BacklogPull::Behind => panic!("cursor at cover_from must be served"),
}
assert!(matches!(b.pull(cv(98), cv(100), ByteSize::from_mib(1)), BacklogPull::Behind));
}
#[test]
fn irrelevant_versions_extend_coverage_without_entries() {
let b = backlog(u64::MAX);
b.publish(cv(5), None);
match b.pull(cv(4), cv(9), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
advance_to,
more,
} => {
assert!(items.is_empty());
assert_eq!(advance_to, cv(9), "an empty pull must advance to the caller's bound");
assert!(!more);
}
BacklogPull::Behind => panic!("published coverage must serve the empty range"),
}
}
#[test]
fn pull_up_to_at_or_below_cursor_is_an_empty_hit() {
let b = backlog(u64::MAX);
b.publish(cv(5), Some(cdc_with_payload(5, 10)));
match b.pull(cv(5), cv(5), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
advance_to,
..
} => {
assert!(items.is_empty());
assert_eq!(advance_to, cv(5));
}
BacklogPull::Behind => panic!("nothing to pull is not Behind"),
}
}
#[test]
fn budget_truncation_reports_more_and_advances_only_to_the_last_taken() {
let b = backlog(u64::MAX);
for v in 1..=4 {
b.publish(cv(v), Some(cdc_with_payload(v, 100)));
}
let two = entry_bytes() * 2;
match b.pull(cv(0), cv(4), ByteSize::from_bytes(two)) {
BacklogPull::Hit {
items,
advance_to,
more,
} => {
assert_eq!(items.len(), 2);
assert_eq!(advance_to, cv(2));
assert!(more, "a truncated pull must tell the caller to come back");
}
BacklogPull::Behind => panic!("expected Hit"),
}
}
#[test]
fn a_single_oversized_entry_is_still_served() {
let b = backlog(u64::MAX);
b.publish(cv(1), Some(cdc_with_payload(1, 4096)));
match b.pull(cv(0), cv(1), ByteSize::from_bytes(1)) {
BacklogPull::Hit {
items,
advance_to,
more,
} => {
assert_eq!(items.len(), 1);
assert_eq!(advance_to, cv(1));
assert!(!more);
}
BacklogPull::Behind => panic!("expected Hit"),
}
}
#[test]
fn evict_below_raises_the_floor_and_later_pulls_go_behind() {
let b = backlog(u64::MAX);
for v in 1..=4 {
b.publish(cv(v), Some(cdc_with_payload(v, 100)));
}
b.evict_below(cv(2));
assert!(
matches!(b.pull(cv(1), cv(4), ByteSize::from_mib(1)), BacklogPull::Behind),
"a cursor below the raised floor must be sent to the loader"
);
match b.pull(cv(2), cv(4), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
..
} => assert_eq!(items.len(), 2),
BacklogPull::Behind => panic!("entries above the floor must survive evict_below"),
}
}
#[test]
fn ceiling_eviction_drops_lowest_versions_first_and_raises_the_floor() {
let one = entry_bytes();
let b = backlog(one * 2);
for v in 1..=3 {
b.publish(cv(v), Some(cdc_with_payload(v, 100)));
}
assert!(matches!(b.pull(cv(0), cv(3), ByteSize::from_mib(1)), BacklogPull::Behind));
match b.pull(cv(1), cv(3), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
..
} => assert_eq!(items.len(), 2, "the two newest entries must survive"),
BacklogPull::Behind => panic!("expected Hit above the evicted floor"),
}
}
#[test]
fn ceiling_eviction_never_crosses_the_scan_anchor() {
let one = entry_bytes();
let b = FlowBacklog::new(ByteSize::from_bytes(one));
b.set_anchor(cv(1));
for v in 1..=3 {
b.publish(cv(v), Some(cdc_with_payload(v, 100)));
}
match b.pull(cv(1), cv(3), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
..
} => assert_eq!(items.len(), 2, "unscanned entries must survive the ceiling"),
BacklogPull::Behind => panic!("entries above the anchor must not be evicted"),
}
b.set_anchor(cv(3));
b.publish(cv(4), Some(cdc_with_payload(4, 100)));
assert!(
matches!(b.pull(cv(1), cv(4), ByteSize::from_mib(1)), BacklogPull::Behind),
"once the anchor passes them, over-ceiling entries must be evicted lowest-first"
);
}
#[test]
fn out_of_order_publish_below_the_floor_is_ignored() {
let b = backlog(u64::MAX);
b.publish(cv(101), Some(cdc_with_payload(101, 10)));
b.publish(cv(99), Some(cdc_with_payload(99, 10)));
assert!(matches!(b.pull(cv(98), cv(101), ByteSize::from_mib(1)), BacklogPull::Behind));
match b.pull(cv(100), cv(101), ByteSize::from_mib(1)) {
BacklogPull::Hit {
items,
..
} => assert_eq!(items.len(), 1),
BacklogPull::Behind => panic!("expected Hit"),
}
}
#[test]
fn notify_fires_once_until_disarmed() {
let fired = Arc::new(AtomicUsize::new(0));
let b = backlog(u64::MAX);
let counter = fired.clone();
b.set_waker(move || {
counter.fetch_add(1, Ordering::SeqCst);
});
b.notify();
b.notify();
b.notify();
assert_eq!(fired.load(Ordering::SeqCst), 1, "repeat notifies while armed must coalesce");
b.disarm();
b.notify();
assert_eq!(fired.load(Ordering::SeqCst), 2, "a disarmed backlog must wake again");
}
#[test]
fn byte_accounting_balances_across_publish_replace_and_eviction() {
let one = entry_bytes();
let b = backlog(u64::MAX);
b.publish(cv(1), Some(cdc_with_payload(1, 100)));
b.publish(cv(2), Some(cdc_with_payload(2, 100)));
b.publish(cv(2), Some(cdc_with_payload(2, 300)));
let mut out = Vec::new();
b.collect(&mut out);
let bytes = out
.iter()
.find(|s| s.scope == "flow_backlog" && s.metric == "bytes")
.map(|s| s.reading.as_f64())
.expect("bytes sample");
assert_eq!(bytes, (one + one + 200) as f64, "replacing an entry must swap its tally, not add");
b.evict_below(cv(2));
let mut out = Vec::new();
b.collect(&mut out);
let bytes = out
.iter()
.find(|s| s.scope == "flow_backlog" && s.metric == "bytes")
.map(|s| s.reading.as_f64())
.expect("bytes sample");
assert_eq!(bytes, 0.0, "evicting every entry must zero the tally");
}
}