use cid::Cid;
use libp2p::PeerId;
use parking_lot::RwLock;
use serde::Serialize;
use std::collections::{HashMap, HashSet};
use std::time::{Duration, Instant};
use tracing::{debug, info};
const DEFAULT_TTL: Duration = Duration::from_secs(3600);
const DEFAULT_MAX_ENTRIES: usize = 10000;
#[derive(Debug, Clone)]
pub struct ProviderCacheConfig {
pub ttl: Duration,
pub max_entries: usize,
pub min_providers: usize,
}
impl Default for ProviderCacheConfig {
fn default() -> Self {
Self {
ttl: DEFAULT_TTL,
max_entries: DEFAULT_MAX_ENTRIES,
min_providers: 1,
}
}
}
#[derive(Debug, Clone)]
struct CachedProviders {
providers: HashSet<PeerId>,
cached_at: Instant,
last_accessed: Instant,
access_count: u64,
}
impl CachedProviders {
fn new(providers: HashSet<PeerId>) -> Self {
let now = Instant::now();
Self {
providers,
cached_at: now,
last_accessed: now,
access_count: 0,
}
}
fn is_expired(&self, ttl: Duration) -> bool {
self.cached_at.elapsed() > ttl
}
fn touch(&mut self) {
self.last_accessed = Instant::now();
self.access_count += 1;
}
}
pub struct ProviderCache {
config: ProviderCacheConfig,
cache: RwLock<HashMap<Cid, CachedProviders>>,
stats: RwLock<CacheStats>,
}
#[derive(Default)]
struct CacheStats {
hits: u64,
misses: u64,
evictions: u64,
expirations: u64,
}
impl ProviderCache {
pub fn new() -> Self {
Self::with_config(ProviderCacheConfig::default())
}
pub fn with_config(config: ProviderCacheConfig) -> Self {
Self {
config,
cache: RwLock::new(HashMap::new()),
stats: RwLock::new(CacheStats::default()),
}
}
pub fn get(&self, cid: &Cid) -> Option<Vec<PeerId>> {
let mut cache = self.cache.write();
let mut stats = self.stats.write();
if let Some(entry) = cache.get_mut(cid) {
if entry.is_expired(self.config.ttl) {
cache.remove(cid);
stats.expirations += 1;
stats.misses += 1;
debug!("Provider cache expired for {}", cid);
return None;
}
entry.touch();
stats.hits += 1;
debug!(
"Provider cache hit for {} ({} providers)",
cid,
entry.providers.len()
);
return Some(entry.providers.iter().cloned().collect());
}
stats.misses += 1;
None
}
pub fn has_providers(&self, cid: &Cid) -> bool {
let cache = self.cache.read();
if let Some(entry) = cache.get(cid) {
!entry.is_expired(self.config.ttl) && !entry.providers.is_empty()
} else {
false
}
}
pub fn needs_refresh(&self, cid: &Cid) -> bool {
let cache = self.cache.read();
if let Some(entry) = cache.get(cid) {
entry.is_expired(self.config.ttl) || entry.providers.len() < self.config.min_providers
} else {
true
}
}
pub fn put(&self, cid: Cid, providers: Vec<PeerId>) {
let provider_set: HashSet<PeerId> = providers.into_iter().collect();
if provider_set.is_empty() {
debug!("Not caching empty provider list for {}", cid);
return;
}
let mut cache = self.cache.write();
if cache.len() >= self.config.max_entries {
self.evict_lru(&mut cache);
}
let count = provider_set.len();
cache.insert(cid, CachedProviders::new(provider_set));
info!("Cached {} providers for {}", count, cid);
}
pub fn add_provider(&self, cid: &Cid, provider: PeerId) {
let mut cache = self.cache.write();
if let Some(entry) = cache.get_mut(cid) {
if !entry.is_expired(self.config.ttl) {
entry.providers.insert(provider);
entry.touch();
debug!("Added provider {} to cache for {}", provider, cid);
}
}
}
pub fn remove_provider(&self, cid: &Cid, provider: &PeerId) {
let mut cache = self.cache.write();
if let Some(entry) = cache.get_mut(cid) {
entry.providers.remove(provider);
debug!("Removed provider {} from cache for {}", provider, cid);
}
}
pub fn invalidate(&self, cid: &Cid) {
let mut cache = self.cache.write();
cache.remove(cid);
debug!("Invalidated cache for {}", cid);
}
pub fn cleanup_expired(&self) {
let mut cache = self.cache.write();
let mut stats = self.stats.write();
let ttl = self.config.ttl;
let before = cache.len();
cache.retain(|_, entry| !entry.is_expired(ttl));
let removed = before - cache.len();
if removed > 0 {
stats.expirations += removed as u64;
info!("Cleaned up {} expired provider cache entries", removed);
}
}
pub fn clear(&self) {
let mut cache = self.cache.write();
cache.clear();
info!("Provider cache cleared");
}
pub fn stats(&self) -> ProviderCacheStats {
let cache = self.cache.read();
let stats = self.stats.read();
let total_providers: usize = cache.values().map(|e| e.providers.len()).sum();
let hit_rate = if stats.hits + stats.misses > 0 {
stats.hits as f64 / (stats.hits + stats.misses) as f64
} else {
0.0
};
ProviderCacheStats {
entries: cache.len(),
max_entries: self.config.max_entries,
total_providers,
hits: stats.hits,
misses: stats.misses,
hit_rate,
evictions: stats.evictions,
expirations: stats.expirations,
}
}
pub fn len(&self) -> usize {
self.cache.read().len()
}
pub fn is_empty(&self) -> bool {
self.cache.read().is_empty()
}
fn evict_lru(&self, cache: &mut HashMap<Cid, CachedProviders>) {
let to_evict = (self.config.max_entries / 10).max(1);
let mut entries: Vec<_> = cache
.iter()
.map(|(cid, entry)| (*cid, entry.last_accessed))
.collect();
entries.sort_by_key(|a| a.1);
let mut stats = self.stats.write();
for (cid, _) in entries.into_iter().take(to_evict) {
cache.remove(&cid);
stats.evictions += 1;
}
debug!("Evicted {} LRU cache entries", to_evict);
}
}
impl Default for ProviderCache {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize)]
pub struct ProviderCacheStats {
pub entries: usize,
pub max_entries: usize,
pub total_providers: usize,
pub hits: u64,
pub misses: u64,
pub hit_rate: f64,
pub evictions: u64,
pub expirations: u64,
}
#[cfg(test)]
mod tests {
use super::*;
use multihash_codetable::{Code, MultihashDigest};
fn make_cid(data: &[u8]) -> Cid {
let hash = Code::Sha2_256.digest(data);
Cid::new_v1(0x55, hash)
}
fn random_peer_id() -> PeerId {
PeerId::random()
}
#[test]
fn test_provider_cache_basic() {
let cache = ProviderCache::new();
let cid = make_cid(b"test data");
let peer1 = random_peer_id();
let peer2 = random_peer_id();
assert!(cache.get(&cid).is_none());
assert!(cache.needs_refresh(&cid));
cache.put(cid, vec![peer1, peer2]);
let providers = cache
.get(&cid)
.expect("test: cache should contain providers after put");
assert_eq!(providers.len(), 2);
assert!(providers.contains(&peer1));
assert!(providers.contains(&peer2));
assert!(cache.has_providers(&cid));
assert!(!cache.needs_refresh(&cid));
}
#[test]
fn test_provider_cache_add_remove() {
let cache = ProviderCache::new();
let cid = make_cid(b"test");
let peer1 = random_peer_id();
let peer2 = random_peer_id();
let peer3 = random_peer_id();
cache.put(cid, vec![peer1, peer2]);
cache.add_provider(&cid, peer3);
let providers = cache
.get(&cid)
.expect("test: cache should have 3 providers after add");
assert_eq!(providers.len(), 3);
cache.remove_provider(&cid, &peer1);
let providers = cache
.get(&cid)
.expect("test: cache should have 2 providers after remove");
assert_eq!(providers.len(), 2);
assert!(!providers.contains(&peer1));
}
#[test]
fn test_provider_cache_expiration() {
let config = ProviderCacheConfig {
ttl: Duration::from_millis(50),
..Default::default()
};
let cache = ProviderCache::with_config(config);
let cid = make_cid(b"expiring");
let peer = random_peer_id();
cache.put(cid, vec![peer]);
assert!(cache.get(&cid).is_some());
std::thread::sleep(Duration::from_millis(100));
assert!(cache.get(&cid).is_none());
assert!(cache.needs_refresh(&cid));
let stats = cache.stats();
assert!(stats.expirations > 0);
}
#[test]
fn test_provider_cache_lru_eviction() {
let config = ProviderCacheConfig {
ttl: Duration::from_secs(3600),
max_entries: 5,
..Default::default()
};
let cache = ProviderCache::with_config(config);
let peer = random_peer_id();
for i in 0..5 {
let cid = make_cid(&[i as u8]);
cache.put(cid, vec![peer]);
}
assert_eq!(cache.len(), 5);
let cid_2 = make_cid(&[2]);
let cid_3 = make_cid(&[3]);
cache.get(&cid_2);
cache.get(&cid_3);
let new_cid = make_cid(&[100]);
cache.put(new_cid, vec![peer]);
assert!(cache.len() <= 5);
let stats = cache.stats();
assert!(stats.evictions > 0);
}
#[test]
fn test_provider_cache_stats() {
let cache = ProviderCache::new();
let cid1 = make_cid(b"one");
let cid2 = make_cid(b"two");
let peer = random_peer_id();
cache.get(&cid1);
cache.put(cid1, vec![peer]);
cache.get(&cid1);
cache.get(&cid2);
let stats = cache.stats();
assert_eq!(stats.hits, 1);
assert_eq!(stats.misses, 2);
assert_eq!(stats.entries, 1);
}
#[test]
fn test_provider_cache_invalidate() {
let cache = ProviderCache::new();
let cid = make_cid(b"invalidate me");
let peer = random_peer_id();
cache.put(cid, vec![peer]);
assert!(cache.has_providers(&cid));
cache.invalidate(&cid);
assert!(!cache.has_providers(&cid));
}
#[test]
fn test_provider_cache_cleanup() {
let config = ProviderCacheConfig {
ttl: Duration::from_millis(10),
..Default::default()
};
let cache = ProviderCache::with_config(config);
let peer = random_peer_id();
for i in 0..5 {
let cid = make_cid(&[i as u8]);
cache.put(cid, vec![peer]);
}
assert_eq!(cache.len(), 5);
std::thread::sleep(Duration::from_millis(50));
cache.cleanup_expired();
assert_eq!(cache.len(), 0);
}
#[test]
fn test_provider_cache_empty_providers_not_cached() {
let cache = ProviderCache::new();
let cid = make_cid(b"empty");
cache.put(cid, vec![]);
assert!(!cache.has_providers(&cid));
assert_eq!(cache.len(), 0);
}
}