use std::collections::{BTreeSet, HashMap};
use crate::dht::routing::NodeId;
use decdn_protocol::Coverage;
use decdn_protocol::MAX_PROVIDERS_PER_HASH;
use decdn_protocol::dht::Provider;
pub use decdn_protocol::ContentHash as Hash;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum InsertOutcome {
Inserted,
Refreshed,
RejectedQuotaExceeded,
}
impl InsertOutcome {
#[must_use]
pub const fn accepted(self) -> bool {
matches!(self, Self::Inserted | Self::Refreshed)
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct HolderEntry {
holder: NodeId,
receive_us: u64,
sequence: u64,
expiry_us: u64,
coverage: Coverage,
}
#[derive(Debug, Clone)]
pub struct RecordStoreConfig {
pub max_records_per_publisher: usize,
pub max_records_global: usize,
pub max_providers_per_hash: usize,
pub ttl_us: u64,
}
impl Default for RecordStoreConfig {
fn default() -> Self {
Self {
max_records_per_publisher: 100_000,
max_records_global: 1_000_000,
max_providers_per_hash: MAX_PROVIDERS_PER_HASH,
ttl_us: 3_600_000_000,
}
}
}
#[derive(Debug)]
pub struct RecordStore {
cfg: RecordStoreConfig,
by_hash: HashMap<Hash, Vec<HolderEntry>>,
by_publisher_count: HashMap<NodeId, usize>,
global_lru: BTreeSet<(u64, u64, Hash, NodeId)>,
next_sequence: u64,
}
impl RecordStore {
#[must_use]
pub fn new(cfg: RecordStoreConfig) -> Self {
Self {
cfg,
by_hash: HashMap::new(),
by_publisher_count: HashMap::new(),
global_lru: BTreeSet::new(),
next_sequence: 0,
}
}
#[must_use]
pub fn len(&self) -> usize {
self.global_lru.len()
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.global_lru.is_empty()
}
#[must_use]
pub const fn capacity(&self) -> usize {
self.cfg.max_records_global
}
#[must_use]
pub fn publisher_record_count(&self, holder: &NodeId) -> usize {
self.by_publisher_count.get(holder).copied().unwrap_or(0)
}
pub fn insert_at(
&mut self,
holder: NodeId,
hash: Hash,
coverage: Coverage,
receive_us: u64,
) -> InsertOutcome {
let sequence = self.next_sequence;
self.next_sequence = self.next_sequence.wrapping_add(1);
if self.remove_entry(holder, hash).is_some() {
self.add_entry(holder, hash, coverage, receive_us, sequence);
return InsertOutcome::Refreshed;
}
if self.publisher_record_count(&holder) >= self.cfg.max_records_per_publisher {
return InsertOutcome::RejectedQuotaExceeded;
}
let evict_holder = self.by_hash.get(&hash).and_then(|entries| {
(entries.len() >= self.cfg.max_providers_per_hash)
.then(|| {
entries
.iter()
.min_by_key(|e| e.receive_us)
.map(|e| e.holder)
})
.flatten()
});
let grows_global = match evict_holder {
Some(old_holder) => {
self.remove_entry(old_holder, hash);
false
}
None => true,
};
if grows_global && self.global_lru.len() >= self.cfg.max_records_global {
self.evict_globally_oldest();
}
self.add_entry(holder, hash, coverage, receive_us, sequence);
InsertOutcome::Inserted
}
pub fn providers_at(&mut self, hash: &Hash, now_us: u64) -> Vec<Provider> {
self.scrub_hash_expired(hash, now_us);
self.by_hash
.get(hash)
.map(|entries| {
entries
.iter()
.map(|e| Provider {
node: e.holder,
coverage: e.coverage.clone(),
})
.collect()
})
.unwrap_or_default()
}
pub fn gc(&mut self, now_us: u64) -> usize {
let mut removed = 0usize;
while let Some(&key @ (receive_us, _, hash, holder)) = self.global_lru.iter().next() {
let expiry_us = receive_us.saturating_add(self.cfg.ttl_us);
if expiry_us > now_us {
break;
}
if self.remove_entry(holder, hash).is_none() {
self.global_lru.remove(&key);
tracing::error!(
?holder,
?hash,
"RecordStore::gc: global_lru key with no by_hash entry; dropping orphan"
);
}
removed += 1;
}
removed
}
fn scrub_hash_expired(&mut self, hash: &Hash, now_us: u64) -> usize {
let Some(entries) = self.by_hash.get(hash) else {
return 0;
};
let expired: Vec<NodeId> = entries
.iter()
.filter(|e| e.expiry_us <= now_us)
.map(|e| e.holder)
.collect();
for holder in &expired {
self.remove_entry(*holder, *hash);
}
expired.len()
}
fn evict_globally_oldest(&mut self) {
let Some(&(_, _, hash, holder)) = self.global_lru.iter().next() else {
return;
};
self.remove_entry(holder, hash);
}
fn add_entry(
&mut self,
holder: NodeId,
hash: Hash,
coverage: Coverage,
receive_us: u64,
sequence: u64,
) {
let expiry_us = receive_us.saturating_add(self.cfg.ttl_us);
self.by_hash.entry(hash).or_default().push(HolderEntry {
holder,
receive_us,
sequence,
expiry_us,
coverage,
});
self.global_lru.insert((receive_us, sequence, hash, holder));
*self.by_publisher_count.entry(holder).or_insert(0) += 1;
}
fn remove_entry(&mut self, holder: NodeId, hash: Hash) -> Option<HolderEntry> {
let entries = self.by_hash.get_mut(&hash)?;
let idx = entries.iter().position(|e| e.holder == holder)?;
let Some(removed) = entries.get(idx).cloned() else {
tracing::error!(
?holder,
?hash,
idx,
"RecordStore::remove_entry: position->Some but get->None; tri-index corruption"
);
return None;
};
entries.swap_remove(idx);
if entries.is_empty() {
self.by_hash.remove(&hash);
}
self.global_lru
.remove(&(removed.receive_us, removed.sequence, hash, removed.holder));
Self::decrement_publisher_count(&mut self.by_publisher_count, &removed.holder);
Some(removed)
}
fn decrement_publisher_count(counts: &mut HashMap<NodeId, usize>, holder: &NodeId) {
if let Some(c) = counts.get_mut(holder) {
*c = c.saturating_sub(1);
if *c == 0 {
counts.remove(holder);
}
}
}
}
#[cfg(test)]
#[allow(
clippy::unwrap_used,
clippy::expect_used,
clippy::panic,
clippy::indexing_slicing,
clippy::cast_possible_truncation
)]
mod tests;