use crate::expiration_tracker::ExpirationTracker;
use std::collections::{BTreeMap, HashMap};
#[derive(Debug, Default, Clone)]
pub struct ExpirationProfile {
buckets: BTreeMap<u32, u64>,
total_bytes: u64,
}
impl ExpirationProfile {
pub fn new() -> Self {
Self::default()
}
pub fn add(&mut self, expiration_time: u32, size: u64) {
if expiration_time != 0 {
*self.buckets.entry(expiration_time).or_default() += size;
}
self.total_bytes += size;
}
pub fn is_all_expired(&self, current_time: u32) -> bool {
if self.total_bytes == 0 {
return true;
}
self.buckets.keys().all(|&t| t <= current_time)
}
pub fn expired_fraction(&self, current_time: u32) -> f64 {
if self.total_bytes == 0 {
return 1.0;
}
let expired: u64 = self
.buckets
.iter()
.filter(|&(&t, _)| t <= current_time)
.map(|(_, &b)| b)
.sum();
expired as f64 / self.total_bytes as f64
}
pub fn total_bytes(&self) -> u64 {
self.total_bytes
}
}
#[derive(Debug, Default)]
pub struct ExpirationProfileStore {
trackers: HashMap<u32, ExpirationTracker>,
}
impl ExpirationProfileStore {
pub fn new() -> Self {
Self::default()
}
pub fn put_file(&mut self, tracker: ExpirationTracker) {
let file_number = tracker.get_file_number();
self.trackers.insert(file_number, tracker);
}
pub fn remove_file(&mut self, file_number: u32) {
self.trackers.remove(&file_number);
}
pub fn get_expired_bytes(
&self,
file_number: u32,
current_time_hours: u64,
) -> i64 {
self.trackers
.get(&file_number)
.map(|t| t.get_expired_bytes(current_time_hours))
.unwrap_or(0)
}
pub fn has_file(&self, file_number: u32) -> bool {
self.trackers.contains_key(&file_number)
}
pub fn file_count(&self) -> usize {
self.trackers.len()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_is_empty() {
let ep = ExpirationProfile::new();
assert_eq!(ep.total_bytes(), 0);
assert!(ep.is_all_expired(0));
assert_eq!(ep.expired_fraction(0), 1.0);
}
#[test]
fn test_add_no_expiration() {
let mut ep = ExpirationProfile::new();
ep.add(0, 100);
assert_eq!(ep.total_bytes(), 100);
assert!(ep.is_all_expired(0));
}
#[test]
fn test_add_with_expiration() {
let mut ep = ExpirationProfile::new();
ep.add(100, 500);
ep.add(200, 300);
assert_eq!(ep.total_bytes(), 800);
assert!(!ep.is_all_expired(99));
assert_eq!(ep.expired_fraction(99), 0.0);
assert!(!ep.is_all_expired(100));
assert!((ep.expired_fraction(100) - 500.0 / 800.0).abs() < 1e-9);
assert!(ep.is_all_expired(200));
assert_eq!(ep.expired_fraction(200), 1.0);
}
#[test]
fn test_expired_fraction_mixed_never_and_ttl() {
let mut ep = ExpirationProfile::new();
ep.add(0, 400); ep.add(10, 600); assert_eq!(ep.total_bytes(), 1000);
let frac = ep.expired_fraction(10);
assert!((frac - 0.6).abs() < 1e-9);
assert!(ep.is_all_expired(10));
}
}
#[cfg(test)]
mod store_tests {
use super::*;
#[test]
fn test_cln9_put_and_get_expired_bytes() {
let mut store = ExpirationProfileStore::new();
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 500); tracker.track(200, 300);
store.put_file(tracker);
assert!(store.has_file(1));
assert_eq!(store.get_expired_bytes(1, 100), 500);
assert_eq!(store.get_expired_bytes(1, 200), 800);
}
#[test]
fn test_cln9_remove_file() {
let mut store = ExpirationProfileStore::new();
let tracker = ExpirationTracker::new(42);
store.put_file(tracker);
assert!(store.has_file(42));
store.remove_file(42);
assert!(!store.has_file(42));
assert_eq!(store.get_expired_bytes(42, 1000), 0);
}
#[test]
fn test_cln9_missing_file_returns_zero() {
let store = ExpirationProfileStore::new();
assert_eq!(store.get_expired_bytes(99, 1000), 0);
}
#[test]
fn test_cln9_replaces_on_put() {
let mut store = ExpirationProfileStore::new();
let mut t1 = ExpirationTracker::new(5);
t1.track(100, 1000);
store.put_file(t1);
let mut t2 = ExpirationTracker::new(5);
t2.track(100, 2000); store.put_file(t2);
assert_eq!(store.get_expired_bytes(5, 100), 2000);
}
}