use crate::beacon::{GeographicBeacon, HierarchyLevel};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
#[derive(Debug, Clone)]
pub struct PartitionConfig {
pub min_detection_attempts: u32,
pub initial_backoff: Duration,
pub max_backoff: Duration,
pub backoff_multiplier: f64,
pub min_higher_level_beacons: usize,
}
impl Default for PartitionConfig {
fn default() -> Self {
Self {
min_detection_attempts: 3,
initial_backoff: Duration::from_secs(2),
max_backoff: Duration::from_secs(30),
backoff_multiplier: 2.0,
min_higher_level_beacons: 1,
}
}
}
impl PartitionConfig {
pub fn calculate_backoff(&self, attempt: u32) -> Duration {
let multiplier = self.backoff_multiplier.powi(attempt as i32);
let backoff_secs = self.initial_backoff.as_secs_f64() * multiplier;
let max_secs = self.max_backoff.as_secs_f64();
Duration::from_secs_f64(backoff_secs.min(max_secs))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PartitionEvent {
PartitionDetected {
attempts: u32,
detection_duration: Duration,
},
PartitionHealed {
visible_beacons: usize,
partition_duration: Duration,
},
}
pub trait PartitionHandler: Send + Sync + std::fmt::Debug {
fn on_partition_detected(&self, event: &PartitionEvent);
fn on_partition_healed(&self, event: &PartitionEvent);
}
#[derive(Debug)]
pub struct PartitionDetector {
config: PartitionConfig,
current_level: HierarchyLevel,
handler: Option<Arc<dyn PartitionHandler>>,
partitioned: bool,
detection_attempts: u32,
last_detection_attempt: Option<Instant>,
partition_start: Option<Instant>,
}
impl PartitionDetector {
pub fn new(current_level: HierarchyLevel, config: PartitionConfig) -> Self {
Self {
config,
current_level,
handler: None,
partitioned: false,
detection_attempts: 0,
last_detection_attempt: None,
partition_start: None,
}
}
pub fn with_handler(mut self, handler: Arc<dyn PartitionHandler>) -> Self {
self.handler = Some(handler);
self
}
pub fn check_partition(
&mut self,
beacons: &HashMap<String, GeographicBeacon>,
) -> Option<PartitionEvent> {
let higher_level_count = self.count_higher_level_beacons(beacons);
let has_connectivity = higher_level_count >= self.config.min_higher_level_beacons;
if has_connectivity {
if self.partitioned {
return self.handle_partition_healed(higher_level_count);
} else {
self.reset_detection_state();
return None;
}
}
if !self.should_attempt_detection() {
return None;
}
self.detection_attempts += 1;
self.last_detection_attempt = Some(Instant::now());
if self.detection_attempts >= self.config.min_detection_attempts {
return self.handle_partition_detected();
}
None
}
fn count_higher_level_beacons(&self, beacons: &HashMap<String, GeographicBeacon>) -> usize {
beacons
.values()
.filter(|beacon| beacon.operational && beacon.hierarchy_level > self.current_level)
.count()
}
fn should_attempt_detection(&self) -> bool {
match self.last_detection_attempt {
None => true, Some(last_attempt) => {
let backoff = self.config.calculate_backoff(self.detection_attempts);
last_attempt.elapsed() >= backoff
}
}
}
fn handle_partition_detected(&mut self) -> Option<PartitionEvent> {
if self.partitioned {
return None; }
let partition_start = self.last_detection_attempt.unwrap_or_else(Instant::now);
let detection_duration = partition_start.elapsed();
let event = PartitionEvent::PartitionDetected {
attempts: self.detection_attempts,
detection_duration,
};
self.partitioned = true;
self.partition_start = Some(partition_start);
if let Some(ref handler) = self.handler {
handler.on_partition_detected(&event);
}
Some(event)
}
fn handle_partition_healed(&mut self, visible_beacons: usize) -> Option<PartitionEvent> {
if !self.partitioned {
return None; }
let partition_duration = self
.partition_start
.map(|start| start.elapsed())
.unwrap_or(Duration::ZERO);
let event = PartitionEvent::PartitionHealed {
visible_beacons,
partition_duration,
};
self.reset_detection_state();
if let Some(ref handler) = self.handler {
handler.on_partition_healed(&event);
}
Some(event)
}
fn reset_detection_state(&mut self) {
self.partitioned = false;
self.detection_attempts = 0;
self.last_detection_attempt = None;
self.partition_start = None;
}
pub fn is_partitioned(&self) -> bool {
self.partitioned
}
pub fn detection_attempts(&self) -> u32 {
self.detection_attempts
}
}
#[cfg(test)]
mod tests {
use super::*;
fn create_beacon(node_id: &str, level: HierarchyLevel, operational: bool) -> GeographicBeacon {
let position = crate::beacon::GeoPosition::new(37.7749, -122.4194);
let mut beacon = GeographicBeacon::new(node_id.to_string(), position, level);
beacon.operational = operational;
beacon
}
#[test]
fn test_partition_config_default() {
let config = PartitionConfig::default();
assert_eq!(config.min_detection_attempts, 3);
assert_eq!(config.initial_backoff, Duration::from_secs(2));
assert_eq!(config.max_backoff, Duration::from_secs(30));
assert_eq!(config.backoff_multiplier, 2.0);
assert_eq!(config.min_higher_level_beacons, 1);
}
#[test]
fn test_backoff_calculation() {
let config = PartitionConfig::default();
assert_eq!(config.calculate_backoff(0), Duration::from_secs(2));
assert_eq!(config.calculate_backoff(1), Duration::from_secs(4));
assert_eq!(config.calculate_backoff(2), Duration::from_secs(8));
assert_eq!(config.calculate_backoff(10), Duration::from_secs(30));
}
#[test]
fn test_count_higher_level_beacons() {
let detector = PartitionDetector::new(HierarchyLevel::Platform, PartitionConfig::default());
let mut beacons = HashMap::new();
beacons.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
beacons.insert(
"platoon1".to_string(),
create_beacon("platoon1", HierarchyLevel::Platoon, true),
);
beacons.insert(
"platform2".to_string(),
create_beacon("platform2", HierarchyLevel::Platform, true),
);
assert_eq!(detector.count_higher_level_beacons(&beacons), 2);
}
#[test]
fn test_count_excludes_non_operational_beacons() {
let detector = PartitionDetector::new(HierarchyLevel::Platform, PartitionConfig::default());
let mut beacons = HashMap::new();
beacons.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
beacons.insert(
"squad2".to_string(),
create_beacon("squad2", HierarchyLevel::Squad, false),
);
assert_eq!(detector.count_higher_level_beacons(&beacons), 1);
}
#[test]
fn test_no_partition_with_connectivity() {
let mut detector =
PartitionDetector::new(HierarchyLevel::Platform, PartitionConfig::default());
let mut beacons = HashMap::new();
beacons.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
let event = detector.check_partition(&beacons);
assert!(event.is_none());
assert!(!detector.is_partitioned());
assert_eq!(detector.detection_attempts(), 0);
}
#[test]
fn test_partition_detection_after_min_attempts() {
let config = PartitionConfig {
min_detection_attempts: 2,
initial_backoff: Duration::from_millis(1), ..Default::default()
};
let mut detector = PartitionDetector::new(HierarchyLevel::Platform, config);
let beacons = HashMap::new();
let event1 = detector.check_partition(&beacons);
assert!(event1.is_none());
assert_eq!(detector.detection_attempts(), 1);
std::thread::sleep(Duration::from_millis(2));
let event2 = detector.check_partition(&beacons);
assert!(event2.is_some());
assert!(detector.is_partitioned());
if let Some(PartitionEvent::PartitionDetected { attempts, .. }) = event2 {
assert_eq!(attempts, 2);
} else {
panic!("Expected PartitionDetected event");
}
}
#[test]
fn test_partition_healed_event() {
let config = PartitionConfig {
min_detection_attempts: 1,
initial_backoff: Duration::from_millis(1),
..Default::default()
};
let mut detector = PartitionDetector::new(HierarchyLevel::Platform, config);
let beacons_empty = HashMap::new();
let _ = detector.check_partition(&beacons_empty);
assert!(detector.is_partitioned());
let mut beacons_with_parent = HashMap::new();
beacons_with_parent.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
let event = detector.check_partition(&beacons_with_parent);
assert!(event.is_some());
if let Some(PartitionEvent::PartitionHealed {
visible_beacons, ..
}) = event
{
assert_eq!(visible_beacons, 1);
} else {
panic!("Expected PartitionHealed event");
}
assert!(!detector.is_partitioned());
assert_eq!(detector.detection_attempts(), 0);
}
#[test]
fn test_backoff_prevents_rapid_detection_attempts() {
let config = PartitionConfig {
min_detection_attempts: 3,
initial_backoff: Duration::from_secs(10), ..Default::default()
};
let mut detector = PartitionDetector::new(HierarchyLevel::Platform, config);
let beacons = HashMap::new();
let _ = detector.check_partition(&beacons);
assert_eq!(detector.detection_attempts(), 1);
let _ = detector.check_partition(&beacons);
assert_eq!(detector.detection_attempts(), 1);
assert_eq!(detector.detection_attempts(), 1);
}
#[test]
fn test_company_level_node_has_no_higher_levels() {
let detector = PartitionDetector::new(HierarchyLevel::Company, PartitionConfig::default());
let mut beacons = HashMap::new();
beacons.insert(
"platoon1".to_string(),
create_beacon("platoon1", HierarchyLevel::Platoon, true),
);
beacons.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
assert_eq!(detector.count_higher_level_beacons(&beacons), 0);
}
#[test]
fn test_min_higher_level_beacons_threshold() {
let config = PartitionConfig {
min_higher_level_beacons: 2, min_detection_attempts: 1,
initial_backoff: Duration::from_millis(1),
..Default::default()
};
let mut detector = PartitionDetector::new(HierarchyLevel::Platform, config);
let mut beacons = HashMap::new();
beacons.insert(
"squad1".to_string(),
create_beacon("squad1", HierarchyLevel::Squad, true),
);
let event = detector.check_partition(&beacons);
assert!(event.is_some());
}
}