use std::hash::Hash;
use std::hash::RandomState;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ByteWeightStats {
pub hits: u64,
pub misses: u64,
pub entry_count: u64,
}
pub struct ByteWeightCache<K, V, W>
where
K: Hash + Eq + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
W: Fn(&V) -> u64 + Send + Sync + 'static,
{
inner: moka::sync::Cache<K, V, RandomState>,
max_entry_bytes: Option<u64>,
weigher: Arc<W>,
hits: AtomicU64,
misses: AtomicU64,
}
impl<K, V, W> ByteWeightCache<K, V, W>
where
K: Hash + Eq + Send + Sync + 'static,
V: Clone + Send + Sync + 'static,
W: Fn(&V) -> u64 + Send + Sync + 'static,
{
pub fn new(max_capacity_bytes: u64, weigher: W) -> Self {
let weigher = Arc::new(weigher);
let moka_weigher = weigher.clone();
let inner = moka::sync::Cache::builder()
.max_capacity(max_capacity_bytes)
.weigher(move |_k, v: &V| (moka_weigher)(v).min(u32::MAX as u64) as u32)
.build();
Self {
inner,
max_entry_bytes: None,
weigher,
hits: AtomicU64::new(0),
misses: AtomicU64::new(0),
}
}
pub fn with_max_entry_bytes(mut self, max_entry_bytes: u64) -> Self {
self.max_entry_bytes = Some(max_entry_bytes);
self
}
pub fn get(&self, key: &K) -> Option<V> {
match self.inner.get(key) {
Some(v) => {
self.hits.fetch_add(1, Ordering::Relaxed);
Some(v)
}
None => {
self.misses.fetch_add(1, Ordering::Relaxed);
None
}
}
}
pub fn insert(&self, key: K, value: V) {
if let Some(max) = self.max_entry_bytes
&& (self.weigher)(&value) > max
{
return;
}
self.inner.insert(key, value);
}
pub fn get_or_compute<F, E>(&self, key: K, compute: F) -> Result<V, E>
where
F: FnOnce() -> Result<V, E>,
E: Clone + Send + Sync + 'static,
{
let ran = AtomicBool::new(false);
let outcome = self.inner.try_get_with(key, || {
ran.store(true, Ordering::Relaxed);
compute()
});
if ran.load(Ordering::Relaxed) {
self.misses.fetch_add(1, Ordering::Relaxed);
} else {
self.hits.fetch_add(1, Ordering::Relaxed);
}
match outcome {
Ok(v) => Ok(v),
Err(e) => Err(Arc::unwrap_or_clone(e)),
}
}
pub fn entry_count(&self) -> u64 {
self.inner.run_pending_tasks();
self.inner.entry_count()
}
pub fn stats(&self) -> ByteWeightStats {
self.inner.run_pending_tasks();
ByteWeightStats {
hits: self.hits.load(Ordering::Relaxed),
misses: self.misses.load(Ordering::Relaxed),
entry_count: self.inner.entry_count(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::sync::Barrier;
use std::thread;
use std::time::Duration;
fn small_cache() -> ByteWeightCache<String, String, impl Fn(&String) -> u64> {
ByteWeightCache::new(1024, |v: &String| v.len() as u64)
}
#[test]
fn test_get_hit_miss_and_clone_semantics() {
let cache = small_cache();
assert_eq!(cache.get(&"k".to_string()), None);
cache.insert("k".to_string(), "v".to_string());
assert_eq!(cache.get(&"k".to_string()), Some("v".to_string()));
let stats = cache.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.misses, 1);
}
#[test]
fn test_oversized_entry_skipped_by_admission() {
let cache =
ByteWeightCache::new(64 * 1024, |v: &String| v.len() as u64).with_max_entry_bytes(256);
let big = "x".repeat(512);
cache.insert("big".to_string(), big.clone());
assert_eq!(cache.get(&"big".to_string()), None);
let got = cache
.get_or_compute("big".to_string(), || Ok::<String, String>(big.clone()))
.unwrap();
assert_eq!(got.len(), 512);
}
#[test]
fn test_byte_budget_bounds_entries() {
let cache = ByteWeightCache::new(100, |v: &String| v.len() as u64);
for i in 0..6 {
cache.insert(format!("k{i}"), "x".repeat(64));
}
assert!(
cache.entry_count() <= 2,
"字节预算应限制条目数,实际 {}",
cache.entry_count()
);
}
#[test]
fn test_error_not_cached_and_propagates() {
let cache = small_cache();
let err = cache
.get_or_compute("k".to_string(), || Err("boom".to_string()))
.unwrap_err();
assert_eq!(err, "boom");
assert_eq!(cache.get(&"k".to_string()), None, "错误结果不写入缓存");
}
#[test]
fn test_single_flight_compute_exactly_once() {
let cache = Arc::new(small_cache());
let count = Arc::new(AtomicU64::new(0));
let barrier = Arc::new(Barrier::new(8));
let handles: Vec<_> = (0..8)
.map(|_| {
let cache = cache.clone();
let count = count.clone();
let barrier = barrier.clone();
thread::spawn(move || {
barrier.wait();
cache
.get_or_compute("hot".to_string(), || {
count.fetch_add(1, Ordering::SeqCst);
thread::sleep(Duration::from_millis(20));
Ok::<_, String>("v".to_string())
})
.unwrap()
})
})
.collect();
for h in handles {
assert_eq!(h.join().unwrap(), "v");
}
assert_eq!(
count.load(Ordering::SeqCst),
1,
"single-flight:compute 应恰执行 1 次"
);
let stats = cache.stats();
assert_eq!(stats.misses, 1);
assert_eq!(stats.hits, 7);
}
#[test]
fn test_stats_snapshot_fields() {
let cache = small_cache();
let _ = cache.get(&"miss".to_string());
cache.insert("hit".to_string(), "v".to_string());
let _ = cache.get(&"hit".to_string());
let stats = cache.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.misses, 1);
assert_eq!(stats.entry_count, 1);
}
}