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};
#[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,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum PartitionState {
#[default]
Healthy,
Suspected,
Partitioned,
Recovering,
}
#[derive(Debug, Clone)]
pub struct PartitionConfig {
pub failure_threshold: usize,
pub failure_window: Duration,
pub probe_interval: Duration,
pub max_queued_requests: usize,
pub recovery_probe_count: usize,
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),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct PartitionStats {
pub partitions_detected: u64,
pub recoveries: u64,
pub queued_requests: usize,
pub dropped_requests: u64,
pub avg_partition_duration: Option<Duration>,
pub state: PartitionState,
}
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct QueuedRequest {
peer: SocketAddr,
data: Vec<u8>,
queued_at: Instant,
}
#[derive(Debug, Clone)]
struct PeerHealth {
last_success: Option<Instant>,
last_failure: Option<Instant>,
failure_count: usize,
failures: VecDeque<Instant>,
}
impl PeerHealth {
fn new() -> Self {
Self {
last_success: None,
last_failure: None,
failure_count: 0,
failures: VecDeque::new(),
}
}
fn record_success(&mut self) {
self.last_success = Some(Instant::now());
self.failure_count = 0;
self.failures.clear();
}
fn record_failure(&mut self, window: Duration) {
let now = Instant::now();
self.last_failure = Some(now);
self.failures.push_back(now);
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();
}
fn is_unhealthy(&self, threshold: usize, timeout: Duration) -> bool {
if self.failure_count >= threshold {
return true;
}
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
}
}
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 {
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)),
}
}
pub fn record_success(&self, peer: &SocketAddr) {
{
let mut health = self
.peer_health
.entry(*peer)
.or_insert_with(PeerHealth::new);
health.record_success();
}
if *self.state.read() != PartitionState::Healthy {
self.check_recovery();
}
}
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);
}
self.check_partition();
}
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 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();
}
}
}
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 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();
}
}
}
fn transition_to_suspected(&self) {
*self.state.write() = PartitionState::Suspected;
let _ = self.state_tx.send(PartitionState::Suspected);
warn!("Network partition suspected");
}
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");
}
fn transition_to_recovering(&self) {
*self.state.write() = PartitionState::Recovering;
let _ = self.state_tx.send(PartitionState::Recovering);
info!("Network partition recovering");
}
fn transition_to_healthy(&self) {
*self.state.write() = PartitionState::Healthy;
let _ = self.state_tx.send(PartitionState::Healthy);
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");
self.flush_queue();
}
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(())
}
fn flush_queue(&self) {
let requests: Vec<_> = {
let mut queue = self.queued_requests.write();
queue.drain(..).collect()
};
info!("Flushing {} queued requests", requests.len());
self.stats.write().queued_requests = 0;
}
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()
}
pub fn state(&self) -> PartitionState {
*self.state.read()
}
pub fn stats(&self) -> PartitionStats {
self.stats.read().clone()
}
pub async fn wait_state_change(&self) -> PartitionState {
let mut rx = self.state_rx.clone();
let _ = rx.changed().await;
let state = *rx.borrow();
state
}
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 }
}
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()
}
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);
health.record_failure(window);
assert_eq!(health.failure_count, 1);
health.record_failure(window);
assert_eq!(health.failure_count, 2);
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().expect("test: valid socket addr");
assert_eq!(detector.state(), PartitionState::Healthy);
detector.record_failure(&peer);
detector.record_failure(&peer);
detector.record_failure(&peer);
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().expect("test: valid socket addr");
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().expect("test: valid socket addr");
detector
.queue_request(peer, vec![1])
.expect("test: queue request");
detector
.queue_request(peer, vec![2])
.expect("test: queue request");
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().expect("test: valid socket addr");
detector
.queue_request(peer, vec![1, 2, 3])
.expect("test: queue request");
detector
.queue_request(peer, vec![4, 5, 6])
.expect("test: queue request");
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().expect("test: valid socket addr");
assert_eq!(detector.state(), PartitionState::Healthy);
detector.record_failure(&peer);
assert!(detector.state() != PartitionState::Healthy);
detector.record_success(&peer);
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
);
}
}