1use std::sync::Arc;
16
17use tokio::time::{Duration, interval};
18use tokio_util::sync::CancellationToken;
19
20use crate::error::MemoryError;
21use crate::store::SqliteStore;
22use crate::types::MessageId;
23
24#[derive(Debug, Clone)]
28pub struct EvictionEntry {
29 pub id: MessageId,
31 pub created_at: String,
33 pub last_accessed: Option<String>,
35 pub access_count: u32,
37}
38
39pub trait EvictionPolicy: Send + Sync {
43 fn score(&self, entry: &EvictionEntry) -> f64;
48}
49
50#[derive(Debug, Clone, serde::Deserialize, serde::Serialize)]
52#[serde(default)]
53pub struct EvictionConfig {
54 pub policy: String,
56 pub max_entries: usize,
58 pub sweep_interval_secs: u64,
60}
61
62impl Default for EvictionConfig {
63 fn default() -> Self {
64 Self {
65 policy: "ebbinghaus".to_owned(),
66 max_entries: 0,
67 sweep_interval_secs: 3600,
68 }
69 }
70}
71
72pub struct EbbinghausPolicy {
87 retention_strength: f64,
88}
89
90impl EbbinghausPolicy {
91 #[must_use]
95 pub fn new(retention_strength: f64) -> Self {
96 Self { retention_strength }
97 }
98}
99
100impl Default for EbbinghausPolicy {
101 fn default() -> Self {
102 Self::new(86_400.0) }
104}
105
106impl EvictionPolicy for EbbinghausPolicy {
107 fn score(&self, entry: &EvictionEntry) -> f64 {
108 let now_secs = unix_now_secs();
109
110 let reference_secs = entry
111 .last_accessed
112 .as_deref()
113 .and_then(parse_sqlite_timestamp_secs)
114 .unwrap_or_else(|| parse_sqlite_timestamp_secs(&entry.created_at).unwrap_or(now_secs));
115
116 #[allow(clippy::cast_precision_loss)]
118 let t = now_secs.saturating_sub(reference_secs) as f64;
119 let n = f64::from(entry.access_count);
120
121 let denominator = (self.retention_strength * (1.0_f64 + n).ln()).max(1.0);
123 (-t / denominator).exp()
124 }
125}
126
127fn unix_now_secs() -> u64 {
128 std::time::SystemTime::now()
129 .duration_since(std::time::UNIX_EPOCH)
130 .map(|d| d.as_secs())
131 .unwrap_or(0)
132}
133
134fn parse_sqlite_timestamp_secs(s: &str) -> Option<u64> {
138 let s = s.trim();
140 if s.len() < 19 {
141 return None;
142 }
143 let year: u64 = s[0..4].parse().ok()?;
144 let month: u64 = s[5..7].parse().ok()?;
145 let day: u64 = s[8..10].parse().ok()?;
146 let hour: u64 = s[11..13].parse().ok()?;
147 let min: u64 = s[14..16].parse().ok()?;
148 let sec: u64 = s[17..19].parse().ok()?;
149
150 let is_leap = |y: u64| (y.is_multiple_of(4) && !y.is_multiple_of(100)) || y.is_multiple_of(400);
153 let days_in_month = |y: u64, m: u64| -> u64 {
154 match m {
155 1 | 3 | 5 | 7 | 8 | 10 | 12 => 31,
156 4 | 6 | 9 | 11 => 30,
157 2 => {
158 if is_leap(y) {
159 29
160 } else {
161 28
162 }
163 }
164 _ => 0,
165 }
166 };
167
168 let mut days: u64 = 0;
169 for y in 1970..year {
170 days += if is_leap(y) { 366 } else { 365 };
171 }
172 for m in 1..month {
173 days += days_in_month(year, m);
174 }
175 days += day.saturating_sub(1);
176
177 Some(days * 86400 + hour * 3600 + min * 60 + sec)
178}
179
180pub async fn start_eviction_loop(
197 store: Arc<SqliteStore>,
198 config: EvictionConfig,
199 policy: Arc<dyn EvictionPolicy + 'static>,
200 cancel: CancellationToken,
201) {
202 if config.max_entries == 0 {
203 tracing::debug!("eviction disabled (max_entries = 0)");
204 return;
205 }
206
207 let mut ticker = interval(Duration::from_secs(config.sweep_interval_secs));
208 ticker.tick().await;
210
211 loop {
212 tokio::select! {
213 () = cancel.cancelled() => {
214 tracing::debug!("eviction loop shutting down");
215 return;
216 }
217 _ = ticker.tick() => {}
218 }
219
220 tracing::debug!(max_entries = config.max_entries, "running eviction sweep");
221
222 match run_eviction_phase1(&store, &*policy, config.max_entries).await {
224 Ok(deleted) => {
225 if deleted > 0 {
226 tracing::info!(deleted, "eviction phase 1: soft-deleted entries");
227 }
228 }
229 Err(e) => {
230 tracing::warn!(error = %e, "eviction phase 1 failed, will retry next sweep");
231 }
232 }
233
234 match run_eviction_phase2(&store).await {
237 Ok(cleaned) => {
238 if cleaned > 0 {
239 tracing::info!(cleaned, "eviction phase 2: removed Qdrant vectors");
240 }
241 }
242 Err(e) => {
243 tracing::warn!(error = %e, "eviction phase 2 failed, will retry next sweep");
244 }
245 }
246 }
247}
248
249#[cfg_attr(
250 feature = "profiling",
251 tracing::instrument(name = "memory.eviction_phase1", skip_all)
252)]
253async fn run_eviction_phase1(
254 store: &SqliteStore,
255 policy: &dyn EvictionPolicy,
256 max_entries: usize,
257) -> Result<usize, MemoryError> {
258 let candidates = store.get_eviction_candidates().await?;
259 let total = candidates.len();
260
261 if total <= max_entries {
262 return Ok(0);
263 }
264
265 let excess = total - max_entries;
266 let mut scored: Vec<(f64, MessageId)> = candidates
267 .into_iter()
268 .map(|e| (policy.score(&e), e.id))
269 .collect();
270
271 scored.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap_or(std::cmp::Ordering::Equal));
273
274 let ids_to_delete: Vec<MessageId> = scored.into_iter().take(excess).map(|(_, id)| id).collect();
275 store.soft_delete_messages(&ids_to_delete).await?;
276
277 Ok(ids_to_delete.len())
278}
279
280#[cfg_attr(
281 feature = "profiling",
282 tracing::instrument(name = "memory.eviction_phase2", skip_all)
283)]
284async fn run_eviction_phase2(store: &SqliteStore) -> Result<usize, MemoryError> {
285 let ids = store.get_soft_deleted_message_ids().await?;
287 if ids.is_empty() {
288 return Ok(0);
289 }
290
291 tracing::warn!(
296 count = ids.len(),
297 "eviction phase 2: Qdrant vector removal not yet wired — marking cleaned without actual deletion (MVP)"
298 );
299
300 store.mark_qdrant_cleaned(&ids).await?;
303 Ok(ids.len())
304}
305
306#[cfg(test)]
309mod tests {
310 use super::*;
311
312 fn ts_ago(seconds_ago: u64) -> String {
316 let ts = unix_now_secs().saturating_sub(seconds_ago);
317 let sec = ts % 60;
319 let min = (ts / 60) % 60;
320 let hour = (ts / 3600) % 24;
321 let mut total_days = ts / 86400;
322 let is_leap =
323 |y: u64| (y.is_multiple_of(4) && !y.is_multiple_of(100)) || y.is_multiple_of(400);
324 let mut year = 1970u64;
325 loop {
326 let days_in_year = if is_leap(year) { 366 } else { 365 };
327 if total_days < days_in_year {
328 break;
329 }
330 total_days -= days_in_year;
331 year += 1;
332 }
333 let month_days = [
334 0u64,
335 31,
336 28 + u64::from(is_leap(year)),
337 31,
338 30,
339 31,
340 30,
341 31,
342 31,
343 30,
344 31,
345 30,
346 31,
347 ];
348 let mut month = 1u64;
349 while month <= 12 {
350 let month_idx = usize::try_from(month).unwrap_or_else(|_| unreachable!());
351 if total_days < month_days[month_idx] {
352 break;
353 }
354 total_days -= month_days[month_idx];
355 month += 1;
356 }
357 let day = total_days + 1;
358 format!("{year:04}-{month:02}-{day:02} {hour:02}:{min:02}:{sec:02}")
359 }
360
361 fn make_entry(access_count: u32, seconds_ago: u64) -> EvictionEntry {
362 let ts = ts_ago(seconds_ago);
363 EvictionEntry {
364 id: MessageId(1),
365 created_at: ts.clone(),
366 last_accessed: Some(ts),
367 access_count,
368 }
369 }
370
371 #[test]
372 fn ebbinghaus_recent_high_access_scores_near_one() {
373 let policy = EbbinghausPolicy::default();
374 let entry = make_entry(10, 1);
376 let score = policy.score(&entry);
377 assert!(
379 score > 0.99,
380 "score should be near 1.0 for recently accessed entry, got {score}"
381 );
382 }
383
384 #[test]
385 fn ebbinghaus_old_zero_access_scores_lower() {
386 let policy = EbbinghausPolicy::default();
387 let old = make_entry(0, 7 * 24 * 3600); let recent = make_entry(0, 60); assert!(
390 policy.score(&old) < policy.score(&recent),
391 "old entry must score lower than recent"
392 );
393 }
394
395 #[test]
396 fn ebbinghaus_high_access_decays_slower() {
397 let policy = EbbinghausPolicy::default();
398 let low = make_entry(1, 3600); let high = make_entry(20, 3600); assert!(
401 policy.score(&high) > policy.score(&low),
402 "high access count should yield higher score"
403 );
404 }
405
406 #[test]
407 fn ebbinghaus_never_accessed_uses_created_at_as_reference() {
408 let policy = EbbinghausPolicy::default();
409 let old_with_no_last_accessed = EvictionEntry {
412 id: MessageId(2),
413 created_at: ts_ago(7 * 24 * 3600),
414 last_accessed: None,
415 access_count: 0,
416 };
417 let old_with_same_last_accessed = make_entry(0, 7 * 24 * 3600);
418 let score_no_access = policy.score(&old_with_no_last_accessed);
419 let score_same = policy.score(&old_with_same_last_accessed);
420 let diff = (score_no_access - score_same).abs();
422 assert!(diff < 1e-6, "scores should match; diff = {diff}");
423 }
424
425 #[test]
426 fn eviction_config_default_is_disabled() {
427 let config = EvictionConfig::default();
428 assert_eq!(
429 config.max_entries, 0,
430 "eviction must be disabled by default"
431 );
432 }
433
434 #[test]
435 fn parse_sqlite_timestamp_known_value() {
436 let ts = parse_sqlite_timestamp_secs("2024-01-01 00:00:00").unwrap();
438 assert_eq!(
441 ts, 1_704_067_200,
442 "2024-01-01 must parse to known timestamp"
443 );
444 }
445}