use std::collections::{BTreeSet, HashMap};
use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
use crate::Key;
use crate::error::{DbError, DbResult};
use crate::hook::TypedWriteHook;
use crate::sync;
use crate::tree_meta::CollectionMeta;
pub const TTL_RETRY_BACKOFF_MS: u64 = 60_000;
#[derive(Clone, Copy, Debug)]
pub(crate) enum TtlAt {
Precomputed(Option<u64>),
Derive,
}
pub(crate) fn catch_deadline<R>(f: impl FnOnce() -> R) -> Option<R> {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)).ok()
}
#[derive(Default, Debug, Clone, Copy, PartialEq, Eq)]
pub struct SweepOutcome {
pub deleted: usize,
pub failed: usize,
pub more: bool,
}
pub(crate) fn log_sweep_failures(collection: Option<&str>, outcome: &SweepOutcome) {
if outcome.failed == 0 {
return;
}
match collection {
Some(name) => tracing::warn!(
collection = name,
failed = outcome.failed,
"ttl sweep rescheduled entries after read/decode failure"
),
None => tracing::warn!(
failed = outcome.failed,
"ttl sweep rescheduled entries after read/decode failure"
),
}
}
pub struct TtlSweepGate {
enabled: AtomicBool,
active: AtomicUsize,
drained: sync::Condvar,
drained_lock: sync::Mutex<()>,
#[cfg(test)]
disable_wait_pause: std::sync::Mutex<Option<ttl_test_support::ParkPause>>,
}
impl TtlSweepGate {
pub fn new() -> Self {
Self {
enabled: AtomicBool::new(true),
active: AtomicUsize::new(0),
drained: sync::Condvar::new(),
drained_lock: sync::Mutex::new(()),
#[cfg(test)]
disable_wait_pause: std::sync::Mutex::new(None),
}
}
pub fn begin_pass(&self) -> Option<TtlSweepPass<'_>> {
if !self.enabled.load(Ordering::Acquire) {
return None;
}
self.active.fetch_add(1, Ordering::SeqCst);
if !self.enabled.load(Ordering::SeqCst) {
self.end_pass();
return None;
}
Some(TtlSweepPass { gate: self })
}
fn end_pass(&self) {
if self.active.fetch_sub(1, Ordering::SeqCst) == 1 {
let _g = sync::lock(&self.drained_lock);
self.drained.notify_all();
}
}
pub fn set_enabled(&self, enabled: bool) {
if enabled {
self.enabled.store(true, Ordering::Release);
return;
}
self.enabled.store(false, Ordering::SeqCst);
let mut g = sync::lock(&self.drained_lock);
while self.active.load(Ordering::SeqCst) > 0 {
#[cfg(test)]
ttl_test_support::wait_pause(&self.disable_wait_pause);
g = sync::cv_wait(&self.drained, g);
}
}
#[cfg(test)]
pub(crate) fn arm_disable_wait(&self) -> ttl_test_support::ParkPause {
ttl_test_support::arm_pause(&self.disable_wait_pause)
}
}
pub struct TtlSweepPass<'a> {
gate: &'a TtlSweepGate,
}
impl Drop for TtlSweepPass<'_> {
fn drop(&mut self) {
self.gate.end_pass();
}
}
impl Default for TtlSweepGate {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
pub(crate) mod ttl_test_support {
use std::sync::{Arc, Barrier, Mutex};
pub(crate) type ParkPause = (Arc<Barrier>, Arc<Barrier>);
pub(super) fn arm_pause(slot: &Mutex<Option<ParkPause>>) -> ParkPause {
let pause = (Arc::new(Barrier::new(2)), Arc::new(Barrier::new(2)));
*slot.lock().expect("ttl pause lock") = Some(pause.clone());
pause
}
pub(super) fn wait_pause(slot: &Mutex<Option<ParkPause>>) {
if let Some((arrived, release)) = slot.lock().expect("ttl pause lock").take() {
arrived.wait();
release.wait();
}
}
}
pub struct TtlSweepState {
pub gate: TtlSweepGate,
pub cursor: AtomicUsize,
needs_rebuild_before_enable: AtomicBool,
}
impl TtlSweepState {
pub fn new() -> Self {
Self {
gate: TtlSweepGate::new(),
cursor: AtomicUsize::new(0),
needs_rebuild_before_enable: AtomicBool::new(false),
}
}
pub fn set_enabled(&self, enabled: bool) -> DbResult<()> {
if enabled && self.needs_rebuild_before_enable.load(Ordering::Acquire) {
return Err(DbError::Client(
"TTL sweep requires rebuild_ttl_index after replication follower duty",
));
}
self.gate.set_enabled(enabled);
Ok(())
}
pub fn mark_replication_follower(&self) {
self.gate.set_enabled(false);
self.needs_rebuild_before_enable
.store(true, Ordering::Release);
}
pub fn clear_rebuild_requirement(&self) {
self.needs_rebuild_before_enable
.store(false, Ordering::Release);
}
}
impl Default for TtlSweepState {
fn default() -> Self {
Self::new()
}
}
#[derive(Clone, Copy)]
pub struct TtlPolicy<T: CollectionMeta + ?Sized> {
pub revision: u32,
pub deadline: fn(&T, &T::SelfId) -> Option<u64>,
}
#[repr(transparent)]
pub struct TtlHook<T, H> {
pub inner: H,
_marker: core::marker::PhantomData<fn() -> T>,
}
impl<T, H> TtlHook<T, H> {
pub const fn new(inner: H) -> Self {
Self {
inner,
_marker: core::marker::PhantomData,
}
}
}
#[cfg(feature = "typed-tree")]
impl<T, H> TypedWriteHook<T::SelfId, T> for TtlHook<T, H>
where
T: CollectionMeta,
H: TypedWriteHook<T::SelfId, T>,
{
const WATCHES: bool = H::WATCHES;
const EXPIRES: bool = T::TTL.is_some();
const TTL_REVISION: u32 = match T::TTL {
Some(p) => p.revision,
None => 0,
};
fn watch_channel(&self) -> Option<&std::sync::Arc<crate::watch::WatchChannel<T::SelfId>>> {
self.inner.watch_channel()
}
fn watch_old(&self) -> bool {
self.inner.watch_old()
}
const NEEDS_OLD_VALUE: bool = H::NEEDS_OLD_VALUE;
const NEEDS_INIT: bool = H::NEEDS_INIT;
const NEEDS_WRITE: bool = H::NEEDS_WRITE;
const PROJECTS: bool = H::PROJECTS;
const COUNTS: bool = H::COUNTS;
const CLAIMS: bool = H::CLAIMS;
#[cfg(feature = "armour")]
const EMITS: bool = H::EMITS;
fn deadline(&self, key: &T::SelfId, value: &T) -> Option<u64> {
match T::TTL {
Some(p) => (p.deadline)(value, key),
None => None,
}
}
fn on_write(&self, key: &T::SelfId, old: Option<&T>, new: Option<&T>) {
self.inner.on_write(key, old, new);
}
fn on_init(&self, key: &T::SelfId, value: &T) {
self.inner.on_init(key, value);
}
fn project(
&self,
key: &T::SelfId,
value: &T,
sink: &mut crate::projection::ProjectionSink<'_>,
) -> DbResult<()> {
self.inner.project(key, value, sink)
}
fn projection_targets(&self) -> Option<&crate::projection::IndexTargets> {
self.inner.projection_targets()
}
fn projection_limits(&self) -> (usize, usize) {
self.inner.projection_limits()
}
fn counter_project(
&self,
key: &T::SelfId,
value: &T,
sink: &mut crate::counter::CounterSink<'_>,
) -> DbResult<()> {
self.inner.counter_project(key, value, sink)
}
fn counter_targets(&self) -> Option<&crate::counter::CounterTargets> {
self.inner.counter_targets()
}
fn claim_project(
&self,
key: &T::SelfId,
value: &T,
sink: &mut crate::unique::ClaimSink<'_>,
) -> DbResult<()> {
self.inner.claim_project(key, value, sink)
}
fn claim_targets(&self) -> Option<&crate::unique::UniqueTargets> {
self.inner.claim_targets()
}
#[cfg(feature = "armour")]
fn stream_project(
&self,
key: &T::SelfId,
old: Option<&T>,
new: Option<&T>,
sink: &mut crate::EmitSink<'_>,
) -> DbResult<()> {
self.inner.stream_project(key, old, new, sink)
}
#[cfg(feature = "armour")]
fn stream_targets(&self) -> Option<&crate::QueueTargets> {
self.inner.stream_targets()
}
}
#[derive(Copy, Clone)]
pub struct ByBytes<K: Key>(pub K);
impl<K: Key> PartialEq for ByBytes<K> {
fn eq(&self, other: &Self) -> bool {
self.0.as_bytes() == other.0.as_bytes()
}
}
impl<K: Key> Eq for ByBytes<K> {}
impl<K: Key> PartialOrd for ByBytes<K> {
fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
Some(self.cmp(other))
}
}
impl<K: Key> Ord for ByBytes<K> {
fn cmp(&self, other: &Self) -> std::cmp::Ordering {
self.0.as_bytes().cmp(other.0.as_bytes())
}
}
impl<K: Key> std::hash::Hash for ByBytes<K> {
fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
self.0.as_bytes().hash(state);
}
}
pub struct TtlIndex<K: Key> {
due: BTreeSet<(u64, ByBytes<K>)>,
at: HashMap<ByBytes<K>, u64>,
}
impl<K: Key> TtlIndex<K> {
pub fn new() -> Self {
Self {
due: BTreeSet::new(),
at: HashMap::new(),
}
}
pub fn record(&mut self, key: &K, expires_at: Option<u64>) {
let kb = ByBytes(*key);
if let Some(old) = self.at.remove(&kb) {
self.due.remove(&(old, kb));
}
if let Some(d) = expires_at {
self.at.insert(kb, d);
self.due.insert((d, kb));
}
}
pub fn len(&self) -> usize {
self.at.len()
}
pub fn deadline_of(&self, key: &K) -> Option<u64> {
self.at.get(&ByBytes(*key)).copied()
}
pub fn pop_due(&mut self, now_ms: u64) -> Option<(u64, K)> {
loop {
let &(d, _) = self.due.iter().next()?;
if d > now_ms {
return None;
}
let Some((d, kb)) = self.due.pop_first() else {
continue;
};
match self.at.get(&kb) {
Some(&at_d) if at_d == d => {
self.at.remove(&kb);
return Some((d, kb.0));
}
_ => continue,
}
}
}
pub fn has_valid_due(&self, now_ms: u64) -> bool {
for &(d, ref kb) in self.due.iter() {
if d > now_ms {
break;
}
if self.at.get(kb) == Some(&d) {
return true;
}
}
false
}
}
pub fn any_shard_has_valid_due<K: Key>(
slots: &[crate::sync::Mutex<TtlIndex<K>>],
now_ms: u64,
) -> bool {
slots
.iter()
.any(|m| crate::sync::lock(m).has_valid_due(now_ms))
}
pub(crate) struct TtlSlots<K: Key> {
slots: std::sync::OnceLock<Box<[crate::sync::Mutex<TtlIndex<K>>]>>,
sweep: TtlSweepState,
}
impl<K: Key> TtlSlots<K> {
pub fn new() -> Self {
Self {
slots: std::sync::OnceLock::new(),
sweep: TtlSweepState::new(),
}
}
#[inline]
pub fn get(&self) -> Option<&[crate::sync::Mutex<TtlIndex<K>>]> {
self.slots.get().map(|s| &s[..])
}
pub fn install_and_clear(&self, n: usize) -> &[crate::sync::Mutex<TtlIndex<K>>] {
if self.slots.get().is_none() {
let fresh: Box<[crate::sync::Mutex<TtlIndex<K>>]> = (0..n)
.map(|_| crate::sync::Mutex::new(TtlIndex::new()))
.collect();
let _ = self.slots.set(fresh);
}
let slots = self.slots.get().expect("ttl slots installed");
for m in slots.iter() {
*crate::sync::lock(m) = TtlIndex::new();
}
slots
}
#[inline]
pub fn record(&self, sid: usize, key: &K, expires_at: Option<u64>) {
if let Some(slots) = self.get() {
crate::sync::lock(&slots[sid]).record(key, expires_at);
}
}
#[inline]
pub fn next_shard(&self, shard_count: usize) -> usize {
self.sweep.cursor.fetch_add(1, Ordering::Relaxed) % shard_count
}
pub fn any_due(&self, now_ms: u64) -> bool {
self.get()
.map(|s| any_shard_has_valid_due(s, now_ms))
.unwrap_or(false)
}
pub fn begin_pass(&self) -> Option<TtlSweepPass<'_>> {
self.sweep.gate.begin_pass()
}
pub fn set_enabled(&self, enabled: bool) -> DbResult<()> {
self.sweep.set_enabled(enabled)
}
#[cfg(test)]
pub(crate) fn arm_disable_wait(&self) -> ttl_test_support::ParkPause {
self.sweep.gate.arm_disable_wait()
}
pub fn mark_replication_follower(&self) {
self.sweep.mark_replication_follower();
}
pub fn clear_rebuild_requirement(&self) {
self.sweep.clear_rebuild_requirement();
}
pub fn debug_len(&self) -> usize {
self.get()
.map(|s| s.iter().map(|m| crate::sync::lock(m).len()).sum())
.unwrap_or(0)
}
pub fn debug_deadline(&self, sid: usize, key: &K) -> Option<u64> {
self.get()
.and_then(|s| crate::sync::lock(&s[sid]).deadline_of(key))
}
pub fn debug_inject_stale(&self, sid: usize, deadline: u64, key: K) {
if let Some(slots) = self.get() {
crate::sync::lock(&slots[sid]).inject_stale_due_for_test(deadline, key);
}
}
}
impl<K: Key> Default for TtlSlots<K> {
fn default() -> Self {
Self::new()
}
}
impl<K: Key> TtlIndex<K> {
pub(crate) fn inject_stale_due_for_test(&mut self, deadline: u64, key: K) {
self.due.insert((deadline, ByBytes(key)));
}
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "typed-tree")]
#[test]
fn watch_channel_forwards_through_ttl_hook() {
use crate::hook::NoHook;
use crate::watch::WatchHook;
#[derive(armour_core::GetType)]
struct WatchRow;
impl CollectionMeta for WatchRow {
type SelfId = [u8; 8];
const NAME: &'static str = "ttl_watch_forwarding";
}
let hook =
TtlHook::<WatchRow, _>::new(WatchHook::<WatchRow, NoHook, true>::with_old(NoHook));
type Hook = TtlHook<WatchRow, WatchHook<WatchRow, NoHook, true>>;
assert!(<Hook as TypedWriteHook<[u8; 8], WatchRow>>::WATCHES);
assert!(hook.watch_channel().is_some());
}
#[test]
fn record_insert_replace_remove() {
let mut idx = TtlIndex::<[u8; 4]>::new();
let k = [1, 2, 3, 4];
assert_eq!(idx.len(), 0);
idx.record(&k, Some(100));
assert_eq!(idx.len(), 1);
assert_eq!(idx.deadline_of(&k), Some(100));
idx.record(&k, Some(200));
assert_eq!(idx.len(), 1);
assert_eq!(idx.deadline_of(&k), Some(200));
idx.record(&k, None);
assert_eq!(idx.len(), 0);
assert_eq!(idx.deadline_of(&k), None);
}
#[test]
fn has_valid_due_ignores_stale_heap_head() {
let mut idx = TtlIndex::<[u8; 4]>::new();
let k = [1, 2, 3, 4];
idx.record(&k, Some(200));
idx.inject_stale_due_for_test(100, k);
assert!(!idx.has_valid_due(150));
assert!(idx.has_valid_due(250));
}
#[test]
fn gate_disabled_refuses_new_passes_and_reenables() {
let gate = TtlSweepGate::new();
assert!(gate.begin_pass().is_some());
gate.set_enabled(false);
assert!(
gate.begin_pass().is_none(),
"a disabled gate must refuse new passes"
);
gate.set_enabled(true);
assert!(gate.begin_pass().is_some());
}
#[test]
fn gate_disable_blocks_until_in_flight_pass_drops() {
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::time::{Duration, Instant};
let gate = Arc::new(TtlSweepGate::new());
let released = Arc::new(AtomicBool::new(false));
let pass = gate.begin_pass().expect("gate starts enabled");
let t = std::thread::spawn({
let gate = Arc::clone(&gate);
let released = Arc::clone(&released);
move || {
gate.set_enabled(false);
assert!(
released.load(Ordering::SeqCst),
"set_enabled(false) returned while a pass was still in flight"
);
}
});
std::thread::sleep(Duration::from_millis(50));
let t0 = Instant::now();
released.store(true, Ordering::SeqCst);
drop(pass);
t.join().expect("disabler thread");
assert!(t0.elapsed() < Duration::from_secs(5));
}
#[test]
fn gate_disable_never_races_a_starting_pass() {
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
let gate = Arc::new(TtlSweepGate::new());
let in_flight = Arc::new(AtomicUsize::new(0));
let stop = Arc::new(AtomicBool::new(false));
let sweepers: Vec<_> = (0..4)
.map(|_| {
let gate = Arc::clone(&gate);
let in_flight = Arc::clone(&in_flight);
let stop = Arc::clone(&stop);
std::thread::spawn(move || {
while !stop.load(Ordering::Relaxed) {
if let Some(pass) = gate.begin_pass() {
in_flight.fetch_add(1, Ordering::SeqCst);
std::hint::spin_loop();
in_flight.fetch_sub(1, Ordering::SeqCst);
drop(pass);
}
}
})
})
.collect();
for _ in 0..2_000 {
gate.set_enabled(false);
assert_eq!(
in_flight.load(Ordering::SeqCst),
0,
"a sweep pass was in flight after set_enabled(false) returned"
);
gate.set_enabled(true);
}
stop.store(true, Ordering::Relaxed);
gate.set_enabled(true);
for s in sweepers {
s.join().expect("sweeper thread");
}
}
#[cfg(all(feature = "typed-tree", feature = "armour"))]
mod claim_forwarding {
use std::sync::Arc;
use crate::TypedWriteHook;
use crate::projection::{IndexTargets, Projected};
use crate::unique::{ClaimSink, UniqueTargets};
use crate::unique_tests::{ClaimFwdKey, ClaimFwdRow, open_claim_ns};
use super::TtlHook;
struct StreamProbe(Arc<std::sync::Mutex<Vec<(Option<u64>, Option<u64>)>>>);
impl TypedWriteHook<ClaimFwdKey, ClaimFwdRow> for StreamProbe {
const EMITS: bool = true;
fn on_write(&self, _: &ClaimFwdKey, _: Option<&ClaimFwdRow>, _: Option<&ClaimFwdRow>) {}
fn stream_project(
&self,
_: &ClaimFwdKey,
old: Option<&ClaimFwdRow>,
new: Option<&ClaimFwdRow>,
_: &mut crate::EmitSink<'_>,
) -> crate::DbResult<()> {
self.0
.lock()
.unwrap()
.push((old.map(|row| row.0), new.map(|row| row.0)));
Ok(())
}
}
#[test]
fn ttl_hook_forwards_a_claim_end_to_end() {
let (_dir, _db, ns) = open_claim_ns();
let targets = UniqueTargets::new().with(&ns);
let ns_for_project = Arc::clone(&ns);
let inner =
Projected::<ClaimFwdKey, ClaimFwdRow>::new(IndexTargets::new(), |_, _, _| Ok(()))
.and_unique(targets.clone(), move |_k, _v: &ClaimFwdRow, sink| {
sink.claim(&ns_for_project, ClaimFwdKey::new(7))
});
let hook = TtlHook::<ClaimFwdRow, _>::new(inner);
let mut sink = ClaimSink::new(&targets, 16, 1024);
hook.claim_project(&ClaimFwdKey::new(0), &ClaimFwdRow(1), &mut sink)
.expect("claim_project");
assert_eq!(
sink.claims().len(),
1,
"the inner Projected's claim must reach TtlHook's caller"
);
}
#[test]
fn ttl_hook_forwards_each_optional_stream_snapshot_once() {
let seen = Arc::new(std::sync::Mutex::new(Vec::new()));
let hook = TtlHook::<ClaimFwdRow, _>::new(StreamProbe(Arc::clone(&seen)));
let targets = crate::QueueTargets::new();
let mut sink = crate::EmitSink::new(&targets, 1, 1);
hook.stream_project(&ClaimFwdKey::new(1), Some(&ClaimFwdRow(3)), None, &mut sink)
.unwrap();
assert_eq!(*seen.lock().unwrap(), vec![(Some(3), None)]);
}
}
}