use std::collections::HashMap;
use dig_dht::CandidateAddr;
use dig_nat::{PeerId, TraversalKind};
use crate::quality::PeerQuality;
use crate::types::{Candidate, Provenance};
pub const DISPATCH_TTL_SECS: u64 = 300;
#[derive(Debug, Clone)]
pub struct PeerEntry {
pub peer_id: PeerId,
pub addresses: Vec<CandidateAddr>,
pub connection_class: Option<TraversalKind>,
pub provenance: Provenance,
pub quality: PeerQuality,
pub first_seen: u64,
pub last_outcome_at: Option<u64>,
pub connected: bool,
pub banned: bool,
pub(crate) in_flight_since: Option<u64>,
}
impl PeerEntry {
pub fn cold(peer_id: PeerId, provenance: Provenance, now: u64) -> Self {
PeerEntry {
peer_id,
addresses: Vec::new(),
connection_class: None,
provenance,
quality: PeerQuality::cold(),
first_seen: now,
last_outcome_at: None,
connected: false,
banned: false,
in_flight_since: None,
}
}
pub fn is_eligible(&self) -> bool {
!self.banned
}
pub fn is_evictable(&self) -> bool {
!self.connected && self.quality.in_flight == 0
}
pub fn is_force_evictable(&self) -> bool {
!self.connected
}
pub fn preferred_address(&self) -> Option<&CandidateAddr> {
self.addresses
.iter()
.filter(|a| a.kind.is_dialable())
.min_by_key(|a| a.kind.rank())
}
pub(crate) fn eviction_value(&self, now: u64) -> f64 {
let tput = self.quality.throughput.value().unwrap_or(0.0);
let conf = self.quality.confidence();
let rel = self.quality.reliability.rate().unwrap_or(0.0);
let quality_value = tput * conf * (0.5 + 0.5 * rel);
let reference = self.last_outcome_at.unwrap_or(self.first_seen);
let age = now.saturating_sub(reference) as f64;
let recency = 1.0 / (1.0 + age / 3600.0);
quality_value * recency + recency }
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FeedResult {
Inserted,
Updated,
MarkedDisconnected,
Absent,
}
#[derive(Debug, Default)]
pub struct Registry {
entries: HashMap<PeerId, PeerEntry>,
capacity: usize,
evicted: Vec<PeerId>,
}
impl Registry {
pub fn new(capacity: usize) -> Self {
Registry {
entries: HashMap::new(),
capacity: capacity.max(1),
evicted: Vec::new(),
}
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn drain_evicted(&mut self) -> Vec<PeerId> {
std::mem::take(&mut self.evicted)
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn get(&self, peer: &PeerId) -> Option<&PeerEntry> {
self.entries.get(peer)
}
pub(crate) fn get_mut(&mut self, peer: &PeerId) -> Option<&mut PeerEntry> {
self.entries.get_mut(peer)
}
pub fn iter(&self) -> impl Iterator<Item = &PeerEntry> {
self.entries.values()
}
pub fn upsert_candidate(
&mut self,
candidate: &Candidate,
provenance: Provenance,
now: u64,
) -> FeedResult {
let result = match self.entries.get_mut(&candidate.peer_id) {
Some(existing) => {
if !candidate.addresses.is_empty() {
existing.addresses = candidate.addresses.clone();
}
if candidate.class.is_some() {
existing.connection_class = candidate.class;
}
if provenance.evidence() > existing.provenance.evidence() {
existing.provenance = provenance;
}
FeedResult::Updated
}
None => {
let mut entry = PeerEntry::cold(candidate.peer_id, provenance, now);
entry.addresses = candidate.addresses.clone();
entry.connection_class = candidate.class;
self.seed_class_prior(&mut entry);
self.entries.insert(candidate.peer_id, entry);
FeedResult::Inserted
}
};
if matches!(result, FeedResult::Inserted) {
self.enforce_capacity(now);
}
result
}
pub fn mark_connected(&mut self, peer: PeerId, provenance: Provenance, now: u64) -> FeedResult {
match self.entries.get_mut(&peer) {
Some(existing) => {
existing.connected = true;
existing.banned = false;
if provenance.evidence() > existing.provenance.evidence() {
existing.provenance = provenance;
}
FeedResult::Updated
}
None => {
let mut entry = PeerEntry::cold(peer, provenance, now);
entry.connected = true;
self.entries.insert(peer, entry);
self.enforce_capacity(now);
FeedResult::Inserted
}
}
}
pub fn mark_disconnected(&mut self, peer: &PeerId, banned: bool) -> FeedResult {
match self.entries.get_mut(peer) {
Some(existing) => {
existing.connected = false;
if banned {
existing.banned = true;
}
FeedResult::MarkedDisconnected
}
None => FeedResult::Absent,
}
}
pub fn set_connection_class(
&mut self,
peer: PeerId,
class: TraversalKind,
now: u64,
) -> FeedResult {
match self.entries.get_mut(&peer) {
Some(existing) => {
existing.connection_class = Some(class);
FeedResult::Updated
}
None => {
let mut entry = PeerEntry::cold(peer, Provenance::Nat, now);
entry.connection_class = Some(class);
self.seed_class_prior(&mut entry);
self.entries.insert(peer, entry);
self.enforce_capacity(now);
FeedResult::Inserted
}
}
}
pub fn remove(&mut self, peer: &PeerId) -> FeedResult {
match self.entries.get(peer) {
Some(e) if e.quality.in_flight > 0 => FeedResult::Updated, Some(_) => {
self.entries.remove(peer);
self.evicted.push(*peer);
FeedResult::MarkedDisconnected
}
None => FeedResult::Absent,
}
}
pub(crate) fn add_in_flight(&mut self, peer: &PeerId, count: u32, now: u64) {
if let Some(e) = self.entries.get_mut(peer) {
let was_idle = e.quality.in_flight == 0;
e.quality.in_flight = e.quality.in_flight.saturating_add(count);
if was_idle && e.quality.in_flight > 0 {
e.in_flight_since = Some(now);
}
}
}
pub(crate) fn release_in_flight(&mut self, peer: &PeerId, count: u32) {
if let Some(e) = self.entries.get_mut(peer) {
e.quality.in_flight = e.quality.in_flight.saturating_sub(count);
if e.quality.in_flight == 0 {
e.in_flight_since = None;
}
}
}
pub fn reclaim_stale_in_flight(&mut self, now: u64) {
for e in self.entries.values_mut() {
if let Some(since) = e.in_flight_since {
if now.saturating_sub(since) >= DISPATCH_TTL_SECS {
e.quality.in_flight = 0;
e.in_flight_since = None;
}
}
}
}
fn seed_class_prior(&self, entry: &mut PeerEntry) {
entry.quality.reliability.seed_prior(0.5);
}
fn enforce_capacity(&mut self, now: u64) {
self.reclaim_stale_in_flight(now);
while self.entries.len() > self.capacity {
let victim = self
.entries
.values()
.filter(|e| e.is_evictable())
.min_by(|a, b| {
a.eviction_value(now)
.partial_cmp(&b.eviction_value(now))
.unwrap_or(std::cmp::Ordering::Equal)
})
.map(|e| e.peer_id)
.or_else(|| {
self.entries
.values()
.filter(|e| e.is_force_evictable())
.min_by(|a, b| {
a.eviction_value(now)
.partial_cmp(&b.eviction_value(now))
.unwrap_or(std::cmp::Ordering::Equal)
})
.map(|e| e.peer_id)
});
match victim {
Some(id) => {
self.entries.remove(&id);
self.evicted.push(id);
}
None => break,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use dig_dht::AddressKind;
fn pid(b: u8) -> PeerId {
PeerId::from_bytes([b; 32])
}
fn cand(b: u8) -> Candidate {
Candidate::new(pid(b), vec![CandidateAddr::direct("10.0.0.1", 9444)])
}
#[test]
fn upsert_candidate_inserts_cold_then_updates_preserving_quality() {
let mut r = Registry::new(100);
assert_eq!(
r.upsert_candidate(&cand(1), Provenance::Dht, 100),
FeedResult::Inserted
);
assert!(r.get(&pid(1)).unwrap().quality.is_cold());
r.get_mut(&pid(1))
.unwrap()
.quality
.observe_throughput(500.0);
r.get_mut(&pid(1)).unwrap().quality.bump_samples();
assert_eq!(
r.upsert_candidate(&cand(1), Provenance::Dht, 200),
FeedResult::Updated
);
assert!(!r.get(&pid(1)).unwrap().quality.is_cold());
}
#[test]
fn disconnect_retains_entry_and_quality() {
let mut r = Registry::new(100);
r.mark_connected(pid(1), Provenance::Gossip, 100);
r.get_mut(&pid(1))
.unwrap()
.quality
.observe_throughput(700.0);
r.get_mut(&pid(1)).unwrap().quality.bump_samples();
assert_eq!(
r.mark_disconnected(&pid(1), false),
FeedResult::MarkedDisconnected
);
let e = r.get(&pid(1)).unwrap();
assert!(!e.connected);
assert!(!e.quality.is_cold(), "quality retained across disconnect");
assert!(e.is_eligible(), "a plain disconnect stays eligible");
}
#[test]
fn banned_peer_is_ineligible_until_reconnect() {
let mut r = Registry::new(100);
r.mark_connected(pid(1), Provenance::Gossip, 100);
r.mark_disconnected(&pid(1), true);
assert!(!r.get(&pid(1)).unwrap().is_eligible());
r.mark_connected(pid(1), Provenance::Gossip, 200);
assert!(r.get(&pid(1)).unwrap().is_eligible());
}
#[test]
fn stronger_provenance_wins_weaker_ignored() {
let mut r = Registry::new(100);
r.upsert_candidate(&cand(1), Provenance::Pex, 100);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
assert_eq!(r.get(&pid(1)).unwrap().provenance, Provenance::Dht);
r.upsert_candidate(&cand(1), Provenance::Pex, 100);
assert_eq!(r.get(&pid(1)).unwrap().provenance, Provenance::Dht);
}
#[test]
fn eviction_sheds_lowest_value_never_connected_or_in_flight() {
let mut r = Registry::new(2);
r.mark_connected(pid(1), Provenance::Gossip, 100);
r.mark_connected(pid(2), Provenance::Gossip, 100);
r.upsert_candidate(&cand(3), Provenance::Dht, 100);
assert!(r.get(&pid(1)).is_some());
assert!(r.get(&pid(2)).is_some());
assert!(
r.get(&pid(3)).is_none(),
"connected peers protected; cold candidate evicted"
);
}
#[test]
fn eviction_prefers_shedding_stale_unmeasured_over_fresh_measured() {
let mut r = Registry::new(1);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
let e = r.get_mut(&pid(1)).unwrap();
e.quality.observe_throughput(900.0);
e.quality.bump_samples();
e.last_outcome_at = Some(1000);
r.upsert_candidate(&cand(2), Provenance::Dht, 5000);
assert!(
r.get(&pid(1)).is_some(),
"the valuable measured peer survives"
);
assert!(r.get(&pid(2)).is_none(), "the stale cold peer is evicted");
}
#[test]
fn in_flight_peer_is_not_removed() {
let mut r = Registry::new(100);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
r.add_in_flight(&pid(1), 2, 100);
assert_eq!(r.remove(&pid(1)), FeedResult::Updated); assert!(r.get(&pid(1)).is_some());
r.release_in_flight(&pid(1), 2);
assert_eq!(r.remove(&pid(1)), FeedResult::MarkedDisconnected);
assert!(r.get(&pid(1)).is_none());
}
#[test]
fn set_connection_class_upserts_and_attaches() {
let mut r = Registry::new(100);
assert_eq!(
r.set_connection_class(pid(1), TraversalKind::Relayed, 100),
FeedResult::Inserted
);
assert_eq!(
r.get(&pid(1)).unwrap().connection_class,
Some(TraversalKind::Relayed)
);
}
#[test]
fn capacity_bound_holds_even_when_every_entry_is_pinned_by_stale_in_flight() {
let capacity = 4;
let mut r = Registry::new(capacity);
for b in 0..(capacity as u8 * 3) {
r.upsert_candidate(&cand(b), Provenance::Dht, 100);
r.add_in_flight(&pid(b), 3, 100); }
assert!(
r.len() <= capacity,
"registry must stay bounded at {capacity} even when every entry has stuck in_flight; got {}",
r.len()
);
}
#[test]
fn stale_in_flight_is_reclaimed_after_ttl_making_the_entry_evictable_again() {
let mut r = Registry::new(100);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
r.add_in_flight(&pid(1), 5, 100);
assert!(
!r.get(&pid(1)).unwrap().is_evictable(),
"pinned while in flight"
);
r.reclaim_stale_in_flight(100 + DISPATCH_TTL_SECS + 1);
assert_eq!(
r.get(&pid(1)).unwrap().quality.in_flight,
0,
"a dispatch that never settled must be reclaimed after its TTL"
);
assert!(
r.get(&pid(1)).unwrap().is_evictable(),
"reclaiming stale in_flight must make the entry evictable again"
);
}
#[test]
fn drain_evicted_reports_peers_shed_by_capacity_enforcement() {
let mut r = Registry::new(1);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
assert!(r.drain_evicted().is_empty(), "no eviction yet");
r.upsert_candidate(&cand(2), Provenance::Dht, 200);
let evicted = r.drain_evicted();
assert_eq!(evicted, vec![pid(1)], "capacity eviction must be reported");
assert!(r.drain_evicted().is_empty());
}
#[test]
fn drain_evicted_reports_peers_shed_by_explicit_remove() {
let mut r = Registry::new(100);
r.upsert_candidate(&cand(1), Provenance::Dht, 100);
r.remove(&pid(1));
assert_eq!(r.drain_evicted(), vec![pid(1)]);
}
#[test]
fn preferred_address_is_most_direct() {
let mut r = Registry::new(100);
let c = Candidate::new(
pid(1),
vec![
CandidateAddr {
host: "r".into(),
port: 1,
kind: AddressKind::Reflexive,
},
CandidateAddr::direct("d", 2),
],
);
r.upsert_candidate(&c, Provenance::Dht, 100);
assert_eq!(
r.get(&pid(1)).unwrap().preferred_address().unwrap().kind,
AddressKind::Direct
);
}
}