use std::collections::HashMap;
use std::time::{Duration, Instant};
use crate::lru_cache::LruCache;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum MediaContentType {
VideoSegment {
bitrate: u32,
codec: String,
},
AudioSegment {
bitrate: u32,
},
Image {
width: u32,
height: u32,
},
Manifest,
Thumbnail,
Metadata,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct ContentCachePriority(pub u8);
impl ContentCachePriority {
pub fn for_type(content_type: &MediaContentType) -> Self {
let p = match content_type {
MediaContentType::Manifest => 10,
MediaContentType::Thumbnail => 8,
MediaContentType::VideoSegment { bitrate, .. } => {
if *bitrate >= 4_000_000 {
7
} else {
6
}
}
MediaContentType::AudioSegment { .. } => 5,
MediaContentType::Image { .. } => 4,
MediaContentType::Metadata => 3,
};
Self(p)
}
}
#[derive(Debug, Clone)]
pub struct CacheEntry {
pub key: String,
pub data: Vec<u8>,
pub content_type: MediaContentType,
pub inserted_at: Instant,
pub last_accessed: Instant,
pub access_count: u32,
pub size_bytes: usize,
}
impl CacheEntry {
fn new(key: String, data: Vec<u8>, content_type: MediaContentType) -> Self {
let size = data.len();
let now = Instant::now();
Self {
key,
data,
content_type,
inserted_at: now,
last_accessed: now,
access_count: 0,
size_bytes: size,
}
}
pub fn score_for_eviction(&self) -> f32 {
self.score_for_eviction_weighted(&ScoringWeights::default())
}
pub fn score_for_eviction_weighted(&self, w: &ScoringWeights) -> f32 {
let age_secs = self.last_accessed.elapsed().as_secs_f64();
let recency = (-age_secs / 60.0_f64).exp(); let base_priority = ContentCachePriority::for_type(&self.content_type).0 as f64;
let multiplier = w.priority_multiplier(&self.content_type);
let effective_priority = (base_priority * multiplier).max(0.001);
let size_factor = self.size_bytes as f64 / 1_048_576.0 + 1.0;
let score = (1.0 - recency).powf(w.recency_exp)
* (1.0 / effective_priority).powf(w.priority_exp)
* size_factor.powf(w.size_exp);
score as f32
}
}
pub fn ttl_for_type(content_type: &MediaContentType) -> Duration {
match content_type {
MediaContentType::Manifest => Duration::from_secs(30),
MediaContentType::VideoSegment { .. } => Duration::from_secs(300),
MediaContentType::AudioSegment { .. } => Duration::from_secs(300),
MediaContentType::Image { .. } => Duration::from_secs(3_600),
MediaContentType::Thumbnail => Duration::from_secs(86_400),
MediaContentType::Metadata => Duration::from_secs(600),
}
}
#[derive(Debug, Clone)]
pub struct ScoringWeights {
pub recency_exp: f64,
pub priority_exp: f64,
pub size_exp: f64,
pub per_type_priority: HashMap<String, f64>,
}
impl Default for ScoringWeights {
fn default() -> Self {
Self {
recency_exp: 1.0,
priority_exp: 1.0,
size_exp: 1.0,
per_type_priority: HashMap::new(),
}
}
}
impl ScoringWeights {
pub fn new() -> Self {
Self::default()
}
fn type_key(content_type: &MediaContentType) -> &'static str {
match content_type {
MediaContentType::VideoSegment { .. } => "VideoSegment",
MediaContentType::AudioSegment { .. } => "AudioSegment",
MediaContentType::Image { .. } => "Image",
MediaContentType::Manifest => "Manifest",
MediaContentType::Thumbnail => "Thumbnail",
MediaContentType::Metadata => "Metadata",
}
}
pub fn priority_multiplier(&self, content_type: &MediaContentType) -> f64 {
let key = Self::type_key(content_type);
self.per_type_priority.get(key).copied().unwrap_or(1.0)
}
}
pub struct ContentAwareCache {
inner: LruCache<String, CacheEntry>,
capacity: usize,
total_bytes: usize,
max_bytes: Option<usize>,
scoring_weights: ScoringWeights,
}
impl ContentAwareCache {
pub fn new(capacity: usize) -> Self {
Self {
inner: LruCache::new(capacity),
capacity,
total_bytes: 0,
max_bytes: None,
scoring_weights: ScoringWeights::default(),
}
}
pub fn with_max_bytes(mut self, max_bytes: usize) -> Self {
self.max_bytes = Some(max_bytes);
self
}
pub fn with_scoring_weights(mut self, weights: ScoringWeights) -> Self {
self.scoring_weights = weights;
self
}
pub fn set_scoring_weights(&mut self, weights: ScoringWeights) {
self.scoring_weights = weights;
}
pub fn scoring_weights(&self) -> &ScoringWeights {
&self.scoring_weights
}
pub fn insert_media(&mut self, key: String, data: Vec<u8>, content_type: MediaContentType) {
let size = data.len();
if let Some(old) = self.inner.remove(&key) {
self.total_bytes = self.total_bytes.saturating_sub(old.size_bytes);
}
if let Some(max_bytes) = self.max_bytes {
while self.total_bytes + size > max_bytes && !self.inner.is_empty() {
self.evict_worst();
}
}
if self.inner.len() >= self.capacity {
self.evict_worst();
}
let entry = CacheEntry::new(key.clone(), data, content_type);
self.total_bytes += size;
self.inner.insert(key, entry, size);
}
pub fn get(&mut self, key: &str) -> Option<&CacheEntry> {
let key_owned = key.to_string();
if self.inner.contains(&key_owned) {
let updated_entry = {
let entry = self.inner.peek(&key_owned)?;
let mut e = entry.clone();
e.last_accessed = Instant::now();
e.access_count = e.access_count.saturating_add(1);
e
};
let size = updated_entry.size_bytes;
self.total_bytes = self.total_bytes.saturating_sub(size);
self.inner.insert(key_owned.clone(), updated_entry, size);
self.total_bytes += size;
self.inner.peek(&key_owned)
} else {
None
}
}
pub fn peek(&self, key: &str) -> Option<&CacheEntry> {
self.inner.peek(&key.to_string())
}
pub fn remove(&mut self, key: &str) -> bool {
if let Some(entry) = self.inner.remove(&key.to_string()) {
self.total_bytes = self.total_bytes.saturating_sub(entry.size_bytes);
true
} else {
false
}
}
pub fn evict_expired(&mut self) -> usize {
let mut expired_keys: Vec<String> = Vec::new();
let mut remaining: Vec<(String, CacheEntry)> = Vec::new();
while let Some((k, entry)) = self.inner.evict_lru() {
let ttl = ttl_for_type(&entry.content_type);
if entry.inserted_at.elapsed() > ttl {
expired_keys.push(k);
} else {
remaining.push((k, entry));
}
}
for (k, entry) in remaining {
let size = entry.size_bytes;
self.inner.insert(k, entry, size);
}
self.total_bytes = 0;
self.total_bytes = self.inner.stats().total_size_bytes;
expired_keys.len()
}
pub fn len(&self) -> usize {
self.inner.len()
}
pub fn is_empty(&self) -> bool {
self.inner.is_empty()
}
pub fn total_bytes(&self) -> usize {
self.total_bytes
}
pub fn capacity(&self) -> usize {
self.capacity
}
fn evict_worst(&mut self) {
if self.inner.is_empty() {
return;
}
let weights = self.scoring_weights.clone();
let mut entries: Vec<(String, CacheEntry)> = Vec::with_capacity(self.inner.len());
while let Some((k, entry)) = self.inner.evict_lru() {
entries.push((k, entry));
}
let worst_idx = entries
.iter()
.enumerate()
.max_by(|(_, (_, a)), (_, (_, b))| {
a.score_for_eviction_weighted(&weights)
.partial_cmp(&b.score_for_eviction_weighted(&weights))
.unwrap_or(std::cmp::Ordering::Equal)
})
.map(|(i, _)| i)
.unwrap_or(0);
let (_, evicted) = entries.remove(worst_idx);
self.total_bytes = self.total_bytes.saturating_sub(evicted.size_bytes);
for (k, entry) in entries {
let size = entry.size_bytes;
self.inner.insert(k, entry, size);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_priority_manifest_is_highest() {
let p = ContentCachePriority::for_type(&MediaContentType::Manifest);
assert_eq!(p.0, 10);
}
#[test]
fn test_priority_thumbnail() {
let p = ContentCachePriority::for_type(&MediaContentType::Thumbnail);
assert_eq!(p.0, 8);
}
#[test]
fn test_priority_high_bitrate_video() {
let p = ContentCachePriority::for_type(&MediaContentType::VideoSegment {
bitrate: 5_000_000,
codec: "av1".into(),
});
assert_eq!(p.0, 7);
}
#[test]
fn test_priority_low_bitrate_video() {
let p = ContentCachePriority::for_type(&MediaContentType::VideoSegment {
bitrate: 1_000_000,
codec: "vp9".into(),
});
assert_eq!(p.0, 6);
}
#[test]
fn test_priority_audio() {
let p =
ContentCachePriority::for_type(&MediaContentType::AudioSegment { bitrate: 128_000 });
assert_eq!(p.0, 5);
}
#[test]
fn test_priority_image() {
let p = ContentCachePriority::for_type(&MediaContentType::Image {
width: 1920,
height: 1080,
});
assert_eq!(p.0, 4);
}
#[test]
fn test_priority_metadata() {
let p = ContentCachePriority::for_type(&MediaContentType::Metadata);
assert_eq!(p.0, 3);
}
#[test]
fn test_ttl_manifest() {
assert_eq!(
ttl_for_type(&MediaContentType::Manifest),
Duration::from_secs(30)
);
}
#[test]
fn test_ttl_video_segment() {
let ct = MediaContentType::VideoSegment {
bitrate: 2_000_000,
codec: "av1".into(),
};
assert_eq!(ttl_for_type(&ct), Duration::from_secs(300));
}
#[test]
fn test_ttl_thumbnail() {
assert_eq!(
ttl_for_type(&MediaContentType::Thumbnail),
Duration::from_secs(86_400)
);
}
#[test]
fn test_ttl_image() {
let ct = MediaContentType::Image {
width: 100,
height: 100,
};
assert_eq!(ttl_for_type(&ct), Duration::from_secs(3_600));
}
#[test]
fn test_insert_and_get() {
let mut cache = ContentAwareCache::new(16);
cache.insert_media(
"seg1".into(),
vec![0u8; 1024],
MediaContentType::VideoSegment {
bitrate: 2_000_000,
codec: "av1".into(),
},
);
let entry = cache.get("seg1");
assert!(entry.is_some());
assert_eq!(entry.map(|e| e.size_bytes), Some(1024));
}
#[test]
fn test_get_absent_returns_none() {
let mut cache = ContentAwareCache::new(8);
assert!(cache.get("missing").is_none());
}
#[test]
fn test_len_and_is_empty() {
let mut cache = ContentAwareCache::new(8);
assert!(cache.is_empty());
cache.insert_media("m".into(), vec![1, 2], MediaContentType::Manifest);
assert_eq!(cache.len(), 1);
assert!(!cache.is_empty());
}
#[test]
fn test_remove() {
let mut cache = ContentAwareCache::new(8);
cache.insert_media("key".into(), vec![0u8; 512], MediaContentType::Metadata);
assert!(cache.remove("key"));
assert!(cache.get("key").is_none());
}
#[test]
fn test_remove_absent() {
let mut cache = ContentAwareCache::new(8);
assert!(!cache.remove("ghost"));
}
#[test]
fn test_total_bytes_tracking() {
let mut cache = ContentAwareCache::new(16);
cache.insert_media("a".into(), vec![0u8; 100], MediaContentType::Manifest);
cache.insert_media("b".into(), vec![0u8; 200], MediaContentType::Metadata);
assert_eq!(cache.total_bytes(), 300);
cache.remove("a");
assert_eq!(cache.total_bytes(), 200);
}
#[test]
fn test_capacity_reported() {
let cache = ContentAwareCache::new(32);
assert_eq!(cache.capacity(), 32);
}
#[test]
fn test_score_for_eviction_just_inserted_is_low() {
let entry = CacheEntry::new("k".into(), vec![0u8; 100], MediaContentType::Manifest);
let score = entry.score_for_eviction();
assert!(
score < 0.1,
"fresh entry should have low eviction score, got {score}"
);
}
#[test]
fn test_score_low_priority_higher_than_high_priority() {
let manifest_entry =
CacheEntry::new("m".into(), vec![0u8; 100], MediaContentType::Manifest);
let meta_entry = CacheEntry::new("d".into(), vec![0u8; 100], MediaContentType::Metadata);
let p_manifest = ContentCachePriority::for_type(&MediaContentType::Manifest).0 as f32;
let p_meta = ContentCachePriority::for_type(&MediaContentType::Metadata).0 as f32;
assert!(
1.0 / p_meta > 1.0 / p_manifest,
"metadata entry should evict before manifest"
);
drop(manifest_entry);
drop(meta_entry);
}
#[test]
fn test_eviction_prefers_low_priority_entries() {
let mut cache = ContentAwareCache::new(2);
cache.insert_media("manifest".into(), vec![0u8; 1], MediaContentType::Manifest);
cache.insert_media("meta".into(), vec![0u8; 1], MediaContentType::Metadata);
std::thread::sleep(std::time::Duration::from_millis(100));
let _ = cache.get("manifest");
cache.insert_media(
"new".into(),
vec![0u8; 1],
MediaContentType::VideoSegment {
bitrate: 2_000_000,
codec: "av1".into(),
},
);
assert_eq!(cache.len(), 2);
assert!(
cache.peek("manifest").is_some(),
"manifest should survive eviction"
);
}
#[test]
fn test_access_count_increments_on_get() {
let mut cache = ContentAwareCache::new(8);
cache.insert_media("k".into(), vec![1, 2, 3], MediaContentType::Thumbnail);
cache.get("k");
cache.get("k");
let count = cache.peek("k").map(|e| e.access_count).unwrap_or(0);
assert_eq!(count, 2, "access_count should be 2 after two gets");
}
#[test]
fn test_max_bytes_triggers_eviction() {
let mut cache = ContentAwareCache::new(100).with_max_bytes(500);
for i in 0..5u32 {
cache.insert_media(
format!("seg_{i}"),
vec![0u8; 100],
MediaContentType::AudioSegment { bitrate: 128_000 },
);
}
assert!(cache.total_bytes() <= 500);
cache.insert_media("extra".into(), vec![0u8; 100], MediaContentType::Metadata);
assert!(
cache.total_bytes() <= 500,
"total bytes exceeded budget: {}",
cache.total_bytes()
);
}
#[test]
fn test_peek_does_not_change_access_count() {
let mut cache = ContentAwareCache::new(8);
cache.insert_media("p".into(), vec![99], MediaContentType::Manifest);
let before = cache.peek("p").map(|e| e.access_count).unwrap_or(99);
let _ = cache.peek("p");
let after = cache.peek("p").map(|e| e.access_count).unwrap_or(99);
assert_eq!(before, after, "peek must not change access_count");
}
#[test]
fn test_insert_same_key_updates_value() {
let mut cache = ContentAwareCache::new(8);
cache.insert_media("k".into(), vec![1, 2, 3], MediaContentType::Manifest);
cache.insert_media("k".into(), vec![10, 20], MediaContentType::Manifest);
assert_eq!(cache.len(), 1, "duplicate key should not increase len");
assert_eq!(
cache.total_bytes(),
2,
"total_bytes should reflect updated size"
);
}
#[test]
fn test_evict_expired_no_entries() {
let mut cache = ContentAwareCache::new(8);
assert_eq!(cache.evict_expired(), 0);
}
#[test]
fn test_evict_expired_fresh_entries_survive() {
let mut cache = ContentAwareCache::new(8);
cache.insert_media("fresh".into(), vec![0u8; 10], MediaContentType::Manifest);
let evicted = cache.evict_expired();
assert_eq!(evicted, 0);
assert!(cache.peek("fresh").is_some());
}
}