shadow_dht/

node.rs

//! DHT node implementation combining routing and storage

use shadow_core::{PeerId, PeerInfo};
use shadow_core::error::{Result, ShadowError};
use crate::routing::RoutingTable;
use crate::store::{DHTStore, StoredValue};
use bytes::Bytes;
use tokio::time::{interval, Duration};
use std::sync::{Arc, RwLock};

/// DHT node configuration
#[derive(Debug, Clone)]
pub struct NodeConfig {
    /// K-bucket size
    pub k: usize,
    /// Alpha parameter for parallel queries
    pub alpha: usize,
    /// Storage size limit
    pub max_storage: usize,
    /// Value TTL (seconds)
    pub value_ttl: u64,
    /// Republish interval (seconds)
    pub republish_interval: u64,
}

impl Default for NodeConfig {
    fn default() -> Self {
        Self {
            k: 20,
            alpha: 3,
            max_storage: 100 * 1024 * 1024, // 100 MB
            value_ttl: 3600, // 1 hour
            republish_interval: 3600,
        }
    }
}

/// DHT node
pub struct DHTNode {
    /// Our peer ID
    local_id: PeerId,
    /// Routing table
    routing: Arc<RwLock<RoutingTable>>,
    /// Local storage
    store: Arc<RwLock<DHTStore>>,
    /// Configuration
    config: NodeConfig,
}

impl DHTNode {
    /// Create new DHT node
    pub fn new(local_id: PeerId, config: NodeConfig) -> Self {
        let routing = Arc::new(RwLock::new(RoutingTable::new(local_id, config.k)));
        let store = Arc::new(RwLock::new(DHTStore::new(config.max_storage)));

        Self {
            local_id,
            routing,
            store,
            config,
        }
    }

    /// Get our peer ID
    pub fn local_id(&self) -> PeerId {
        self.local_id
    }

    /// Add peer to routing table
    pub fn add_peer(&self, peer: PeerInfo) -> Result<bool> {
        self.routing.write().unwrap().add_peer(peer)
    }

    /// Remove peer from routing table
    pub fn remove_peer(&self, peer_id: &PeerId) -> Result<bool> {
        self.routing.write().unwrap().remove_peer(peer_id)
    }

    /// Find closest peers to target
    pub fn find_closest_peers(&self, target: &PeerId, count: usize) -> Vec<PeerInfo> {
        self.routing.read().unwrap().find_closest(target, count)
    }

    /// Store value locally
    pub fn store_value(&self, key: [u8; 32], data: Bytes, publisher: [u8; 32]) -> Result<()> {
        let value = StoredValue::new(data, publisher, self.config.value_ttl);
        self.store.write().unwrap().put(key, value)
    }

    /// Get value from local storage
    pub fn get_value(&self, key: &[u8; 32]) -> Option<StoredValue> {
        self.store.read().unwrap().get(key).cloned()
    }

    /// Check if we have a value
    pub fn has_value(&self, key: &[u8; 32]) -> bool {
        self.store.read().unwrap().contains(key)
    }

    /// Remove value from storage
    pub fn remove_value(&self, key: &[u8; 32]) -> Option<StoredValue> {
        self.store.write().unwrap().remove(key)
    }

    /// Get all known peers
    pub fn all_peers(&self) -> Vec<PeerInfo> {
        self.routing.read().unwrap().all_peers()
    }

    /// Get peer count
    pub fn peer_count(&self) -> usize {
        self.routing.read().unwrap().peer_count()
    }

    /// Get stored value count
    pub fn value_count(&self) -> usize {
        self.store.read().unwrap().len()
    }

    /// Get storage size
    pub fn storage_size(&self) -> usize {
        self.store.read().unwrap().size()
    }

    /// Mark peer as failed
    pub fn mark_peer_failed(&self, peer_id: &PeerId) -> Result<()> {
        self.routing.write().unwrap().mark_failed(peer_id)
    }

    /// Cleanup expired values
    pub fn cleanup_expired(&self) -> usize {
        self.store.write().unwrap().cleanup_expired()
    }

    /// Start background maintenance tasks
    pub async fn start_maintenance(self: Arc<Self>) {
        // Spawn cleanup task
        let node_clone = self.clone();
        tokio::spawn(async move {
            let mut cleanup_interval = interval(Duration::from_secs(300)); // 5 minutes
            
            loop {
                cleanup_interval.tick().await;
                let cleaned = node_clone.cleanup_expired();
                tracing::debug!("Cleaned {} expired values", cleaned);
            }
        });

        // Spawn republish task
        let node_clone = self.clone();
        tokio::spawn(async move {
            let mut republish_interval = interval(Duration::from_secs(node_clone.config.republish_interval));
            
            loop {
                republish_interval.tick().await;
                node_clone.republish_values();
            }
        });
    }

    /// Republish values that are expiring soon
    fn republish_values(&self) {
        let threshold = Duration::from_secs(self.config.republish_interval / 2);
        let expiring = self.store.read().unwrap().get_expiring_soon(threshold);
        
        tracing::debug!("Found {} values expiring soon", expiring.len());
        
        // In a real implementation, we would:
        // 1. Find K closest nodes to each key
        // 2. Send STORE RPCs to those nodes
        // For now, just log
    }

    /// Get DHT statistics
    pub fn stats(&self) -> DHTStats {
        DHTStats {
            peer_count: self.peer_count(),
            value_count: self.value_count(),
            storage_size: self.storage_size(),
            max_storage: self.config.max_storage,
        }
    }
}

/// DHT statistics
#[derive(Debug, Clone)]
pub struct DHTStats {
    pub peer_count: usize,
    pub value_count: usize,
    pub storage_size: usize,
    pub max_storage: usize,
}

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

    #[test]
    fn test_dht_node_creation() {
        let local_id = PeerId::random();
        let config = NodeConfig::default();
        let node = DHTNode::new(local_id, config);

        assert_eq!(node.local_id(), local_id);
        assert_eq!(node.peer_count(), 0);
        assert_eq!(node.value_count(), 0);
    }

    #[test]
    fn test_add_peer() {
        let local_id = PeerId::random();
        let node = DHTNode::new(local_id, NodeConfig::default());

        let peer = PeerInfo::new(
            PeerId::random(),
            vec!["127.0.0.1:9000".to_string()],
            [0u8; 32],
            [0u8; 32],
        );

        node.add_peer(peer).unwrap();
        assert_eq!(node.peer_count(), 1);
    }

    #[test]
    fn test_store_retrieve() {
        let local_id = PeerId::random();
        let node = DHTNode::new(local_id, NodeConfig::default());

        let key = [1u8; 32];
        let data = Bytes::from("test data");
        let publisher = [2u8; 32];

        node.store_value(key, data.clone(), publisher).unwrap();

        let retrieved = node.get_value(&key).unwrap();
        assert_eq!(retrieved.data, data);
        assert_eq!(retrieved.publisher, publisher);
    }

    #[test]
    fn test_find_closest() {
        let local_id = PeerId::random();
        let node = DHTNode::new(local_id, NodeConfig::default());

        // Add multiple peers
        for _ in 0..10 {
            let peer = PeerInfo::new(PeerId::random(), vec![], [0u8; 32], [0u8; 32]);
            node.add_peer(peer).unwrap();
        }

        let target = PeerId::random();
        let closest = node.find_closest_peers(&target, 3);

        assert_eq!(closest.len(), 3);
    }

    #[test]
    fn test_stats() {
        let local_id = PeerId::random();
        let node = DHTNode::new(local_id, NodeConfig::default());

        let stats = node.stats();
        assert_eq!(stats.peer_count, 0);
        assert_eq!(stats.value_count, 0);
    }
}