use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Mutex, OnceLock};
use std::time::{SystemTime, UNIX_EPOCH};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::cold_prefix;
const MAX_TRACKED: usize = 4096;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CacheOutcome {
ColdStart,
WarmReuse,
TtlLapse,
PrefixChange,
}
#[derive(Debug, Clone, Copy)]
struct PrefixState {
prefix_hash: u64,
last_touch: u64,
}
static COLD_STARTS: AtomicU64 = AtomicU64::new(0);
static WARM_REUSES: AtomicU64 = AtomicU64::new(0);
static TTL_LAPSES: AtomicU64 = AtomicU64::new(0);
static PREFIX_CHANGES: AtomicU64 = AtomicU64::new(0);
fn store() -> &'static Mutex<HashMap<u64, PrefixState>> {
static STORE: OnceLock<Mutex<HashMap<u64, PrefixState>>> = OnceLock::new();
STORE.get_or_init(|| Mutex::new(HashMap::new()))
}
fn now_secs() -> u64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map_or(0, |d| d.as_secs())
}
#[must_use]
pub fn classify(prev: Option<(u64, u64)>, curr_hash: u64, now: u64, ttl_secs: u64) -> CacheOutcome {
match prev {
None => CacheOutcome::ColdStart,
Some((prev_hash, last_touch)) => {
if prev_hash != curr_hash {
CacheOutcome::PrefixChange
} else if now.saturating_sub(last_touch) > ttl_secs {
CacheOutcome::TtlLapse
} else {
CacheOutcome::WarmReuse
}
}
}
}
fn bump(outcome: CacheOutcome) {
let counter = match outcome {
CacheOutcome::ColdStart => &COLD_STARTS,
CacheOutcome::WarmReuse => &WARM_REUSES,
CacheOutcome::TtlLapse => &TTL_LAPSES,
CacheOutcome::PrefixChange => &PREFIX_CHANGES,
};
counter.fetch_add(1, Ordering::Relaxed);
}
fn evict_oldest(map: &mut HashMap<u64, PrefixState>) {
if let Some(oldest) = map
.iter()
.min_by_key(|(_, s)| s.last_touch)
.map(|(k, _)| *k)
{
map.remove(&oldest);
}
}
pub fn record_request(messages: &[Value], cached: usize) -> Option<CacheOutcome> {
let conv_key = cold_prefix::conversation_key(messages)?;
let curr_hash = cold_prefix::cached_prefix_hash(messages, cached)?;
let ttl = cold_prefix::resolved_ttl_secs(messages, cached).unwrap_or(0);
let now = now_secs();
let outcome = {
let mut map = store()
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let prev = map.get(&conv_key).map(|s| (s.prefix_hash, s.last_touch));
let outcome = classify(prev, curr_hash, now, ttl);
map.insert(
conv_key,
PrefixState {
prefix_hash: curr_hash,
last_touch: now,
},
);
if map.len() > MAX_TRACKED {
evict_oldest(&mut map);
}
outcome
};
bump(outcome);
Some(outcome)
}
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq)]
pub struct CacheAttribution {
pub cold_starts: u64,
pub warm_reuses: u64,
pub ttl_lapses: u64,
pub prefix_changes: u64,
}
#[must_use]
pub fn snapshot() -> CacheAttribution {
CacheAttribution {
cold_starts: COLD_STARTS.load(Ordering::Relaxed),
warm_reuses: WARM_REUSES.load(Ordering::Relaxed),
ttl_lapses: TTL_LAPSES.load(Ordering::Relaxed),
prefix_changes: PREFIX_CHANGES.load(Ordering::Relaxed),
}
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn classify_distinguishes_ttl_lapse_from_prefix_change() {
assert_eq!(classify(None, 1, 100, 300), CacheOutcome::ColdStart);
assert_eq!(
classify(Some((1, 100)), 1, 350, 300),
CacheOutcome::WarmReuse
);
assert_eq!(
classify(Some((1, 100)), 1, 500, 300),
CacheOutcome::TtlLapse
);
assert_eq!(
classify(Some((1, 100)), 2, 110, 300),
CacheOutcome::PrefixChange
);
assert_eq!(
classify(Some((1, 100)), 2, 9999, 300),
CacheOutcome::PrefixChange
);
}
fn anchored(first_text: &str) -> Vec<Value> {
vec![
json!({"role": "user", "content": [
{"type": "text", "text": first_text, "cache_control": {"type": "ephemeral"}}
]}),
json!({"role": "assistant", "content": "ok"}),
]
}
#[test]
fn unanchored_turn_is_not_attributed() {
let msgs = anchored("unanchored-attribution-test");
assert_eq!(record_request(&msgs, 0), None);
}
#[test]
fn first_anchored_turn_is_cold_start_then_warm() {
let msgs = anchored("cold-then-warm-attribution-test");
assert_eq!(record_request(&msgs, 1), Some(CacheOutcome::ColdStart));
assert_eq!(record_request(&msgs, 1), Some(CacheOutcome::WarmReuse));
}
#[test]
fn prefix_change_detected_with_stable_head() {
let head = json!({"role": "user", "content": [
{"type": "text", "text": "stable-head-attribution", "cache_control": {"type": "ephemeral"}}
]});
let v1 = vec![
head.clone(),
json!({"role": "assistant", "content": "answer one"}),
];
let v2 = vec![
head,
json!({"role": "assistant", "content": "answer two CHANGED"}),
];
assert_eq!(record_request(&v1, 2), Some(CacheOutcome::ColdStart));
assert_eq!(record_request(&v2, 2), Some(CacheOutcome::PrefixChange));
}
}