use std::collections::{BTreeMap, BTreeSet};
use std::fmt;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use crate::branch::current_timestamp;
use crate::store::NodeStore;
use crate::tree::{Hash, Node};
use crate::ttl::entry::{ExpiryTimestamp, StampedEntry, TtlEntry};
use crate::wal::{Mutation, WalBuffer, WalError};
pub(super) const MAX_GENERATION: u64 = (1_u64 << 60) - 1;
static NEXT_GLOBAL_GENERATION: AtomicU64 = AtomicU64::new(0);
pub(super) trait WallClock: fmt::Debug + Send + Sync {
fn now(&self) -> ExpiryTimestamp;
}
#[derive(Debug)]
pub(super) struct SystemWallClock;
impl WallClock for SystemWallClock {
fn now(&self) -> ExpiryTimestamp {
current_timestamp()
}
}
pub(super) trait DeadlineScheduler {
fn schedule(&mut self, delay: Duration, generation: Generation) -> Result<(), ArmError>;
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) struct Generation(u64);
impl Generation {
pub(super) const fn from_value(value: u64) -> Option<Self> {
if value <= MAX_GENERATION {
Some(Self(value))
} else {
None
}
}
pub(super) const fn value(self) -> u64 {
self.0
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Arm {
generation: Generation,
minimum: ExpiryTimestamp,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(super) enum ArmError {
GenerationExhausted,
TokenUnrepresentable,
SchedulerUnavailable,
DelayUnrepresentable,
}
impl std::error::Error for ArmError {}
impl fmt::Display for ArmError {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::GenerationExhausted => formatter.write_str("TTL deadline generation exhausted"),
Self::TokenUnrepresentable => formatter
.write_str("TTL deadline generation is not a beamr immediate small integer"),
Self::SchedulerUnavailable => {
formatter.write_str("beamr timer wheel unavailable for TTL deadline")
}
Self::DelayUnrepresentable => {
formatter.write_str("TTL deadline relative delay is unrepresentable")
}
}
}
}
pub(super) use super::expiry_metrics::ExpiryMetrics;
#[derive(Debug)]
pub(super) struct ExpiryBuild {
buckets: BTreeMap<ExpiryTimestamp, BTreeSet<Vec<u8>>>,
inspected: u64,
}
#[derive(Debug)]
pub(super) struct TransitionPlan {
key: Vec<u8>,
old: Option<ExpiryTimestamp>,
new: Option<ExpiryTimestamp>,
}
#[derive(Debug)]
pub(super) struct AppliedTransition {
plan: TransitionPlan,
previous_arm: Option<Arm>,
previous_pending: bool,
index_mutated: bool,
minimum_changed: bool,
}
#[derive(Debug)]
enum GenerationSource {
Global,
#[cfg(test)]
Local(u64),
}
#[derive(Debug)]
pub(super) struct ExpiryState {
buckets: BTreeMap<ExpiryTimestamp, BTreeSet<Vec<u8>>>,
arm: Option<Arm>,
generation_source: GenerationSource,
arm_pending: bool,
metrics: ExpiryMetrics,
clock: Arc<dyn WallClock>,
}
impl ExpiryState {
pub(super) fn empty(clock: Arc<dyn WallClock>) -> Self {
Self {
buckets: BTreeMap::new(),
arm: None,
generation_source: GenerationSource::Global,
arm_pending: false,
metrics: ExpiryMetrics::default(),
clock,
}
}
#[cfg(test)]
pub(super) fn empty_for_test(clock: Arc<dyn WallClock>) -> Self {
let mut state = Self::empty(clock);
state.generation_source = GenerationSource::Local(0);
state
}
pub(super) fn from_build(clock: Arc<dyn WallClock>, build: ExpiryBuild) -> Self {
let mut state = Self::empty(clock);
state.metrics.rebuild_entries = build.inspected;
state.buckets = build.buckets;
state.arm_pending = !state.buckets.is_empty();
state
}
pub(super) fn clock_now(&self) -> ExpiryTimestamp {
self.clock.now()
}
#[cfg(test)]
pub(super) const fn metrics(&self) -> ExpiryMetrics {
self.metrics
}
pub(super) fn prepare(
&mut self,
key: &[u8],
old_encoded: Option<&[u8]>,
new_encoded: Option<&[u8]>,
) -> Result<TransitionPlan, WalError> {
let old = match old_encoded {
Some(bytes) => {
self.metrics.old_value_decodes = self.metrics.old_value_decodes.saturating_add(1);
self.metrics.old_value_decode_bytes = self
.metrics
.old_value_decode_bytes
.saturating_add(u64::try_from(bytes.len()).unwrap_or(u64::MAX));
let expiry = decode_expiry(bytes)?;
if expiry.is_some() {
self.metrics.old_value_expiring =
self.metrics.old_value_expiring.saturating_add(1);
}
expiry
}
None => None,
};
let new = new_encoded.map(decode_expiry).transpose()?.flatten();
Ok(TransitionPlan {
key: key.to_vec(),
old,
new,
})
}
pub(super) fn apply(&mut self, plan: TransitionPlan) -> AppliedTransition {
let previous_arm = self.arm;
let previous_pending = self.arm_pending;
let index_mutated = plan.old != plan.new;
let mut minimum_changed = false;
if index_mutated {
let before = self.minimum();
if let Some(deadline) = plan.old {
remove_key(&mut self.buckets, deadline, &plan.key);
}
if let Some(deadline) = plan.new {
self.buckets
.entry(deadline)
.or_default()
.insert(plan.key.clone());
}
if before != self.minimum() {
self.invalidate_for_minimum_change();
minimum_changed = true;
}
}
AppliedTransition {
plan,
previous_arm,
previous_pending,
index_mutated,
minimum_changed,
}
}
pub(super) const fn confirm(&mut self, applied: &AppliedTransition) {
if applied.index_mutated {
self.metrics.index_mutations = self.metrics.index_mutations.saturating_add(1);
}
if applied.minimum_changed {
self.metrics.minimum_changes = self.metrics.minimum_changes.saturating_add(1);
}
}
pub(super) fn rollback(&mut self, applied: AppliedTransition) {
if applied.plan.old != applied.plan.new {
if let Some(deadline) = applied.plan.new {
remove_key(&mut self.buckets, deadline, &applied.plan.key);
}
if let Some(deadline) = applied.plan.old {
self.buckets
.entry(deadline)
.or_default()
.insert(applied.plan.key);
}
}
self.arm = applied.previous_arm;
self.arm_pending = applied.previous_pending;
self.metrics.current_arms = u64::from(self.arm.is_some());
}
pub(super) fn install_merge_build(&mut self, build: ExpiryBuild) {
self.buckets = build.buckets;
self.metrics.merge_rebuild_entries = self
.metrics
.merge_rebuild_entries
.saturating_add(build.inspected);
self.invalidate_for_rebuild();
}
pub(super) fn arm_pending<S>(&mut self, scheduler: &mut S) -> Result<(), ArmError>
where
S: DeadlineScheduler,
{
if !self.arm_pending {
return Ok(());
}
self.arm_pending = false;
let Some(minimum) = self.minimum() else {
self.arm = None;
self.metrics.current_arms = 0;
return Ok(());
};
let generation = self.next_generation()?;
let wall_now = self.clock.now();
let delay = relative_delay(minimum, wall_now)?;
self.metrics.scheduler_arm_requests = self.metrics.scheduler_arm_requests.saturating_add(1);
if let Err(error) = scheduler.schedule(delay, generation) {
self.metrics.arm_refusals = self.metrics.arm_refusals.saturating_add(1);
self.arm = None;
self.metrics.current_arms = 0;
return Err(error);
}
self.metrics.physical_arms = self.metrics.physical_arms.saturating_add(1);
self.arm = Some(Arm {
generation,
minimum,
});
self.metrics.current_arms = 1;
Ok(())
}
pub(super) fn begin_delivery(&mut self, generation: Generation) -> bool {
self.metrics.deadline_deliveries = self.metrics.deadline_deliveries.saturating_add(1);
self.metrics.actor_wakes = self.metrics.actor_wakes.saturating_add(1);
if self.arm.map(|arm| arm.generation) != Some(generation) {
self.metrics.stale_drops = self.metrics.stale_drops.saturating_add(1);
return false;
}
self.arm = None;
self.metrics.current_arms = 0;
self.metrics.accepted_deliveries = self.metrics.accepted_deliveries.saturating_add(1);
true
}
pub(super) fn detach_due(&mut self, now: ExpiryTimestamp) -> Vec<(ExpiryTimestamp, Vec<u8>)> {
let mut due = Vec::new();
while self.minimum().is_some_and(|deadline| deadline <= now) {
let Some((deadline, keys)) = self.buckets.pop_first() else {
break;
};
due.extend(keys.into_iter().map(|key| (deadline, key)));
}
due
}
pub(super) fn restore_key(&mut self, deadline: ExpiryTimestamp, key: Vec<u8>) {
self.buckets.entry(deadline).or_default().insert(key);
}
pub(super) fn finish_delivery(&mut self) {
self.arm_pending = !self.buckets.is_empty();
}
pub(super) const fn inspected_key(&mut self) {
self.metrics.inspected_keys = self.metrics.inspected_keys.saturating_add(1);
}
pub(super) const fn delete_attempt(&mut self) {
self.metrics.delete_attempts = self.metrics.delete_attempts.saturating_add(1);
}
pub(super) const fn deleted(&mut self) {
self.metrics.deletes = self.metrics.deletes.saturating_add(1);
}
pub(super) const fn invalidate_shutdown(&mut self) {
self.arm = None;
self.arm_pending = false;
self.metrics.current_arms = 0;
}
#[cfg(test)]
pub(super) fn current_generation(&self) -> Option<Generation> {
self.arm.map(|arm| arm.generation)
}
fn minimum(&self) -> Option<ExpiryTimestamp> {
self.buckets
.first_key_value()
.map(|(deadline, _keys)| *deadline)
}
fn next_generation(&mut self) -> Result<Generation, ArmError> {
match &mut self.generation_source {
GenerationSource::Global => NEXT_GLOBAL_GENERATION
.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
current
.checked_add(1)
.filter(|next| *next <= MAX_GENERATION)
})
.map(|previous| Generation(previous.saturating_add(1)))
.map_err(|_current| ArmError::GenerationExhausted),
#[cfg(test)]
GenerationSource::Local(current) => {
let next = current
.checked_add(1)
.filter(|next| *next <= MAX_GENERATION)
.ok_or(ArmError::GenerationExhausted)?;
*current = next;
Ok(Generation(next))
}
}
}
#[cfg(test)]
pub(super) const fn force_generation_for_test(&mut self, value: u64) {
self.generation_source = GenerationSource::Local(value);
}
fn invalidate_for_minimum_change(&mut self) {
self.arm = None;
self.metrics.current_arms = 0;
self.arm_pending = !self.buckets.is_empty();
}
fn invalidate_for_rebuild(&mut self) {
self.arm = None;
self.metrics.current_arms = 0;
self.arm_pending = !self.buckets.is_empty();
self.metrics.minimum_changes = self.metrics.minimum_changes.saturating_add(1);
}
}
pub(super) fn relative_delay(
minimum: ExpiryTimestamp,
wall_now: ExpiryTimestamp,
) -> Result<Duration, ArmError> {
let nanos = minimum.saturating_sub(wall_now);
let subsec_nanos =
u32::try_from(nanos % 1_000_000_000).map_err(|_error| ArmError::DelayUnrepresentable)?;
Ok(Duration::new(nanos / 1_000_000_000, subsec_nanos))
}
pub(super) fn rebuild<S>(
store: &S,
committed_root: Option<Hash>,
buffer: Option<&WalBuffer>,
) -> Result<ExpiryBuild, WalError>
where
S: NodeStore + ?Sized,
{
let mut view = BTreeMap::new();
if let Some(root) = committed_root {
walk_rebuild_entries(store, root, &mut view)?;
}
if let Some(buffer) = buffer {
for mutation in buffer {
match mutation {
Mutation::Put { key, value } => {
view.insert(key.clone(), value.clone());
}
Mutation::Delete { key } => {
view.remove(key);
}
}
}
}
let inspected = u64::try_from(view.len()).unwrap_or(u64::MAX);
let mut buckets: BTreeMap<ExpiryTimestamp, BTreeSet<Vec<u8>>> = BTreeMap::new();
for (key, value) in view {
if let Some(deadline) = decode_expiry(&value)? {
buckets.entry(deadline).or_default().insert(key);
}
}
Ok(ExpiryBuild { buckets, inspected })
}
pub(super) fn decode_expiry(bytes: &[u8]) -> Result<Option<ExpiryTimestamp>, WalError> {
if let Some(entry) = StampedEntry::decode(bytes).map_err(codec_error)? {
return Ok(entry.expires_at());
}
Ok(TtlEntry::decode(bytes)
.map_err(codec_error)?
.and_then(|entry| entry.expires_at()))
}
fn walk_rebuild_entries<S>(
store: &S,
root: Hash,
out: &mut BTreeMap<Vec<u8>, Vec<u8>>,
) -> Result<(), WalError>
where
S: NodeStore + ?Sized,
{
let mut stack = vec![root];
while let Some(hash) = stack.pop() {
let node = store
.get(&hash)
.map_err(|error| WalError::TreeError(error.to_string()))?
.ok_or_else(|| WalError::TreeError(format!("missing expiry rebuild node {hash:?}")))?;
match &*node {
Node::Leaf(leaf) => {
for (key, value) in leaf.entries() {
out.insert(key.clone(), value.clone());
}
}
Node::Internal(internal) => {
for (_separator, child) in internal.children() {
stack.push(*child);
}
}
}
}
Ok(())
}
fn remove_key(
buckets: &mut BTreeMap<ExpiryTimestamp, BTreeSet<Vec<u8>>>,
deadline: ExpiryTimestamp,
key: &[u8],
) {
if let Some(keys) = buckets.get_mut(&deadline) {
keys.remove(key);
if keys.is_empty() {
buckets.remove(&deadline);
}
}
}
fn codec_error(error: impl fmt::Display) -> WalError {
WalError::TreeError(error.to_string())
}