shadow_dht/

replication.rs

//! DHT data replication manager.
//!
//! Tracks which values need replication across the network, schedules
//! replication tasks, handles periodic refresh, and re-replicates when
//! peers depart.

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

use bytes::Bytes;
use shadow_core::{PeerId, PeerInfo};

use crate::node::DHTNode;

// ── Replication priority ────────────────────────────────────────────────────

/// Priority level for a replication task
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum ReplicationPriority {
    /// Background maintenance (periodic refresh)
    Low = 0,
    /// Standard replication of existing data
    Normal = 1,
    /// Urgent — new data or peer loss requiring immediate re-replication
    High = 2,
}

// ── Replication task ────────────────────────────────────────────────────────

/// A unit of work describing data that must be sent to a remote peer.
#[derive(Debug, Clone)]
pub struct ReplicationTask {
    /// The key of the value to replicate
    pub key: [u8; 32],
    /// The serialised value
    pub value: Bytes,
    /// The peer that should receive the replica
    pub target_peer: PeerId,
    /// How urgently this should be replicated
    pub priority: ReplicationPriority,
    /// When this task was created
    pub created_at: Instant,
    /// Optional deadline — if `None`, best-effort
    pub deadline: Option<Instant>,
}

// ── Replica tracking ────────────────────────────────────────────────────────

#[derive(Debug, Clone)]
struct ReplicaInfo {
    /// Peers known to hold a copy
    holders: HashSet<PeerId>,
    /// When the value was first stored locally
    stored_at: Instant,
    /// When we last refreshed the replicas
    last_refresh: Instant,
    /// Size of the value in bytes
    value_size: usize,
}

// ── Stats ───────────────────────────────────────────────────────────────────

/// Summary statistics for the replication sub-system
#[derive(Debug, Clone, Default)]
pub struct ReplicationStats {
    pub total_keys: usize,
    pub fully_replicated: usize,
    pub under_replicated: usize,
    pub pending_tasks: usize,
    pub total_replicas: usize,
}

// ── Config ──────────────────────────────────────────────────────────────────

/// Configuration for the replication manager
#[derive(Debug, Clone)]
pub struct ReplicationConfig {
    /// Number of copies to maintain across the network (default: 3)
    pub factor: usize,
    /// Interval between refresh sweeps (default: 1 hour)
    pub refresh_interval: Duration,
    /// Value time-to-live (default: 24 hours)
    pub ttl: Duration,
    /// Maximum pending tasks before back-pressure (default: 1000)
    pub max_pending: usize,
}

impl Default for ReplicationConfig {
    fn default() -> Self {
        Self {
            factor: 3,
            refresh_interval: Duration::from_secs(3600),
            ttl: Duration::from_secs(86400),
            max_pending: 1000,
        }
    }
}

// ── ReplicationManager ──────────────────────────────────────────────────────

/// Coordinates data replication across the DHT.
///
/// The manager does **not** perform any networking itself — it produces
/// [`ReplicationTask`]s that a higher-level transport layer executes.
pub struct ReplicationManager {
    config: ReplicationConfig,
    /// key → replica tracking info
    replicas: HashMap<[u8; 32], ReplicaInfo>,
    /// Pending replication work
    pending: VecDeque<ReplicationTask>,
}

impl ReplicationManager {
    /// Create a new replication manager with the given configuration.
    pub fn new(config: ReplicationConfig) -> Self {
        Self {
            config,
            replicas: HashMap::new(),
            pending: VecDeque::new(),
        }
    }

    /// Create with default config but a custom replication factor.
    pub fn with_factor(factor: usize) -> Self {
        Self::new(ReplicationConfig {
            factor,
            ..Default::default()
        })
    }

    // ── Core operations ─────────────────────────────────────────────────

    /// Schedule replication for a newly stored (or updated) value.
    ///
    /// Returns the tasks that should be executed by the transport layer.
    pub fn schedule_replication(
        &mut self,
        key: [u8; 32],
        value: Bytes,
        closest_peers: &[PeerInfo],
    ) -> Vec<ReplicationTask> {
        let now = Instant::now();

        // Insert or update tracking entry
        let info = self.replicas.entry(key).or_insert_with(|| ReplicaInfo {
            holders: HashSet::new(),
            stored_at: now,
            last_refresh: now,
            value_size: value.len(),
        });
        info.last_refresh = now;
        info.value_size = value.len();

        // Determine how many additional replicas we need
        let current = info.holders.len();
        let needed = self.config.factor.saturating_sub(current);

        let mut tasks = Vec::new();

        // Select peers that don't already hold a copy
        for peer in closest_peers.iter().take(needed + current) {
            if !info.holders.contains(&peer.id) && tasks.len() < needed {
                let task = ReplicationTask {
                    key,
                    value: value.clone(),
                    target_peer: peer.id,
                    priority: ReplicationPriority::High,
                    created_at: now,
                    deadline: Some(now + Duration::from_secs(60)),
                };
                tasks.push(task.clone());
                self.pending.push_back(task);
            }
        }

        tasks
    }

    /// Confirm that a replica was successfully stored on a peer.
    pub fn confirm_replication(&mut self, key: [u8; 32], peer: PeerId) {
        if let Some(info) = self.replicas.get_mut(&key) {
            info.holders.insert(peer);
        }
    }

    /// Drain and return all pending tasks.
    pub fn take_pending_tasks(&mut self) -> Vec<ReplicationTask> {
        self.pending.drain(..).collect()
    }

    /// Peek at the number of pending tasks.
    pub fn pending_count(&self) -> usize {
        self.pending.len()
    }

    // ── Maintenance ─────────────────────────────────────────────────────

    /// Scan for keys that need a refresh (their replicas haven't been
    /// refreshed within `refresh_interval`).
    ///
    /// The caller must supply the local DHT node so we can read stored
    /// values and closest peers.
    pub fn check_refresh(&mut self, dht: &DHTNode) -> Vec<ReplicationTask> {
        let now = Instant::now();
        let mut tasks = Vec::new();

        let stale_keys: Vec<[u8; 32]> = self
            .replicas
            .iter()
            .filter(|(_, info)| now.duration_since(info.last_refresh) >= self.config.refresh_interval)
            .map(|(k, _)| *k)
            .collect();

        for key in stale_keys {
            if let Some(stored) = dht.get_value(&key) {
                let key_peer = PeerId::from_bytes(key);
                let closest = dht.find_closest_peers(&key_peer, self.config.factor);
                let new_tasks = self.schedule_replication(key, stored.data.clone(), &closest);
                tasks.extend(new_tasks);
            }
        }

        tasks
    }

    /// Handle the loss of a peer — re-replicate any data it held.
    pub fn handle_peer_lost(&mut self, lost_peer: &PeerId, dht: &DHTNode) -> Vec<ReplicationTask> {
        let now = Instant::now();
        let mut tasks = Vec::new();

        // Find all keys where the lost peer was a replica holder
        let affected_keys: Vec<[u8; 32]> = self
            .replicas
            .iter()
            .filter(|(_, info)| info.holders.contains(lost_peer))
            .map(|(k, _)| *k)
            .collect();

        for key in affected_keys {
            // Remove the lost peer from holders
            if let Some(info) = self.replicas.get_mut(&key) {
                info.holders.remove(lost_peer);
            }

            // Re-replicate if under-replicated
            if let Some(stored) = dht.get_value(&key) {
                let key_peer = PeerId::from_bytes(key);
                let closest = dht.find_closest_peers(&key_peer, self.config.factor * 2);

                if let Some(info) = self.replicas.get(&key) {
                    let current = info.holders.len();
                    let needed = self.config.factor.saturating_sub(current);

                    for peer in closest.iter() {
                        if !info.holders.contains(&peer.id) && tasks.len() < needed {
                            let task = ReplicationTask {
                                key,
                                value: stored.data.clone(),
                                target_peer: peer.id,
                                priority: ReplicationPriority::High,
                                created_at: now,
                                deadline: Some(now + Duration::from_secs(30)),
                            };
                            tasks.push(task.clone());
                            self.pending.push_back(task);
                        }
                    }
                }
            }
        }

        tasks
    }

    /// Remove tracking for expired keys.
    pub fn cleanup_expired(&mut self) -> usize {
        let now = Instant::now();
        let before = self.replicas.len();
        self.replicas.retain(|_, info| {
            now.duration_since(info.stored_at) < self.config.ttl
        });
        before - self.replicas.len()
    }

    // ── Queries ─────────────────────────────────────────────────────────

    /// How many confirmed replicas exist for a key.
    pub fn replica_count(&self, key: &[u8; 32]) -> usize {
        self.replicas.get(key).map(|i| i.holders.len()).unwrap_or(0)
    }

    /// Whether a key is fully replicated.
    pub fn is_fully_replicated(&self, key: &[u8; 32]) -> bool {
        self.replica_count(key) >= self.config.factor
    }

    /// Get all keys that are under-replicated.
    pub fn under_replicated_keys(&self) -> Vec<[u8; 32]> {
        self.replicas
            .iter()
            .filter(|(_, info)| info.holders.len() < self.config.factor)
            .map(|(k, _)| *k)
            .collect()
    }

    /// Get the set of peers holding replicas of a key.
    pub fn holders(&self, key: &[u8; 32]) -> HashSet<PeerId> {
        self.replicas
            .get(key)
            .map(|i| i.holders.clone())
            .unwrap_or_default()
    }

    /// Aggregate statistics.
    pub fn stats(&self) -> ReplicationStats {
        let total_keys = self.replicas.len();
        let fully_replicated = self
            .replicas
            .values()
            .filter(|i| i.holders.len() >= self.config.factor)
            .count();
        let total_replicas: usize = self.replicas.values().map(|i| i.holders.len()).sum();

        ReplicationStats {
            total_keys,
            fully_replicated,
            under_replicated: total_keys - fully_replicated,
            pending_tasks: self.pending.len(),
            total_replicas,
        }
    }

    /// Get the replication configuration
    pub fn config(&self) -> &ReplicationConfig {
        &self.config
    }
}

// ── Tests ───────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::NodeConfig;

    fn make_peer(id_byte: u8) -> PeerInfo {
        PeerInfo::new(
            PeerId::from_bytes([id_byte; 32]),
            vec![format!("/ip4/127.0.0.1/tcp/{}", 9000 + id_byte as u16)],
            [id_byte; 32],
            [id_byte; 32],
        )
    }

    #[test]
    fn test_schedule_replication() {
        let mut mgr = ReplicationManager::with_factor(3);
        let key = [1u8; 32];
        let value = Bytes::from_static(b"hello world");
        let peers: Vec<PeerInfo> = (1..=5).map(make_peer).collect();

        let tasks = mgr.schedule_replication(key, value, &peers);
        assert_eq!(tasks.len(), 3);
        assert!(tasks.iter().all(|t| t.priority == ReplicationPriority::High));
    }

    #[test]
    fn test_confirm_replication() {
        let mut mgr = ReplicationManager::with_factor(3);
        let key = [2u8; 32];
        let value = Bytes::from_static(b"data");
        let peers: Vec<PeerInfo> = (1..=5).map(make_peer).collect();

        let tasks = mgr.schedule_replication(key, value, &peers);
        assert_eq!(mgr.replica_count(&key), 0);

        for task in &tasks {
            mgr.confirm_replication(key, task.target_peer);
        }
        assert_eq!(mgr.replica_count(&key), 3);
        assert!(mgr.is_fully_replicated(&key));
    }

    #[test]
    fn test_handle_peer_lost() {
        let mut mgr = ReplicationManager::with_factor(3);
        let local_id = PeerId::random();
        let dht = DHTNode::new(local_id, NodeConfig::default());
        let key = [3u8; 32];
        let value = Bytes::from_static(b"important data");
        let peers: Vec<PeerInfo> = (1..=5).map(make_peer).collect();

        // Add peers to DHT routing table
        for p in &peers {
            let _ = dht.add_peer(p.clone());
        }

        // Store in DHT
        dht.store_value(key, value.clone(), [0u8; 32]).unwrap();

        // Schedule & confirm 3 replicas
        let tasks = mgr.schedule_replication(key, value, &peers);
        for t in &tasks {
            mgr.confirm_replication(key, t.target_peer);
        }
        assert_eq!(mgr.replica_count(&key), 3);

        // Lose one peer → should produce 1 re-replication task
        let lost = tasks[0].target_peer;
        let re_tasks = mgr.handle_peer_lost(&lost, &dht);
        // After losing one peer, we need 1 more replica
        assert!(!re_tasks.is_empty());
        assert_eq!(mgr.replica_count(&key), 2); // lost peer removed
    }

    #[test]
    fn test_under_replicated_keys() {
        let mut mgr = ReplicationManager::with_factor(3);
        let key = [4u8; 32];
        let value = Bytes::from_static(b"partial");
        let peers: Vec<PeerInfo> = (1..=5).map(make_peer).collect();

        mgr.schedule_replication(key, value, &peers);
        // Nothing confirmed yet → under-replicated
        let under = mgr.under_replicated_keys();
        assert!(under.contains(&key));
    }

    #[test]
    fn test_stats() {
        let mut mgr = ReplicationManager::with_factor(2);
        let peers: Vec<PeerInfo> = (1..=5).map(make_peer).collect();

        // Key A: fully replicated
        let key_a = [10u8; 32];
        let tasks_a = mgr.schedule_replication(key_a, Bytes::from_static(b"aaa"), &peers);
        for t in &tasks_a {
            mgr.confirm_replication(key_a, t.target_peer);
        }

        // Key B: partially replicated
        let key_b = [11u8; 32];
        let tasks_b = mgr.schedule_replication(key_b, Bytes::from_static(b"bbb"), &peers);
        mgr.confirm_replication(key_b, tasks_b[0].target_peer);

        let s = mgr.stats();
        assert_eq!(s.total_keys, 2);
        assert_eq!(s.fully_replicated, 1);
        assert_eq!(s.under_replicated, 1);
        assert_eq!(s.total_replicas, 3); // 2 for A + 1 for B
    }

    #[test]
    fn test_take_pending_tasks() {
        let mut mgr = ReplicationManager::with_factor(2);
        let peers: Vec<PeerInfo> = (1..=3).map(make_peer).collect();

        mgr.schedule_replication([20u8; 32], Bytes::from_static(b"x"), &peers);
        assert!(mgr.pending_count() > 0);

        let taken = mgr.take_pending_tasks();
        assert!(!taken.is_empty());
        assert_eq!(mgr.pending_count(), 0);
    }
}