use std::time::{Duration, Instant};
use dynomite::cluster::apl::{get_apl_ann, ClusterState, NodeRole, RingPoint};
use dynomite::embed::events::PeerId;
use dynomite::events::TokenRange;
use dynomite::hashkit::DynToken;
use gen_fsm::{Action, EventType, FsmHandler, TimeoutKind, Transition};
use throttle_core::{SystemClock, Throttle};
pub const DEFAULT_REAP_TOMBSTONES_AFTER_SECONDS: u64 = 86_400;
pub const DEFAULT_REAP_SIBLINGS_AFTER_SECONDS: u64 = 604_800;
pub const DEFAULT_REAP_MAX_PER_CYCLE: u64 = 10_000;
pub const DEFAULT_REAP_INTERVAL_SECONDS: u64 = 300;
pub const DEFAULT_REAPS_PER_SEC: u64 = 100;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReaperConfig {
pub reap_tombstones_after_seconds: u64,
pub reap_siblings_after_seconds: u64,
pub reap_max_per_cycle: u64,
pub reap_interval_seconds: u64,
pub reaps_per_sec: u64,
}
impl Default for ReaperConfig {
fn default() -> Self {
Self {
reap_tombstones_after_seconds: DEFAULT_REAP_TOMBSTONES_AFTER_SECONDS,
reap_siblings_after_seconds: DEFAULT_REAP_SIBLINGS_AFTER_SECONDS,
reap_max_per_cycle: DEFAULT_REAP_MAX_PER_CYCLE,
reap_interval_seconds: DEFAULT_REAP_INTERVAL_SECONDS,
reaps_per_sec: DEFAULT_REAPS_PER_SEC,
}
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum KeyKind {
Live,
Tombstone,
Sibling,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ScannedKey {
pub partition_idx: usize,
pub key: Vec<u8>,
pub kind: KeyKind,
pub age: Duration,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ReaperCycleComplete {
pub bucket: Vec<u8>,
pub reaped: u64,
pub scanned: u64,
pub duration: Duration,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ReaperOutcome {
Stopped,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum State {
Idle,
Scanning,
Reaping,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Event {
Tick,
KeyScanned(ScannedKey),
NextSegmentDone,
KeyReaped,
BatchAcked,
CycleError(String),
Shutdown,
}
pub struct ReaperHandler {
bucket: Vec<u8>,
config: ReaperConfig,
partitions: Vec<TokenRange>,
throttle: Throttle<SystemClock>,
partition_idx: usize,
batch: Vec<ScannedKey>,
reaped_this_cycle: u64,
scanned_this_cycle: u64,
outstanding_reaps: u64,
cycle_started_at: Option<Instant>,
last_complete: Option<ReaperCycleComplete>,
last_error: Option<String>,
last_state: State,
}
impl ReaperHandler {
#[must_use]
pub fn new(bucket: Vec<u8>) -> Self {
Self::with_config(bucket, ReaperConfig::default())
}
#[must_use]
pub fn with_config(bucket: Vec<u8>, config: ReaperConfig) -> Self {
let rate = config.reaps_per_sec.max(1);
let throttle = Throttle::new(rate, rate);
Self {
bucket,
config,
partitions: Vec::new(),
throttle,
partition_idx: 0,
batch: Vec::new(),
reaped_this_cycle: 0,
scanned_this_cycle: 0,
outstanding_reaps: 0,
cycle_started_at: None,
last_complete: None,
last_error: None,
last_state: State::Idle,
}
}
#[must_use]
pub fn with_partitions(mut self, partitions: Vec<TokenRange>) -> Self {
self.partitions = partitions;
self
}
pub fn set_partitions(&mut self, partitions: Vec<TokenRange>) {
self.partitions = partitions;
}
pub fn refresh_partitions_from_cluster(
&mut self,
cluster: &ClusterState,
local_peer: PeerId,
n: usize,
) {
let ring = cluster.ring();
if ring.is_empty() {
self.partitions.clear();
return;
}
let mut out: Vec<TokenRange> = Vec::new();
for (idx, point) in ring.iter().enumerate() {
let apl = get_apl_ann(cluster, point.token, n);
let is_primary = apl
.iter()
.any(|p| p.peer_id == local_peer && p.role == NodeRole::Primary);
if !is_primary {
continue;
}
let next = (idx + 1) % ring.len();
let start = ring_token_to_dyntoken(point);
let end = ring_token_to_dyntoken(&ring[next]);
out.push(TokenRange::new(start, end));
}
self.partitions = out;
}
#[must_use]
pub fn bucket(&self) -> &[u8] {
&self.bucket
}
#[must_use]
pub const fn config(&self) -> &ReaperConfig {
&self.config
}
#[must_use]
pub fn partitions(&self) -> &[TokenRange] {
&self.partitions
}
#[must_use]
pub const fn partition_idx(&self) -> usize {
self.partition_idx
}
#[must_use]
pub fn current_partition(&self) -> Option<&TokenRange> {
self.partitions.get(self.partition_idx)
}
#[must_use]
pub fn batch_len(&self) -> usize {
self.batch.len()
}
#[must_use]
pub const fn outstanding_reaps(&self) -> u64 {
self.outstanding_reaps
}
#[must_use]
pub const fn reaped_this_cycle(&self) -> u64 {
self.reaped_this_cycle
}
#[must_use]
pub const fn scanned_this_cycle(&self) -> u64 {
self.scanned_this_cycle
}
#[must_use]
pub const fn last_state(&self) -> State {
self.last_state
}
#[must_use]
pub const fn last_complete(&self) -> Option<&ReaperCycleComplete> {
self.last_complete.as_ref()
}
pub fn take_last_complete(&mut self) -> Option<ReaperCycleComplete> {
self.last_complete.take()
}
pub fn take_batch(&mut self) -> Vec<ScannedKey> {
std::mem::take(&mut self.batch)
}
pub fn try_admit_reap(&self) -> bool {
self.throttle.try_acquire(1)
}
#[must_use]
pub fn cycle_elapsed(&self) -> Duration {
match self.cycle_started_at {
Some(t) => t.elapsed(),
None => Duration::ZERO,
}
}
#[must_use]
pub fn is_reap_candidate(&self, key: &ScannedKey) -> bool {
match key.kind {
KeyKind::Live => false,
KeyKind::Tombstone => {
key.age >= Duration::from_secs(self.config.reap_tombstones_after_seconds)
}
KeyKind::Sibling => {
key.age >= Duration::from_secs(self.config.reap_siblings_after_seconds)
}
}
}
fn record_state(&mut self, state: State) {
self.last_state = state;
}
fn idle_state_timeout(&self) -> Duration {
Duration::from_secs(self.config.reap_interval_seconds.max(1))
}
fn begin_cycle(&mut self) {
self.partition_idx = 0;
self.batch.clear();
self.reaped_this_cycle = 0;
self.scanned_this_cycle = 0;
self.outstanding_reaps = 0;
self.cycle_started_at = Some(Instant::now());
self.last_error = None;
}
fn finish_cycle(&mut self) {
let duration = self
.cycle_started_at
.map_or(Duration::ZERO, |t| t.elapsed());
self.last_complete = Some(ReaperCycleComplete {
bucket: self.bucket.clone(),
reaped: self.reaped_this_cycle,
scanned: self.scanned_this_cycle,
duration,
});
self.cycle_started_at = None;
}
fn handle_idle(&mut self, ev: &Event) -> Transition<Self> {
match ev {
Event::Tick => {
if self.partitions.is_empty() {
self.begin_cycle();
self.finish_cycle();
return Transition::Keep(vec![Action::set_state_timeout(
self.idle_state_timeout(),
)]);
}
self.begin_cycle();
Transition::Next(State::Scanning, vec![])
}
Event::Shutdown => Transition::Stop(ReaperOutcome::Stopped),
_ => Transition::Keep(vec![]),
}
}
fn handle_scanning(&mut self, ev: Event) -> Transition<Self> {
match ev {
Event::KeyScanned(key) => {
self.scanned_this_cycle = self.scanned_this_cycle.saturating_add(1);
if self.batch.len() as u64 >= self.config.reap_max_per_cycle {
return Transition::Keep(vec![]);
}
if self.is_reap_candidate(&key) {
self.batch.push(key);
}
Transition::Keep(vec![])
}
Event::NextSegmentDone => {
self.partition_idx = self.partition_idx.saturating_add(1);
if self.partition_idx >= self.partitions.len() {
self.outstanding_reaps = self.batch.len() as u64;
if self.outstanding_reaps == 0 {
self.finish_cycle();
return Transition::Next(
State::Idle,
vec![Action::set_state_timeout(self.idle_state_timeout())],
);
}
return Transition::Next(State::Reaping, vec![]);
}
Transition::Keep(vec![])
}
Event::CycleError(reason) => {
self.last_error = Some(reason);
self.finish_cycle();
Transition::Next(
State::Idle,
vec![Action::set_state_timeout(self.idle_state_timeout())],
)
}
Event::Shutdown => Transition::Stop(ReaperOutcome::Stopped),
_ => Transition::Keep(vec![]),
}
}
fn handle_reaping(&mut self, ev: Event) -> Transition<Self> {
match ev {
Event::KeyReaped => {
self.reaped_this_cycle = self.reaped_this_cycle.saturating_add(1);
self.outstanding_reaps = self.outstanding_reaps.saturating_sub(1);
Transition::Keep(vec![])
}
Event::BatchAcked => {
self.finish_cycle();
Transition::Next(
State::Idle,
vec![Action::set_state_timeout(self.idle_state_timeout())],
)
}
Event::CycleError(reason) => {
self.last_error = Some(reason);
self.finish_cycle();
Transition::Next(
State::Idle,
vec![Action::set_state_timeout(self.idle_state_timeout())],
)
}
Event::Shutdown => Transition::Stop(ReaperOutcome::Stopped),
_ => Transition::Keep(vec![]),
}
}
}
fn ring_token_to_dyntoken(point: &RingPoint) -> DynToken {
let narrow = u32::try_from(point.token).unwrap_or(u32::MAX);
DynToken::from_u32(narrow)
}
impl FsmHandler for ReaperHandler {
type State = State;
type Event = Event;
type Reply = ();
type Stop = ReaperOutcome;
fn initial(&self) -> State {
State::Idle
}
fn on_enter(&mut self, state: State) -> Transition<Self> {
self.record_state(state);
match state {
State::Idle => {
Transition::Keep(vec![Action::set_state_timeout(self.idle_state_timeout())])
}
State::Scanning | State::Reaping => Transition::Keep(vec![]),
}
}
fn handle(&mut self, state: State, _et: EventType, ev: Event) -> Transition<Self> {
self.record_state(state);
match state {
State::Idle => self.handle_idle(&ev),
State::Scanning => self.handle_scanning(ev),
State::Reaping => self.handle_reaping(ev),
}
}
fn on_timeout(&mut self, state: State, kind: TimeoutKind) -> Transition<Self> {
self.record_state(state);
if state == State::Idle && matches!(kind, TimeoutKind::State) {
return Transition::Keep(vec![Action::post_internal(Event::Tick)]);
}
Transition::Keep(vec![])
}
}
#[cfg(test)]
mod tests {
use super::*;
fn cfg() -> ReaperConfig {
ReaperConfig {
reap_tombstones_after_seconds: 10,
reap_siblings_after_seconds: 100,
reap_max_per_cycle: 4,
reap_interval_seconds: 60,
reaps_per_sec: 1_000_000,
}
}
fn handler() -> ReaperHandler {
ReaperHandler::with_config(b"users".to_vec(), cfg()).with_partitions(vec![
TokenRange::new(DynToken::from_u32(0), DynToken::from_u32(100)),
TokenRange::new(DynToken::from_u32(100), DynToken::from_u32(200)),
])
}
fn key(idx: usize, kind: KeyKind, age_secs: u64) -> ScannedKey {
ScannedKey {
partition_idx: idx,
key: format!("k{idx}-{age_secs}").into_bytes(),
kind,
age: Duration::from_secs(age_secs),
}
}
#[test]
fn idle_entry_arms_state_timeout() {
let mut h = handler();
let t = h.on_enter(State::Idle);
match t {
Transition::Keep(actions) => {
let secs = cfg().reap_interval_seconds;
let found = actions.iter().any(
|a| matches!(a, Action::SetStateTimeout(d) if *d == Duration::from_secs(secs)),
);
assert!(found, "expected SetStateTimeout; got {actions:?}");
}
other => panic!("expected Keep, got {other:?}"),
}
}
#[test]
fn idle_tick_with_empty_partitions_stays_idle() {
let mut h = ReaperHandler::with_config(b"empty".to_vec(), cfg());
let t = h.handle(State::Idle, EventType::Cast, Event::Tick);
match t {
Transition::Keep(_) => {}
other => panic!("expected Keep on empty partitions, got {other:?}"),
}
let rec = h
.take_last_complete()
.expect("empty cycle should still emit audit");
assert_eq!(rec.bucket, b"empty");
assert_eq!(rec.scanned, 0);
assert_eq!(rec.reaped, 0);
}
#[test]
fn scanning_old_tombstone_is_queued() {
let mut h = handler();
let _ = h.handle(State::Idle, EventType::Cast, Event::Tick);
let _ = h.handle(
State::Scanning,
EventType::Cast,
Event::KeyScanned(key(0, KeyKind::Tombstone, 60)),
);
assert_eq!(h.batch_len(), 1);
assert_eq!(h.scanned_this_cycle(), 1);
}
#[test]
fn scanning_live_key_is_never_queued() {
let mut h = handler();
let _ = h.handle(State::Idle, EventType::Cast, Event::Tick);
let _ = h.handle(
State::Scanning,
EventType::Cast,
Event::KeyScanned(key(0, KeyKind::Live, 1_000_000)),
);
assert_eq!(h.batch_len(), 0);
assert_eq!(h.scanned_this_cycle(), 1);
}
#[test]
fn batch_capped_at_reap_max_per_cycle() {
let mut h = handler();
let _ = h.handle(State::Idle, EventType::Cast, Event::Tick);
for i in 0..10u64 {
let mut k = key(0, KeyKind::Tombstone, 60);
k.key = format!("k{i}").into_bytes();
let _ = h.handle(State::Scanning, EventType::Cast, Event::KeyScanned(k));
}
assert_eq!(h.batch_len() as u64, cfg().reap_max_per_cycle);
assert_eq!(h.scanned_this_cycle(), 10);
}
#[test]
fn next_segment_done_advances_idx() {
let mut h = handler();
let _ = h.handle(State::Idle, EventType::Cast, Event::Tick);
assert_eq!(h.partition_idx(), 0);
let _ = h.handle(State::Scanning, EventType::Cast, Event::NextSegmentDone);
assert_eq!(h.partition_idx(), 1);
}
#[test]
fn shutdown_stops_from_any_state() {
for state in [State::Idle, State::Scanning, State::Reaping] {
let mut h = handler();
let t = h.handle(state, EventType::Cast, Event::Shutdown);
match t {
Transition::Stop(ReaperOutcome::Stopped) => {}
other => panic!("expected Stop from {state:?}, got {other:?}"),
}
}
}
}