use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use std::collections::BTreeSet;
use std::sync::atomic::{AtomicU64, Ordering};
#[derive(
Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize, Default,
)]
pub struct Epoch(pub u64);
impl Epoch {
pub const ZERO: Epoch = Epoch(0);
#[inline]
pub fn next(self) -> Epoch {
Epoch(self.0.wrapping_add(1))
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct MaintenanceReceipt {
pub epoch: Epoch,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Snapshot {
pub epoch: Epoch,
pub commit_ts: mongreldb_types::hlc::HlcTimestamp,
}
impl Snapshot {
#[inline]
pub fn at(epoch: Epoch) -> Self {
Self {
epoch,
commit_ts: mongreldb_types::hlc::HlcTimestamp::ZERO,
}
}
#[inline]
pub fn at_hlc(epoch: Epoch, commit_ts: mongreldb_types::hlc::HlcTimestamp) -> Self {
Self { epoch, commit_ts }
}
#[inline]
pub fn unbounded() -> Self {
Self::at_hlc(Epoch(u64::MAX), mongreldb_types::hlc::HlcTimestamp::MAX)
}
#[inline]
pub fn uses_hlc_authority(&self) -> bool {
self.commit_ts != mongreldb_types::hlc::HlcTimestamp::ZERO
}
#[inline]
pub fn observes(&self, page_epoch: Epoch) -> bool {
page_epoch <= self.epoch
}
#[inline]
pub fn observes_version(
&self,
row_epoch: Epoch,
row_commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
) -> bool {
match (self.uses_hlc_authority(), row_commit_ts) {
(true, Some(ts)) => ts <= self.commit_ts,
(false, Some(_)) => row_epoch <= self.epoch,
(_, None) => row_epoch <= self.epoch,
}
}
#[inline]
pub fn observes_row(
&self,
row_epoch: Epoch,
row_commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
) -> bool {
self.observes_version(row_epoch, row_commit_ts)
}
#[inline]
pub fn version_is_newer(
a_epoch: Epoch,
a_commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
b_epoch: Epoch,
b_commit_ts: Option<mongreldb_types::hlc::HlcTimestamp>,
) -> bool {
match (a_commit_ts, b_commit_ts) {
(Some(a), Some(b)) => a > b,
_ => a_epoch > b_epoch,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct GcFloor {
pub transaction_snapshot: mongreldb_types::hlc::HlcTimestamp,
pub history_retention: mongreldb_types::hlc::HlcTimestamp,
pub backup_pitr: mongreldb_types::hlc::HlcTimestamp,
pub replication: mongreldb_types::hlc::HlcTimestamp,
pub read_generation: mongreldb_types::hlc::HlcTimestamp,
pub online_index_build: mongreldb_types::hlc::HlcTimestamp,
}
impl Default for GcFloor {
fn default() -> Self {
Self::ZERO
}
}
impl GcFloor {
pub const ZERO: Self = Self {
transaction_snapshot: mongreldb_types::hlc::HlcTimestamp::ZERO,
history_retention: mongreldb_types::hlc::HlcTimestamp::ZERO,
backup_pitr: mongreldb_types::hlc::HlcTimestamp::ZERO,
replication: mongreldb_types::hlc::HlcTimestamp::ZERO,
read_generation: mongreldb_types::hlc::HlcTimestamp::ZERO,
online_index_build: mongreldb_types::hlc::HlcTimestamp::ZERO,
};
#[inline]
pub fn floor(&self) -> mongreldb_types::hlc::HlcTimestamp {
[
self.transaction_snapshot,
self.history_retention,
self.backup_pitr,
self.replication,
self.read_generation,
self.online_index_build,
]
.into_iter()
.filter(|ts| *ts != mongreldb_types::hlc::HlcTimestamp::ZERO)
.min()
.unwrap_or(mongreldb_types::hlc::HlcTimestamp::ZERO)
}
pub fn sources(&self) -> [(&'static str, mongreldb_types::hlc::HlcTimestamp); 6] {
[
("transaction_snapshot", self.transaction_snapshot),
("history_retention", self.history_retention),
("backup_pitr", self.backup_pitr),
("replication", self.replication),
("read_generation", self.read_generation),
("online_index_build", self.online_index_build),
]
}
}
#[derive(Debug, Default)]
pub struct EpochClock {
current: AtomicU64,
}
impl EpochClock {
pub fn new(start: u64) -> Self {
Self {
current: AtomicU64::new(start),
}
}
#[inline]
pub fn now(&self) -> Epoch {
Epoch(self.current.load(Ordering::Acquire))
}
#[inline]
pub fn snapshot(&self) -> Snapshot {
Snapshot::at(self.now())
}
#[inline]
pub fn bump(&self) -> Epoch {
Epoch(self.current.fetch_add(1, Ordering::AcqRel) + 1)
}
pub fn advance_to(&self, e: Epoch) {
loop {
let cur = self.current.load(Ordering::Acquire);
if e.0 <= cur {
return;
}
let _ = self
.current
.compare_exchange(cur, e.0, Ordering::AcqRel, Ordering::Acquire);
}
}
}
#[derive(Debug)]
pub struct EpochAuthority {
assigned: AtomicU64,
visible: AtomicU64,
committed: AtomicU64,
pending: Mutex<BTreeSet<u64>>,
abandoned: Mutex<BTreeSet<u64>>,
}
pub(crate) struct EpochGuard<'a> {
authority: &'a EpochAuthority,
epoch: Epoch,
armed: bool,
}
impl<'a> EpochGuard<'a> {
pub(crate) fn new(authority: &'a EpochAuthority, epoch: Epoch) -> Self {
Self {
authority,
epoch,
armed: true,
}
}
pub(crate) fn disarm(&mut self) {
self.armed = false;
}
}
impl Drop for EpochGuard<'_> {
fn drop(&mut self) {
if self.armed {
self.authority.abandon(self.epoch);
}
}
}
impl EpochAuthority {
pub fn new(start: u64) -> Self {
Self {
assigned: AtomicU64::new(start),
visible: AtomicU64::new(start),
committed: AtomicU64::new(start),
pending: Mutex::new(BTreeSet::new()),
abandoned: Mutex::new(BTreeSet::new()),
}
}
#[inline]
pub fn visible(&self) -> Epoch {
Epoch(self.visible.load(Ordering::Acquire))
}
#[inline]
pub fn bump_assigned(&self) -> Epoch {
Epoch(self.assigned.fetch_add(1, Ordering::AcqRel) + 1)
}
pub fn publish_visible(&self, e: Epoch) {
let mut cur = self.visible.load(Ordering::Acquire);
while e.0 > cur {
match self
.visible
.compare_exchange_weak(cur, e.0, Ordering::AcqRel, Ordering::Acquire)
{
Ok(_) => break,
Err(actual) => cur = actual,
}
}
}
pub fn publish_in_order(&self, e: Epoch) {
raise_to(&self.committed, e.0);
let mut pending = self.pending.lock();
let mut abandoned = self.abandoned.lock();
pending.insert(e.0);
let mut vis = self.visible.load(Ordering::Acquire);
loop {
let next = vis + 1;
if pending.remove(&next) || abandoned.remove(&next) {
vis = next;
} else {
break;
}
}
drop(pending);
drop(abandoned);
self.publish_visible(Epoch(vis));
}
pub fn abandon(&self, e: Epoch) {
let mut pending = self.pending.lock();
let mut abandoned = self.abandoned.lock();
pending.remove(&e.0);
abandoned.insert(e.0);
let mut vis = self.visible.load(Ordering::Acquire);
loop {
let next = vis + 1;
if pending.remove(&next) || abandoned.remove(&next) {
vis = next;
} else {
break;
}
}
drop(pending);
drop(abandoned);
self.publish_visible(Epoch(vis));
}
pub fn set_recovered(&self, e: Epoch) {
self.assigned.store(e.0, Ordering::Release);
self.visible.store(e.0, Ordering::Release);
self.committed.store(e.0, Ordering::Release);
}
pub fn advance_recovered(&self, e: Epoch) {
raise_to(&self.assigned, e.0);
raise_to(&self.visible, e.0);
raise_to(&self.committed, e.0);
}
#[inline]
pub fn assigned(&self) -> Epoch {
Epoch(self.assigned.load(Ordering::Acquire))
}
#[inline]
pub fn committed(&self) -> Epoch {
Epoch(self.committed.load(Ordering::Acquire))
}
}
fn raise_to(cell: &AtomicU64, target: u64) {
let mut cur = cell.load(Ordering::Acquire);
while target > cur {
match cell.compare_exchange_weak(cur, target, Ordering::AcqRel, Ordering::Acquire) {
Ok(_) => break,
Err(actual) => cur = actual,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use mongreldb_types::hlc::HlcTimestamp;
fn hlc(physical_micros: u64) -> HlcTimestamp {
HlcTimestamp {
physical_micros,
logical: 0,
node_tiebreaker: 1,
}
}
#[test]
fn hlc_visibility_is_authoritative_when_stamped() {
let early = hlc(100);
let late = hlc(200);
let snap = Snapshot::at_hlc(Epoch(99), early);
assert!(snap.uses_hlc_authority());
assert!(snap.observes_version(Epoch(1), Some(early)));
assert!(!snap.observes_version(Epoch(2), Some(late)));
assert!(snap.observes_row(Epoch(1), Some(early)));
assert!(!snap.observes_row(Epoch(2), Some(late)));
let snap2 = Snapshot::at_hlc(Epoch(1), late);
assert!(snap2.observes_version(Epoch(1), Some(early)));
assert!(snap2.observes_version(Epoch(2), Some(late)));
}
#[test]
fn snapshot_hlc_hides_later_commit_ts_even_if_epoch_higher() {
let early = hlc(100);
let late = hlc(200);
let snap = Snapshot::at_hlc(Epoch(10), early);
assert!(!snap.observes_version(Epoch(5), Some(late)));
assert!(snap.observes_version(Epoch(5), Some(early)));
let legacy = Snapshot::at(Epoch(99));
assert!(!legacy.uses_hlc_authority());
assert!(legacy.observes_version(Epoch(1), Some(early)));
assert!(legacy.observes_version(Epoch(1), None));
assert!(!legacy.observes_version(Epoch(100), Some(early)));
}
#[test]
fn version_is_newer_prefers_hlc_when_both_stamped() {
let early = hlc(100);
let late = hlc(200);
assert!(Snapshot::version_is_newer(
Epoch(1),
Some(late),
Epoch(50),
Some(early)
));
assert!(!Snapshot::version_is_newer(
Epoch(50),
Some(early),
Epoch(1),
Some(late)
));
assert!(Snapshot::version_is_newer(Epoch(3), None, Epoch(2), None));
assert!(Snapshot::version_is_newer(
Epoch(3),
Some(early),
Epoch(2),
None
));
}
#[test]
fn hlc_clock_order_matches_visibility_order() {
let a = hlc(100);
let b = HlcTimestamp {
physical_micros: 100,
logical: 1,
node_tiebreaker: 1,
};
let c = HlcTimestamp {
physical_micros: 100,
logical: 1,
node_tiebreaker: 2,
};
assert!(a < b && b < c);
let snap = Snapshot::at_hlc(Epoch(1), b);
assert!(snap.observes_row(Epoch(1), Some(a)));
assert!(snap.observes_row(Epoch(1), Some(b)));
assert!(!snap.observes_row(Epoch(1), Some(c)));
}
#[test]
fn gc_floor_min_skips_zero_sources() {
let mut floor = GcFloor::ZERO;
assert_eq!(floor.floor(), HlcTimestamp::ZERO);
floor.transaction_snapshot = hlc(300);
floor.backup_pitr = hlc(100);
floor.replication = hlc(200);
assert_eq!(floor.floor(), hlc(100));
assert_eq!(floor.sources().len(), 6);
}
#[test]
fn clock_skew_rejects_timestamp_allocation() {
use mongreldb_types::hlc::HlcClock;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};
use std::sync::Arc;
use std::time::Duration;
let wall = Arc::new(AtomicU64::new(10_000));
let wall_src = {
let wall = Arc::clone(&wall);
Arc::new(move || wall.load(AtomicOrdering::Relaxed))
as mongreldb_types::hlc::WallClockSource
};
let clock = HlcClock::with_time_source(1, Duration::from_micros(1_000), wall_src);
let remote = HlcTimestamp {
physical_micros: 15_000,
logical: 0,
node_tiebreaker: 2,
};
assert!(
clock.observe(remote).is_err(),
"5ms skew must exceed 1ms bound"
);
assert!(
clock.now().is_err(),
"further allocation stays rejected after skew trip"
);
let _ = clock.next_after(HlcTimestamp::ZERO);
}
#[test]
fn snapshot_visibility_is_monotonic() {
let clock = EpochClock::new(10);
assert_eq!(clock.now(), Epoch(10));
let s = clock.snapshot();
assert!(s.observes(Epoch(10)));
assert!(!s.observes(Epoch(11)));
let next = clock.bump();
assert_eq!(next, Epoch(11));
assert!(clock.snapshot().observes(Epoch(11)));
}
#[test]
fn epoch_authority_assigned_and_visible_advance_in_order() {
let a = EpochAuthority::new(0);
assert_eq!(a.visible(), Epoch(0));
let e1 = a.bump_assigned();
let e2 = a.bump_assigned();
assert_eq!((e1, e2), (Epoch(1), Epoch(2)));
a.publish_visible(Epoch(2));
assert_eq!(a.visible(), Epoch(2));
a.publish_visible(Epoch(1));
assert_eq!(a.visible(), Epoch(2));
}
#[test]
fn publish_in_order_gates_until_gap_filled() {
let a = EpochAuthority::new(0);
let e1 = a.bump_assigned();
let e2 = a.bump_assigned();
let e3 = a.bump_assigned();
assert_eq!((e1, e2, e3), (Epoch(1), Epoch(2), Epoch(3)));
a.publish_in_order(e3);
assert_eq!(a.visible(), Epoch(0), "e3 cannot be visible before e1/e2");
a.publish_in_order(e2);
assert_eq!(a.visible(), Epoch(0), "e2 still gated on e1");
a.publish_in_order(e1);
assert_eq!(a.visible(), Epoch(3));
}
#[test]
fn abandon_unblocks_watermark() {
let a = EpochAuthority::new(0);
let e1 = a.bump_assigned();
let e2 = a.bump_assigned(); let e3 = a.bump_assigned();
a.publish_in_order(e3);
assert_eq!(a.visible(), Epoch(0), "e3 gated on e1 and e2");
a.publish_in_order(e1);
assert_eq!(a.visible(), Epoch(1), "e1 visible but e2 hole remains");
a.abandon(e2);
assert_eq!(a.visible(), Epoch(3), "abandoning e2 drains e3");
assert_eq!(a.committed(), Epoch(3));
}
#[test]
fn abandoned_latest_ticket_does_not_advance_commit_watermark() {
let a = EpochAuthority::new(4);
let abandoned = a.bump_assigned();
a.abandon(abandoned);
assert_eq!(a.visible(), Epoch(5));
assert_eq!(a.committed(), Epoch(4));
}
#[test]
fn abandon_before_publish_of_later_epochs() {
let a = EpochAuthority::new(0);
let e1 = a.bump_assigned();
let e2 = a.bump_assigned();
let e3 = a.bump_assigned();
a.abandon(e1);
assert_eq!(
a.visible(),
Epoch(1),
"e1 abandoned, watermark at 1, e2 pending"
);
a.publish_in_order(e2);
assert_eq!(a.visible(), Epoch(2));
a.publish_in_order(e3);
assert_eq!(a.visible(), Epoch(3));
}
}