Skip to main content

feoxdb/core/
ttl_sweep.rs

1use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
2use std::sync::{Arc, Weak};
3use std::thread::{self, JoinHandle};
4use std::time::{Duration, Instant};
5
6use rand::Rng;
7
8use crate::constants::Operation;
9use crate::core::store::FeoxStore;
10
11/// Configuration for TTL cleaner background thread
12#[derive(Clone, Debug)]
13pub struct TtlConfig {
14    /// Number of keys to sample per batch
15    pub sample_size: usize,
16    /// Continue sampling if expiry rate exceeds this threshold (0.0-1.0)
17    pub expiry_threshold: f32,
18    /// Maximum iterations per cleaning run
19    pub max_iterations: usize,
20    /// Maximum time to spend per cleaning run
21    pub max_time_per_run: Duration,
22    /// Sleep interval between cleaning runs
23    pub sleep_interval: Duration,
24    /// Whether TTL cleaner is enabled
25    pub enabled: bool,
26}
27
28impl Default for TtlConfig {
29    fn default() -> Self {
30        Self {
31            sample_size: 100,
32            expiry_threshold: 0.25,
33            max_iterations: 16,
34            max_time_per_run: Duration::from_millis(1),
35            sleep_interval: Duration::from_millis(1000),
36            enabled: false,
37        }
38    }
39}
40
41impl TtlConfig {
42    /// Create a default configuration for persistent stores
43    pub fn default_persistent() -> Self {
44        Self {
45            enabled: true,
46            ..Default::default()
47        }
48    }
49
50    /// Create a default configuration for memory-only stores
51    pub fn default_memory() -> Self {
52        Self {
53            enabled: true,
54            ..Default::default()
55        }
56    }
57}
58
59/// Background thread that periodically sweeps expired TTL keys
60pub struct TtlSweeper {
61    /// Weak reference to the store to avoid circular references
62    store: Weak<FeoxStore>,
63    /// Configuration
64    config: TtlConfig,
65    /// Shutdown flag
66    shutdown: Arc<AtomicBool>,
67    /// Thread handle
68    handle: Option<JoinHandle<()>>,
69    /// Statistics
70    stats: TtlSweeperStats,
71}
72
73/// Statistics for TTL sweeper operations
74pub struct TtlSweeperStats {
75    /// Total keys sampled
76    pub total_sampled: Arc<AtomicU64>,
77    /// Total keys expired
78    pub total_expired: Arc<AtomicU64>,
79    /// Total cleaning runs
80    pub total_runs: Arc<AtomicU64>,
81    /// Last run timestamp (nanoseconds)
82    pub last_run: Arc<AtomicU64>,
83}
84
85impl TtlSweeperStats {
86    fn new() -> Self {
87        Self {
88            total_sampled: Arc::new(AtomicU64::new(0)),
89            total_expired: Arc::new(AtomicU64::new(0)),
90            total_runs: Arc::new(AtomicU64::new(0)),
91            last_run: Arc::new(AtomicU64::new(0)),
92        }
93    }
94}
95
96impl TtlSweeper {
97    /// Create a new TTL sweeper
98    pub fn new(store: Weak<FeoxStore>, config: TtlConfig) -> Self {
99        Self {
100            store,
101            config,
102            shutdown: Arc::new(AtomicBool::new(false)),
103            handle: None,
104            stats: TtlSweeperStats::new(),
105        }
106    }
107
108    /// Start the background sweeper thread
109    pub fn start(&mut self) {
110        if !self.config.enabled {
111            return;
112        }
113
114        let store = self.store.clone();
115        let config = self.config.clone();
116        let shutdown = self.shutdown.clone();
117        let stats = TtlSweeperStats {
118            total_sampled: self.stats.total_sampled.clone(),
119            total_expired: self.stats.total_expired.clone(),
120            total_runs: self.stats.total_runs.clone(),
121            last_run: self.stats.last_run.clone(),
122        };
123
124        let handle = thread::spawn(move || {
125            run_sweeper_loop(store, config, shutdown, stats);
126        });
127
128        self.handle = Some(handle);
129    }
130
131    /// Stop the background sweeper thread
132    pub fn stop(&mut self) {
133        self.shutdown.store(true, Ordering::Release);
134
135        if let Some(handle) = self.handle.take() {
136            if handle.thread().id() != thread::current().id() {
137                let _ = handle.join();
138            }
139        }
140    }
141
142    /// Get sweeper statistics
143    pub fn stats(&self) -> SweeperSnapshot {
144        SweeperSnapshot {
145            total_sampled: self.stats.total_sampled.load(Ordering::Relaxed),
146            total_expired: self.stats.total_expired.load(Ordering::Relaxed),
147            total_runs: self.stats.total_runs.load(Ordering::Relaxed),
148            last_run: self.stats.last_run.load(Ordering::Relaxed),
149        }
150    }
151}
152
153impl Drop for TtlSweeper {
154    fn drop(&mut self) {
155        self.stop();
156    }
157}
158
159/// Snapshot of sweeper statistics
160#[derive(Debug, Clone)]
161pub struct SweeperSnapshot {
162    pub total_sampled: u64,
163    pub total_expired: u64,
164    pub total_runs: u64,
165    pub last_run: u64,
166}
167
168/// Main sweeper loop that runs in the background thread
169fn run_sweeper_loop(
170    store: Weak<FeoxStore>,
171    config: TtlConfig,
172    shutdown: Arc<AtomicBool>,
173    stats: TtlSweeperStats,
174) {
175    while !shutdown.load(Ordering::Acquire) {
176        // Sleep between runs
177        thread::sleep(config.sleep_interval);
178
179        // Try to get strong reference to store
180        let Some(store) = store.upgrade() else {
181            // Store has been dropped, exit
182            break;
183        };
184
185        // Perform sweeping run
186        let start = Instant::now();
187        let mut iterations = 0;
188        let mut total_sampled = 0;
189        let mut total_expired = 0;
190
191        loop {
192            // Sample and expire a batch
193            let (sampled, expired) = sample_and_expire_batch(&store, &config);
194            total_sampled += sampled;
195            total_expired += expired;
196            iterations += 1;
197
198            // Calculate expiry rate
199            let expiry_rate = if sampled > 0 {
200                expired as f32 / sampled as f32
201            } else {
202                0.0
203            };
204
205            // Check stop conditions
206            if expiry_rate < config.expiry_threshold {
207                break; // Few expired keys, we're done
208            }
209            if iterations >= config.max_iterations {
210                break; // Bounded iterations
211            }
212            if start.elapsed() > config.max_time_per_run {
213                break; // Bounded time
214            }
215        }
216
217        // Update statistics
218        if total_sampled > 0 {
219            stats
220                .total_sampled
221                .fetch_add(total_sampled, Ordering::Relaxed);
222            stats
223                .total_expired
224                .fetch_add(total_expired, Ordering::Relaxed);
225            stats.total_runs.fetch_add(1, Ordering::Relaxed);
226            stats.last_run.store(
227                std::time::SystemTime::now()
228                    .duration_since(std::time::UNIX_EPOCH)
229                    .unwrap()
230                    .as_nanos() as u64,
231                Ordering::Relaxed,
232            );
233        }
234
235        // Check shutdown flag again
236        if shutdown.load(Ordering::Acquire) {
237            break;
238        }
239    }
240}
241
242/// Sample keys and expire those that have exceeded their TTL
243fn sample_and_expire_batch(store: &Arc<FeoxStore>, config: &TtlConfig) -> (u64, u64) {
244    let now = store.get_timestamp_pub();
245    let mut expired = 0;
246    let mut rng = rand::rng();
247
248    // Get access to the hash table
249    let hash_table = store.get_hash_table();
250
251    let candidates = sample_ttl_entries(hash_table, config.sample_size, &mut rng);
252    let sampled = candidates.len() as u64;
253
254    for (key, record) in candidates {
255        let ttl_expiry = record.ttl_expiry.load(Ordering::Relaxed);
256
257        if ttl_expiry > 0 && ttl_expiry < now {
258            #[cfg(test)]
259            crate::test_hooks::pause_at(crate::test_hooks::TTL_AFTER_EXPIRED_SAMPLE);
260
261            let old_value_len = record.value_len;
262            let record_size = record.calculate_size();
263            let retired = match hash_table.entry(key.clone()) {
264                scc::hash_map::Entry::Occupied(entry) => {
265                    let current_expiry = record.ttl_expiry.load(Ordering::Acquire);
266                    if Arc::ptr_eq(entry.get(), &record)
267                        && current_expiry > 0
268                        && current_expiry < now
269                    {
270                        record.retired_at.store(now, Ordering::Release);
271                        record.refcount.store(0, Ordering::Release);
272                        store.remove_from_tree(&key);
273                        let _ = entry.remove();
274                        true
275                    } else {
276                        false
277                    }
278                }
279                scc::hash_map::Entry::Vacant(_) => false,
280            };
281
282            if retired {
283                store.remove_cached(&key, &record);
284                store.note_expired_record(record_size);
285                expired += 1;
286
287                if let Some(wb) = store.get_write_buffer() {
288                    let _ = wb.add_write(Operation::Delete, record, old_value_len);
289                }
290            }
291        }
292    }
293
294    (sampled, expired)
295}
296
297#[cfg(test)]
298pub(crate) fn sample_and_expire_for_test(store: &Arc<FeoxStore>) -> (u64, u64) {
299    sample_and_expire_batch(
300        store,
301        &TtlConfig {
302            sample_size: 1,
303            ..TtlConfig::default()
304        },
305    )
306}
307
308fn sample_ttl_entries<R: Rng + ?Sized>(
309    hash_table: &scc::HashMap<Vec<u8>, Arc<crate::core::record::Record>, ahash::RandomState>,
310    sample_size: usize,
311    rng: &mut R,
312) -> Vec<(Vec<u8>, Arc<crate::core::record::Record>)> {
313    if sample_size == 0 {
314        return Vec::new();
315    }
316
317    let mut candidates = Vec::with_capacity(sample_size.min(hash_table.len()));
318    let mut seen = 0usize;
319    hash_table.scan(|key: &Vec<u8>, value: &Arc<crate::core::record::Record>| {
320        if value.ttl_expiry.load(Ordering::Relaxed) > 0 {
321            seen += 1;
322            if candidates.len() < sample_size {
323                candidates.push((key.clone(), Arc::clone(value)));
324            } else {
325                let index = rng.random_range(0..seen);
326                if index < sample_size {
327                    candidates[index] = (key.clone(), Arc::clone(value));
328                }
329            }
330        }
331    });
332
333    candidates
334}
335
336#[cfg(test)]
337pub(crate) fn sample_ttl_keys_for_test(store: &FeoxStore, sample_size: usize) -> Vec<Vec<u8>> {
338    let mut rng = rand::rng();
339    sample_ttl_entries(store.get_hash_table(), sample_size, &mut rng)
340        .into_iter()
341        .map(|(key, _)| key)
342        .collect()
343}
344
345#[cfg(test)]
346#[path = "../tests/ttl_sweep_safety_tests.rs"]
347mod tests;