ipfrs_transport/

partition.rs

//! Network partition detection and handling
//!
//! Provides mechanisms to:
//! - Detect network partitions
//! - Queue requests during partitions
//! - Automatically recover when partition heals
//! - Monitor peer health and connectivity

use dashmap::DashMap;
use parking_lot::RwLock;
use std::collections::VecDeque;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};
use thiserror::Error;
use tokio::sync::watch;
use tracing::{debug, info, warn};

/// Partition detection error
#[derive(Error, Debug)]
pub enum PartitionError {
    #[error("Network partition detected")]
    PartitionDetected,

    #[error("Peer unreachable: {0}")]
    PeerUnreachable(String),

    #[error("Queue full: cannot accept more requests")]
    QueueFull,

    #[error("Recovery timeout")]
    RecoveryTimeout,
}

/// Network partition state
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PartitionState {
    /// Network is healthy
    #[default]
    Healthy,
    /// Partition suspected (some failures)
    Suspected,
    /// Partition confirmed
    Partitioned,
    /// Recovering from partition
    Recovering,
}

/// Partition detection configuration
#[derive(Debug, Clone)]
pub struct PartitionConfig {
    /// Number of consecutive failures to trigger suspicion
    pub failure_threshold: usize,
    /// Time window for failure counting
    pub failure_window: Duration,
    /// Probe interval when partition suspected
    pub probe_interval: Duration,
    /// Maximum queued requests during partition
    pub max_queued_requests: usize,
    /// Recovery probe count before declaring healthy
    pub recovery_probe_count: usize,
    /// Peer timeout duration
    pub peer_timeout: Duration,
}

impl Default for PartitionConfig {
    fn default() -> Self {
        Self {
            failure_threshold: 3,
            failure_window: Duration::from_secs(10),
            probe_interval: Duration::from_secs(5),
            max_queued_requests: 1000,
            recovery_probe_count: 3,
            peer_timeout: Duration::from_secs(30),
        }
    }
}

/// Partition statistics
#[derive(Debug, Clone, Default)]
pub struct PartitionStats {
    /// Total partitions detected
    pub partitions_detected: u64,
    /// Total recoveries
    pub recoveries: u64,
    /// Current queued requests
    pub queued_requests: usize,
    /// Requests dropped due to queue full
    pub dropped_requests: u64,
    /// Average partition duration
    pub avg_partition_duration: Option<Duration>,
    /// Current state
    pub state: PartitionState,
}

/// Queued request during partition
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct QueuedRequest {
    peer: SocketAddr,
    data: Vec<u8>,
    queued_at: Instant,
}

/// Peer health information
#[derive(Debug, Clone)]
struct PeerHealth {
    /// Last successful contact
    last_success: Option<Instant>,
    /// Last failure
    last_failure: Option<Instant>,
    /// Recent failure count
    failure_count: usize,
    /// Failure timestamps in window
    failures: VecDeque<Instant>,
}

impl PeerHealth {
    fn new() -> Self {
        Self {
            last_success: None,
            last_failure: None,
            failure_count: 0,
            failures: VecDeque::new(),
        }
    }

    /// Record a successful contact
    fn record_success(&mut self) {
        self.last_success = Some(Instant::now());
        self.failure_count = 0;
        self.failures.clear();
    }

    /// Record a failure
    fn record_failure(&mut self, window: Duration) {
        let now = Instant::now();
        self.last_failure = Some(now);
        self.failures.push_back(now);

        // Clean old failures outside window
        while let Some(&first) = self.failures.front() {
            if now.duration_since(first) > window {
                self.failures.pop_front();
            } else {
                break;
            }
        }

        self.failure_count = self.failures.len();
    }

    /// Check if peer is unhealthy
    fn is_unhealthy(&self, threshold: usize, timeout: Duration) -> bool {
        // Too many recent failures
        if self.failure_count >= threshold {
            return true;
        }

        // No recent success and timeout exceeded
        if let Some(last_success) = self.last_success {
            if last_success.elapsed() > timeout {
                return true;
            }
        } else if let Some(last_failure) = self.last_failure {
            if last_failure.elapsed() > timeout {
                return true;
            }
        }

        false
    }
}

/// Network partition detector and handler
pub struct PartitionDetector {
    config: PartitionConfig,
    state: Arc<RwLock<PartitionState>>,
    peer_health: Arc<DashMap<SocketAddr, PeerHealth>>,
    queued_requests: Arc<RwLock<VecDeque<QueuedRequest>>>,
    stats: Arc<RwLock<PartitionStats>>,
    state_tx: watch::Sender<PartitionState>,
    state_rx: watch::Receiver<PartitionState>,
    partition_start: Arc<RwLock<Option<Instant>>>,
}

impl PartitionDetector {
    /// Create a new partition detector
    pub fn new(config: PartitionConfig) -> Self {
        let (state_tx, state_rx) = watch::channel(PartitionState::Healthy);

        Self {
            config,
            state: Arc::new(RwLock::new(PartitionState::Healthy)),
            peer_health: Arc::new(DashMap::new()),
            queued_requests: Arc::new(RwLock::new(VecDeque::new())),
            stats: Arc::new(RwLock::new(PartitionStats::default())),
            state_tx,
            state_rx,
            partition_start: Arc::new(RwLock::new(None)),
        }
    }

    /// Record a successful peer interaction
    pub fn record_success(&self, peer: &SocketAddr) {
        {
            let mut health = self
                .peer_health
                .entry(*peer)
                .or_insert_with(PeerHealth::new);
            health.record_success();
        } // Release DashMap entry guard here

        // Check if we should transition to healthy
        if *self.state.read() != PartitionState::Healthy {
            self.check_recovery();
        }
    }

    /// Record a peer failure
    pub fn record_failure(&self, peer: &SocketAddr) {
        {
            let mut health = self
                .peer_health
                .entry(*peer)
                .or_insert_with(PeerHealth::new);
            health.record_failure(self.config.failure_window);

            debug!("Peer {} failure count: {}", peer, health.failure_count);
        } // Release DashMap entry guard here

        // Check if we should transition to partitioned state
        self.check_partition();
    }

    /// Check if network partition should be declared
    fn check_partition(&self) {
        let unhealthy_count = self
            .peer_health
            .iter()
            .filter(|entry| {
                entry
                    .value()
                    .is_unhealthy(self.config.failure_threshold, self.config.peer_timeout)
            })
            .count();

        let total_peers = self.peer_health.len();

        // If majority of peers are unhealthy, declare partition
        if total_peers > 0 && unhealthy_count * 2 > total_peers {
            let current_state = *self.state.read();

            if current_state == PartitionState::Healthy {
                self.transition_to_suspected();
            } else if current_state == PartitionState::Suspected {
                self.transition_to_partitioned();
            }
        }
    }

    /// Check if network has recovered
    fn check_recovery(&self) {
        let healthy_count = self
            .peer_health
            .iter()
            .filter(|entry| {
                !entry
                    .value()
                    .is_unhealthy(self.config.failure_threshold, self.config.peer_timeout)
            })
            .count();

        let total_peers = self.peer_health.len();

        // If majority of peers are healthy, start recovery
        if total_peers > 0 && healthy_count * 2 > total_peers {
            let current_state = *self.state.read();

            if current_state == PartitionState::Partitioned {
                self.transition_to_recovering();
            } else if current_state == PartitionState::Recovering {
                self.transition_to_healthy();
            }
        }
    }

    /// Transition to suspected state
    fn transition_to_suspected(&self) {
        *self.state.write() = PartitionState::Suspected;
        let _ = self.state_tx.send(PartitionState::Suspected);
        warn!("Network partition suspected");
    }

    /// Transition to partitioned state
    fn transition_to_partitioned(&self) {
        *self.state.write() = PartitionState::Partitioned;
        *self.partition_start.write() = Some(Instant::now());
        let _ = self.state_tx.send(PartitionState::Partitioned);

        let mut stats = self.stats.write();
        stats.partitions_detected += 1;
        stats.state = PartitionState::Partitioned;

        warn!("Network partition detected - queueing requests");
    }

    /// Transition to recovering state
    fn transition_to_recovering(&self) {
        *self.state.write() = PartitionState::Recovering;
        let _ = self.state_tx.send(PartitionState::Recovering);
        info!("Network partition recovering");
    }

    /// Transition to healthy state
    fn transition_to_healthy(&self) {
        *self.state.write() = PartitionState::Healthy;
        let _ = self.state_tx.send(PartitionState::Healthy);

        // Update partition duration stats
        if let Some(start) = *self.partition_start.read() {
            let duration = start.elapsed();
            let mut stats = self.stats.write();

            stats.avg_partition_duration = Some(
                stats
                    .avg_partition_duration
                    .map(|avg| (avg + duration) / 2)
                    .unwrap_or(duration),
            );
            stats.recoveries += 1;
            stats.state = PartitionState::Healthy;
        }

        *self.partition_start.write() = None;

        info!("Network partition recovered - processing queued requests");

        // Process queued requests
        self.flush_queue();
    }

    /// Queue a request during partition
    pub fn queue_request(&self, peer: SocketAddr, data: Vec<u8>) -> Result<(), PartitionError> {
        let mut queue = self.queued_requests.write();

        if queue.len() >= self.config.max_queued_requests {
            self.stats.write().dropped_requests += 1;
            return Err(PartitionError::QueueFull);
        }

        queue.push_back(QueuedRequest {
            peer,
            data,
            queued_at: Instant::now(),
        });

        self.stats.write().queued_requests = queue.len();

        Ok(())
    }

    /// Flush queued requests
    fn flush_queue(&self) {
        let requests: Vec<_> = {
            let mut queue = self.queued_requests.write();
            queue.drain(..).collect()
        };

        info!("Flushing {} queued requests", requests.len());

        // Requests would be processed here by the caller
        self.stats.write().queued_requests = 0;
    }

    /// Get queued requests for processing
    pub fn drain_queue(&self) -> Vec<(SocketAddr, Vec<u8>)> {
        let requests: Vec<_> = {
            let mut queue = self.queued_requests.write();
            queue.drain(..).collect()
        };

        self.stats.write().queued_requests = 0;

        requests
            .into_iter()
            .map(|req| (req.peer, req.data))
            .collect()
    }

    /// Get current partition state
    pub fn state(&self) -> PartitionState {
        *self.state.read()
    }

    /// Get statistics
    pub fn stats(&self) -> PartitionStats {
        self.stats.read().clone()
    }

    /// Wait for state change
    pub async fn wait_state_change(&self) -> PartitionState {
        let mut rx = self.state_rx.clone();
        let _ = rx.changed().await;
        let state = *rx.borrow();
        state
    }

    /// Check if a specific peer is reachable
    pub fn is_peer_reachable(&self, peer: &SocketAddr) -> bool {
        if let Some(health) = self.peer_health.get(peer) {
            !health.is_unhealthy(self.config.failure_threshold, self.config.peer_timeout)
        } else {
            true // Unknown peers are assumed reachable
        }
    }

    /// Get unhealthy peers
    pub fn unhealthy_peers(&self) -> Vec<SocketAddr> {
        self.peer_health
            .iter()
            .filter(|entry| {
                entry
                    .value()
                    .is_unhealthy(self.config.failure_threshold, self.config.peer_timeout)
            })
            .map(|entry| *entry.key())
            .collect()
    }

    /// Clear peer health data
    pub fn clear_peer_health(&self) {
        self.peer_health.clear();
        info!("Cleared peer health data");
    }
}

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

    #[test]
    fn test_peer_health() {
        let mut health = PeerHealth::new();
        let window = Duration::from_secs(10);

        // Record some failures
        health.record_failure(window);
        assert_eq!(health.failure_count, 1);

        health.record_failure(window);
        assert_eq!(health.failure_count, 2);

        // Success should reset
        health.record_success();
        assert_eq!(health.failure_count, 0);
    }

    #[test]
    fn test_partition_detection() {
        let config = PartitionConfig {
            failure_threshold: 2,
            ..Default::default()
        };

        let detector = PartitionDetector::new(config);
        let peer: SocketAddr = "127.0.0.1:8080".parse().unwrap();

        assert_eq!(detector.state(), PartitionState::Healthy);

        // Record failures
        detector.record_failure(&peer);
        detector.record_failure(&peer);
        detector.record_failure(&peer);

        // Should transition to suspected or partitioned
        let state = detector.state();
        assert!(state == PartitionState::Suspected || state == PartitionState::Partitioned);
    }

    #[test]
    fn test_queue_request() {
        let detector = PartitionDetector::new(PartitionConfig::default());
        let peer: SocketAddr = "127.0.0.1:8080".parse().unwrap();

        let result = detector.queue_request(peer, vec![1, 2, 3]);
        assert!(result.is_ok());

        let stats = detector.stats();
        assert_eq!(stats.queued_requests, 1);
    }

    #[test]
    fn test_queue_full() {
        let config = PartitionConfig {
            max_queued_requests: 2,
            ..Default::default()
        };

        let detector = PartitionDetector::new(config);
        let peer: SocketAddr = "127.0.0.1:8080".parse().unwrap();

        detector.queue_request(peer, vec![1]).unwrap();
        detector.queue_request(peer, vec![2]).unwrap();

        // Third request should fail
        let result = detector.queue_request(peer, vec![3]);
        assert!(result.is_err());
    }

    #[test]
    fn test_drain_queue() {
        let detector = PartitionDetector::new(PartitionConfig::default());
        let peer: SocketAddr = "127.0.0.1:8080".parse().unwrap();

        detector.queue_request(peer, vec![1, 2, 3]).unwrap();
        detector.queue_request(peer, vec![4, 5, 6]).unwrap();

        let drained = detector.drain_queue();
        assert_eq!(drained.len(), 2);
        assert_eq!(detector.stats().queued_requests, 0);
    }

    #[tokio::test]
    async fn test_state_transitions() {
        let config = PartitionConfig {
            failure_threshold: 1,
            ..Default::default()
        };

        let detector = PartitionDetector::new(config);
        let peer: SocketAddr = "127.0.0.1:8080".parse().unwrap();

        // Start healthy
        assert_eq!(detector.state(), PartitionState::Healthy);

        // Cause partition
        detector.record_failure(&peer);

        // Should transition
        assert!(detector.state() != PartitionState::Healthy);

        // Recover
        detector.record_success(&peer);

        // May still be recovering or might still be in suspected/partitioned state
        // since we only have one peer and one success after one failure
        let state = detector.state();
        assert!(
            state == PartitionState::Recovering
                || state == PartitionState::Healthy
                || state == PartitionState::Suspected
                || state == PartitionState::Partitioned,
            "Expected one of the valid states, got: {:?}",
            state
        );
    }
}