use hashbrown::HashMap;
#[derive(Debug)]
pub struct ExpirationTracker {
file_number: u32,
bins: HashMap<u64, ExpirationBin>,
}
#[derive(Debug, Clone, Default)]
pub struct ExpirationBin {
pub expiration_time: u64,
pub count: i32,
pub size: i32,
}
impl ExpirationTracker {
pub fn new(file_number: u32) -> Self {
Self { file_number, bins: HashMap::new() }
}
pub fn get_file_number(&self) -> u32 {
self.file_number
}
pub fn track(&mut self, expiration_time: u64, size: i32) {
if expiration_time == 0 {
return;
}
self.bins
.entry(expiration_time)
.and_modify(|bin| {
bin.count += 1;
bin.size += size;
})
.or_insert(ExpirationBin { expiration_time, count: 1, size });
}
pub fn get_expired_bytes(&self, current_time: u64) -> i64 {
let mut expired_size = 0i64;
for bin in self.bins.values() {
if bin.expiration_time > 0 && bin.expiration_time <= current_time {
expired_size += bin.size as i64;
}
}
expired_size
}
pub fn get_total_tracked_size(&self) -> i64 {
self.bins.values().map(|bin| bin.size as i64).sum()
}
pub fn get_bin_count(&self) -> usize {
self.bins.len()
}
pub fn get_bins(&self) -> &HashMap<u64, ExpirationBin> {
&self.bins
}
pub fn clear(&mut self) {
self.bins.clear();
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_new_tracker() {
let tracker = ExpirationTracker::new(5);
assert_eq!(tracker.get_file_number(), 5);
assert_eq!(tracker.get_bin_count(), 0);
assert_eq!(tracker.get_total_tracked_size(), 0);
}
#[test]
fn test_track_single_entry() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1024);
assert_eq!(tracker.get_bin_count(), 1);
assert_eq!(tracker.get_total_tracked_size(), 1024);
}
#[test]
fn test_track_never_expires() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(0, 1024);
assert_eq!(tracker.get_bin_count(), 0);
assert_eq!(tracker.get_total_tracked_size(), 0);
}
#[test]
fn test_track_multiple_same_expiration() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 512);
tracker.track(100, 256);
tracker.track(100, 128);
assert_eq!(tracker.get_bin_count(), 1);
assert_eq!(tracker.get_total_tracked_size(), 896);
let bin = tracker.get_bins().get(&100).unwrap();
assert_eq!(bin.count, 3);
assert_eq!(bin.size, 896);
}
#[test]
fn test_track_different_expirations() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
tracker.track(300, 3000);
assert_eq!(tracker.get_bin_count(), 3);
assert_eq!(tracker.get_total_tracked_size(), 6000);
}
#[test]
fn test_get_expired_bytes_none_expired() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
tracker.track(300, 3000);
let expired = tracker.get_expired_bytes(50);
assert_eq!(expired, 0);
}
#[test]
fn test_get_expired_bytes_some_expired() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
tracker.track(300, 3000);
let expired = tracker.get_expired_bytes(250);
assert_eq!(expired, 3000); }
#[test]
fn test_get_expired_bytes_all_expired() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
tracker.track(300, 3000);
let expired = tracker.get_expired_bytes(400);
assert_eq!(expired, 6000);
}
#[test]
fn test_get_expired_bytes_exact_boundary() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
let expired = tracker.get_expired_bytes(100);
assert_eq!(expired, 1000);
let expired = tracker.get_expired_bytes(200);
assert_eq!(expired, 3000);
}
#[test]
fn test_clear() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
assert_eq!(tracker.get_bin_count(), 2);
tracker.clear();
assert_eq!(tracker.get_bin_count(), 0);
assert_eq!(tracker.get_total_tracked_size(), 0);
assert_eq!(tracker.get_expired_bytes(1000), 0);
}
#[test]
fn test_large_values() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(1_000_000, 100_000_000);
tracker.track(2_000_000, 200_000_000);
assert_eq!(tracker.get_total_tracked_size(), 300_000_000);
assert_eq!(tracker.get_expired_bytes(1_500_000), 100_000_000);
assert_eq!(tracker.get_expired_bytes(3_000_000), 300_000_000);
}
#[test]
fn test_bins_independent() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(100, 1000);
tracker.track(200, 2000);
tracker.track(300, 3000);
let bin100 = tracker.get_bins().get(&100).unwrap();
let bin200 = tracker.get_bins().get(&200).unwrap();
let bin300 = tracker.get_bins().get(&300).unwrap();
assert_eq!(bin100.size, 1000);
assert_eq!(bin200.size, 2000);
assert_eq!(bin300.size, 3000);
assert_eq!(bin100.count, 1);
assert_eq!(bin200.count, 1);
assert_eq!(bin300.count, 1);
}
#[test]
fn test_expiration_bin_default() {
let bin = ExpirationBin::default();
assert_eq!(bin.expiration_time, 0);
assert_eq!(bin.count, 0);
assert_eq!(bin.size, 0);
}
#[test]
fn test_mixed_tracking() {
let mut tracker = ExpirationTracker::new(1);
tracker.track(0, 1000); tracker.track(100, 500);
tracker.track(0, 2000); tracker.track(100, 500);
tracker.track(200, 1000);
assert_eq!(tracker.get_bin_count(), 2); assert_eq!(tracker.get_total_tracked_size(), 2000); }
}