syncweb_core/sync/

peer_tracker.rs

use std::{
    collections::{HashMap, HashSet},
    time::{Duration, Instant},
};

use iroh::PublicKey;
use iroh_blobs::Hash;

/// Policy used when the peer availability cache reaches its size limit.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[non_exhaustive]
pub enum EvictionStrategy {
    Lru,
    Fifo,
}

#[derive(Clone, Debug)]
struct Availability {
    peers: HashSet<PublicKey>,
    inserted_at: u64,
    accessed_at: u64,
    last_seen: Instant,
}

/// Bounded cache mapping blob hashes to the peers observed serving them.
///
/// `len()` returns the number of unique blob hashes tracked, not the total
/// number of peers. Use [`peer_count`](Self::peer_count) for per-blob counts.
#[derive(Clone, Debug)]
pub struct PeerTracker {
    entries: HashMap<Hash, Availability>,
    max_cache_size: usize,
    strategy: EvictionStrategy,
    expiry: Duration,
    clock: u64,
}

impl PeerTracker {
    #[must_use]
    pub fn new(max_cache_size: usize, strategy: EvictionStrategy) -> Self {
        Self::with_expiry(max_cache_size, strategy, Duration::from_mins(5))
    }

    #[must_use]
    pub fn with_expiry(max_cache_size: usize, strategy: EvictionStrategy, expiry: Duration) -> Self {
        Self {
            entries: HashMap::new(),
            max_cache_size,
            strategy,
            expiry,
            clock: 0,
        }
    }

    pub fn record_peer(&mut self, blob_hash: Hash, node_id: PublicKey) {
        self.clock = self.clock.saturating_add(1);
        let now = Instant::now();
        let entry = self.entries.entry(blob_hash).or_insert_with(|| Availability {
            peers: HashSet::new(),
            inserted_at: self.clock,
            accessed_at: self.clock,
            last_seen: now,
        });
        entry.peers.insert(node_id);
        entry.accessed_at = self.clock;
        entry.last_seen = now;
        self.evict_to_limit();
    }

    #[must_use]
    pub fn get_peers(&mut self, blob_hash: &Hash) -> Vec<PublicKey> {
        self.clock = self.clock.saturating_add(1);
        self.entries.get_mut(blob_hash).map_or_else(Vec::new, |entry| {
            if self.strategy == EvictionStrategy::Lru {
                entry.accessed_at = self.clock;
            }
            entry.peers.iter().copied().collect()
        })
    }

    #[must_use]
    pub fn peer_count(&self, blob_hash: &Hash) -> usize {
        self.entries.get(blob_hash).map_or(0, |entry| entry.peers.len())
    }

    #[must_use]
    pub fn contains(&self, blob_hash: &Hash, node_id: &PublicKey) -> bool {
        self.entries
            .get(blob_hash)
            .is_some_and(|entry| entry.peers.contains(node_id))
    }

    /// Remove a peer from all cached blob availability records.
    pub fn on_peer_disconnected(&mut self, node_id: &PublicKey) {
        self.entries.retain(|_, entry| {
            entry.peers.remove(node_id);
            !entry.peers.is_empty()
        });
    }

    #[must_use]
    pub fn is_fresh(&self, blob_hash: &Hash) -> bool {
        self.entries
            .get(blob_hash)
            .is_some_and(|entry| entry.last_seen.elapsed() <= self.expiry)
    }

    #[must_use]
    pub fn len(&self) -> usize {
        self.entries.len()
    }

    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    pub fn tick_and_maybe_evict(&mut self) {
        let now = Instant::now();
        self.entries
            .retain(|_, entry| now.saturating_duration_since(entry.last_seen) <= self.expiry);
        self.evict_to_limit();
    }

    fn evict_to_limit(&mut self) {
        while self.entries.len() > self.max_cache_size {
            let oldest = self
                .entries
                .iter()
                .min_by_key(|(_, entry)| match self.strategy {
                    EvictionStrategy::Lru => entry.accessed_at,
                    EvictionStrategy::Fifo => entry.inserted_at,
                })
                .map(|(hash, _)| *hash);
            if let Some(hash) = oldest {
                self.entries.remove(&hash);
            } else {
                break;
            }
        }
    }
}

/// Compact peer presence cache using stable integer indices.
#[derive(Clone, Debug, Default)]
pub struct EfficientPeerCache {
    hash_to_indices: HashMap<Hash, Vec<u16>>,
    presence: HashMap<Hash, Vec<u64>>,
    peer_to_index: HashMap<PublicKey, u16>,
    index_peer: Vec<PublicKey>,
    last_seen: HashMap<PublicKey, Instant>,
}

impl EfficientPeerCache {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }

    /// Record that a peer can serve a blob.
    ///
    /// # Errors
    ///
    /// Returns an error after all `u16` peer indices have been allocated.
    pub fn add_presence(&mut self, hash: Hash, peer: PublicKey) -> crate::Result<()> {
        let index = if let Some(index) = self.peer_to_index.get(&peer) {
            *index
        } else {
            let index = u16::try_from(self.index_peer.len())
                .map_err(|error| crate::SyncwebError::operation("peer cache capacity exceeded", error))?;
            self.index_peer.push(peer);
            self.peer_to_index.insert(peer, index);
            index
        };
        self.last_seen.insert(peer, Instant::now());
        let indices = self.hash_to_indices.entry(hash).or_default();
        match indices.binary_search(&index) {
            Ok(_) => {}
            Err(position) => indices.insert(position, index),
        }
        let index_usize = usize::from(index);
        let word_index = index_usize.checked_div(64).unwrap_or(0);
        let bit_index = index_usize.checked_rem(64).unwrap_or(0);
        let mask = 1_u64.checked_shl(u32::try_from(bit_index).unwrap_or(0)).unwrap_or(0);
        let presence = self.presence.entry(hash).or_default();
        presence.resize(word_index.saturating_add(1), 0);
        if let Some(word) = presence.get_mut(word_index) {
            *word |= mask;
        }
        Ok(())
    }

    pub fn remove_presence(&mut self, hash: &Hash, peer: &PublicKey) -> bool {
        let Some(index) = self.peer_to_index.get(peer) else {
            return false;
        };
        let Some(indices) = self.hash_to_indices.get_mut(hash) else {
            return false;
        };
        let removed = indices.binary_search(index).is_ok_and(|position| {
            indices.remove(position);
            true
        });
        if indices.is_empty() {
            self.hash_to_indices.remove(hash);
        }
        if removed {
            let index_usize = usize::from(*index);
            let word_index = index_usize.checked_div(64).unwrap_or(0);
            let bit_index = index_usize.checked_rem(64).unwrap_or(0);
            let mask = 1_u64.checked_shl(u32::try_from(bit_index).unwrap_or(0)).unwrap_or(0);
            if let Some(word) = self.presence.get_mut(hash).and_then(|words| words.get_mut(word_index)) {
                *word &= !mask;
            }
        }
        removed
    }

    #[must_use]
    pub fn has_peer(&self, hash: &Hash, peer: &PublicKey) -> bool {
        self.peer_to_index.get(peer).is_some_and(|index| {
            let index_usize = usize::from(*index);
            let word_index = index_usize.checked_div(64).unwrap_or(0);
            let bit_index = index_usize.checked_rem(64).unwrap_or(0);
            let mask = 1_u64.checked_shl(u32::try_from(bit_index).unwrap_or(0)).unwrap_or(0);
            self.presence
                .get(hash)
                .and_then(|words| words.get(word_index))
                .is_some_and(|word| *word & mask != 0)
        })
    }

    #[must_use]
    pub fn peers(&self, hash: &Hash) -> Vec<PublicKey> {
        self.hash_to_indices.get(hash).map_or_else(Vec::new, |indices| {
            indices
                .iter()
                .filter_map(|index| self.index_peer.get(usize::from(*index)).copied())
                .collect()
        })
    }

    /// Return the total number of *unique* peers registered across all blobs.
    /// For per-blob counts see [`active_peer_count`](Self::active_peer_count).
    #[must_use]
    pub const fn total_peers(&self) -> usize {
        self.index_peer.len()
    }

    #[must_use]
    pub fn active_peer_count(&self, hash: &Hash) -> usize {
        self.hash_to_indices.get(hash).map_or(0, Vec::len)
    }

    #[must_use]
    pub fn memory_usage(&self) -> usize {
        let word_bytes = self
            .presence
            .values()
            .map(|words| words.len().saturating_mul(8))
            .sum::<usize>();
        let index_bytes = self
            .hash_to_indices
            .values()
            .map(|indices| indices.len().saturating_mul(2))
            .sum::<usize>();
        word_bytes
            .saturating_add(index_bytes)
            .saturating_add(
                self.peer_to_index
                    .len()
                    .saturating_mul(std::mem::size_of::<PublicKey>().saturating_add(2)),
            )
            .saturating_add(self.index_peer.len().saturating_mul(std::mem::size_of::<PublicKey>()))
    }

    /// Remove a peer from every blob while retaining stable indices for peers
    /// that may be observed again later.
    pub fn remove_peer(&mut self, peer: &PublicKey) {
        if !self.peer_to_index.contains_key(peer) {
            return;
        }
        let hashes = self.hash_to_indices.keys().copied().collect::<Vec<_>>();
        for hash in hashes {
            let _removed = self.remove_presence(&hash, peer);
        }
        self.last_seen.remove(peer);
    }

    /// Remove peers whose `last_seen` timestamp exceeds the given TTL.
    ///
    /// Fully reclaims each removed peer's index, so the total number of
    /// tracked peers stays bounded.
    ///
    /// Returns the number of peers removed.
    pub fn retain_active(&mut self, ttl: Duration) -> usize {
        self.retain_active_at(ttl, Instant::now())
    }

    /// Like [`retain_active`](Self::retain_active) but with an explicit
    /// observation time for deterministic testing.
    fn retain_active_at(&mut self, ttl: Duration, now: Instant) -> usize {
        let cutoff = now.checked_sub(ttl).unwrap_or(now);
        let stale: Vec<PublicKey> = self
            .last_seen
            .iter()
            .filter(|(_, last_seen)| **last_seen < cutoff)
            .map(|(peer, _)| *peer)
            .collect();
        let count = stale.len();
        for peer in stale {
            self.remove_peer(&peer);
            if let Some(stale_index) = self.peer_to_index.remove(&peer) {
                let last_index = self.index_peer.len().saturating_sub(1);
                if usize::from(stale_index) != last_index
                    && let Some(last_peer) = self.index_peer.get(last_index).copied()
                {
                    // Update presence bitmask for the moved peer: its bit
                    // is still at last_index but must be moved to stale_index.
                    self.remap_presence(last_index, usize::from(stale_index));
                    if let Some(slot) = self.index_peer.get_mut(usize::from(stale_index)) {
                        *slot = last_peer;
                    }
                    self.peer_to_index.insert(last_peer, stale_index);
                }
                self.index_peer.pop();
            }
        }
        count
    }

    /// Move all presence bits from `old_index` to `new_index` after a
    /// swap-remove compaction of `index_peer`.
    fn remap_presence(&mut self, old_index: usize, new_index: usize) {
        let old_u16 = u16::try_from(old_index).unwrap_or(u16::MAX);
        let new_u16 = u16::try_from(new_index).unwrap_or(u16::MAX);
        for (hash, indices) in &mut self.hash_to_indices {
            // binary_search returns a valid index in the Ok case
            let Ok(pos) = indices.binary_search(&old_u16) else {
                continue;
            };
            if let Some(slot) = indices.get_mut(pos) {
                *slot = new_u16;
            }
            let mut p = pos;
            while p > 0 {
                let prev = p.wrapping_sub(1);
                if indices.get(prev).is_none_or(|&v| v <= new_u16) {
                    break;
                }
                indices.swap(prev, p);
                p = prev;
            }
            // Update bitvector: clear old bit, set new bit
            if let Some(words) = self.presence.get_mut(hash) {
                let old_word = old_index.checked_div(64).unwrap_or(0);
                let old_bit = old_index.checked_rem(64).unwrap_or(0);
                if let Some(word) = words.get_mut(old_word) {
                    *word &= !(1_u64 << old_bit);
                }
                let new_word = new_index.checked_div(64).unwrap_or(0);
                let new_bit = new_index.checked_rem(64).unwrap_or(0);
                if new_word >= words.len() {
                    words.resize(new_word.saturating_add(1), 0);
                }
                if let Some(word) = words.get_mut(new_word) {
                    *word |= 1_u64 << new_bit;
                }
            }
        }
    }
}