use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use reifydb_runtime::context::clock::Clock;
use reifydb_value::value::duration::Duration;
use crate::lifecycle::{class::RetentionClass, progress::Progress, task::LifecycleTask};
#[derive(Clone)]
pub struct RetentionStartupGate {
inner: Arc<Inner>,
}
struct Inner {
clock: Clock,
armed_at_nanos: u64,
grace: Duration,
skipped_slices: AtomicU64,
}
impl RetentionStartupGate {
pub fn arm(clock: Clock, grace: Duration) -> Self {
let armed_at_nanos = clock.now().to_nanos();
Self {
inner: Arc::new(Inner {
clock,
armed_at_nanos,
grace,
skipped_slices: AtomicU64::new(0),
}),
}
}
pub fn open(clock: Clock) -> Self {
Self::arm(clock, Duration::zero())
}
pub fn is_open(&self) -> bool {
if self.inner.grace.is_zero() {
return true;
}
let now = self.inner.clock.now();
match now.checked_sub(self.inner.grace) {
Some(released) => released.to_nanos() >= self.inner.armed_at_nanos,
None => false,
}
}
pub fn record_skip(&self) {
self.inner.skipped_slices.fetch_add(1, Ordering::Relaxed);
}
pub fn skipped_slices(&self) -> u64 {
self.inner.skipped_slices.load(Ordering::Relaxed)
}
pub fn grace(&self) -> Duration {
self.inner.grace
}
}
pub struct Gated<T: LifecycleTask> {
inner: T,
gate: RetentionStartupGate,
}
impl<T: LifecycleTask> Gated<T> {
pub fn new(inner: T, gate: RetentionStartupGate) -> Self {
Self {
inner,
gate,
}
}
}
impl<T: LifecycleTask> LifecycleTask for Gated<T> {
fn name(&self) -> &'static str {
self.inner.name()
}
fn interval(&self) -> Duration {
self.inner.interval()
}
fn classes(&self) -> &'static [RetentionClass] {
self.inner.classes()
}
fn run_slice(&mut self) -> Progress {
if !self.gate.is_open() {
self.gate.record_skip();
return Progress::Exhausted;
}
self.inner.run_slice()
}
}
#[cfg(test)]
mod tests {
use std::sync::{
Arc,
atomic::{AtomicU64, Ordering},
};
use reifydb_runtime::context::clock::{Clock, MockClock};
use reifydb_value::value::duration::Duration;
use super::{Gated, RetentionStartupGate};
use crate::lifecycle::{class::RetentionClass, progress::Progress, task::LifecycleTask};
fn mock() -> (Clock, MockClock) {
let mock = MockClock::from_millis(0);
(Clock::Mock(mock.clone()), mock)
}
#[test]
fn a_gate_armed_with_a_grace_period_starts_closed() {
let (clock, _mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(300).unwrap());
assert!(!gate.is_open(), "a freshly armed gate must hold reclamation back for its grace period");
}
#[test]
fn a_zero_grace_gate_is_open_immediately() {
let (clock, _mock) = mock();
let gate = RetentionStartupGate::open(clock);
assert!(gate.is_open(), "zero grace must mean no gating at all");
}
#[test]
fn the_gate_opens_once_the_grace_period_has_elapsed() {
let (clock, mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(60).unwrap());
assert!(!gate.is_open(), "precondition: still inside the grace window");
mock.advance_secs(59);
assert!(!gate.is_open(), "one second short of the grace period must still gate");
mock.advance_secs(1);
assert!(gate.is_open(), "once the grace period elapses the gate must release reclamation");
}
#[test]
fn the_gate_stays_open_once_released() {
let (clock, mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(60).unwrap());
mock.advance_secs(60);
assert!(gate.is_open(), "precondition: released");
mock.advance_secs(3600);
assert!(gate.is_open(), "a released gate must remain open for the life of the process");
}
struct CountingTask {
slices: Arc<AtomicU64>,
}
impl LifecycleTask for CountingTask {
fn name(&self) -> &'static str {
"counting"
}
fn interval(&self) -> Duration {
Duration::from_seconds(1).unwrap()
}
fn classes(&self) -> &'static [RetentionClass] {
&[RetentionClass::RowTtl]
}
fn run_slice(&mut self) -> Progress {
self.slices.fetch_add(1, Ordering::SeqCst);
Progress::Exhausted
}
}
#[test]
fn a_gated_task_does_no_work_while_the_gate_is_closed() {
let (clock, _mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(300).unwrap());
let slices = Arc::new(AtomicU64::new(0));
let mut task = Gated::new(
CountingTask {
slices: slices.clone(),
},
gate.clone(),
);
assert_eq!(task.run_slice(), Progress::Exhausted, "a gated slice must not ask the lane for a catch-up");
assert_eq!(slices.load(Ordering::SeqCst), 0, "the wrapped task must not run at all while gated");
assert_eq!(gate.skipped_slices(), 1, "the skip must be counted so a gated class is not read as idle");
}
#[test]
fn a_gated_task_runs_normally_once_the_gate_opens() {
let (clock, mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(60).unwrap());
let slices = Arc::new(AtomicU64::new(0));
let mut task = Gated::new(
CountingTask {
slices: slices.clone(),
},
gate,
);
task.run_slice();
assert_eq!(slices.load(Ordering::SeqCst), 0, "precondition: gated");
mock.advance_secs(60);
task.run_slice();
assert_eq!(slices.load(Ordering::SeqCst), 1, "once released the wrapped task must run");
}
#[test]
fn gating_preserves_the_wrapped_class_identity() {
let (clock, _mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(300).unwrap());
let task = Gated::new(
CountingTask {
slices: Arc::new(AtomicU64::new(0)),
},
gate,
);
assert_eq!(task.name(), "counting", "gating must not rename the class");
assert_eq!(task.interval(), Duration::from_seconds(1).unwrap(), "gating must not change the cadence");
}
#[test]
fn the_gate_counts_the_slices_it_turned_away() {
let (clock, _mock) = mock();
let gate = RetentionStartupGate::arm(clock, Duration::from_seconds(300).unwrap());
for _ in 0..3 {
if !gate.is_open() {
gate.record_skip();
}
}
assert_eq!(gate.skipped_slices(), 3, "every skipped slice must be counted, not silently dropped");
}
}