use std::sync::Mutex;
use std::time::{Duration, Instant};
use indexmap::IndexMap;
use tracing::warn;
use crate::dht::routing::NodeId;
pub use crate::dht::records::Hash;
const DEFAULT_TTL: Duration = Duration::from_mins(5);
const DEFAULT_CAPACITY: usize = 1024;
type Key = (NodeId, Hash);
#[derive(Debug)]
struct Inner {
entries: IndexMap<Key, Instant>,
cap: usize,
ttl: Duration,
}
#[derive(Debug)]
pub struct NegativeProbeCache {
inner: Mutex<Inner>,
}
impl Default for NegativeProbeCache {
fn default() -> Self {
Self::new()
}
}
impl NegativeProbeCache {
#[must_use]
pub fn new() -> Self {
Self::with_capacity_and_ttl(DEFAULT_CAPACITY, DEFAULT_TTL)
}
#[cfg(test)]
#[must_use]
fn with_capacity(cap: usize) -> Self {
Self::with_capacity_and_ttl(cap, DEFAULT_TTL)
}
#[must_use]
pub fn with_capacity_and_ttl(cap: usize, ttl: Duration) -> Self {
let cap = cap.max(1);
Self {
inner: Mutex::new(Inner {
entries: IndexMap::with_capacity(cap),
cap,
ttl,
}),
}
}
#[must_use]
pub fn contains_active(&self, node_id: &NodeId, hash: &Hash) -> bool {
let key = (*node_id, *hash);
let now = Instant::now();
let mut guard = self.lock();
let Some(expiry) = guard.entries.get(&key).copied() else {
return false;
};
if expiry <= now {
guard.entries.shift_remove(&key);
return false;
}
guard.entries.shift_insert(0, key, expiry);
true
}
pub fn record_failure(&self, node_id: NodeId, hash: Hash) {
let ttl = self.lock().ttl;
self.record_failure_with_ttl(node_id, hash, ttl);
}
pub fn record_failure_with_ttl(&self, node_id: NodeId, hash: Hash, ttl: Duration) {
let key = (node_id, hash);
let mut guard = self.lock();
let expiry = Instant::now() + ttl;
guard.entries.shift_insert(0, key, expiry);
if guard.entries.len() > guard.cap {
guard.entries.pop();
}
}
#[must_use]
pub fn len(&self) -> usize {
self.lock().entries.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.lock().entries.is_empty()
}
fn lock(&self) -> std::sync::MutexGuard<'_, Inner> {
match self.inner.lock() {
Ok(g) => g,
Err(poisoned) => {
warn!("NegativeProbeCache mutex poisoned; recovering inner state");
poisoned.into_inner()
}
}
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::indexing_slicing
)]
mod tests;