use std::collections::HashSet;
use parking_lot::Mutex;
#[derive(Default)]
pub struct PartitionController {
isolated: Mutex<HashSet<u64>>,
}
impl PartitionController {
#[must_use]
pub fn new() -> Self {
Self::default()
}
pub fn isolate(&self, node_id: u64) {
self.isolated.lock().insert(node_id);
}
pub fn restore(&self, node_id: u64) {
self.isolated.lock().remove(&node_id);
}
pub fn heal(&self) {
self.isolated.lock().clear();
}
#[must_use]
pub fn is_blocked(&self, from: u64, to: u64) -> bool {
let guard = self.isolated.lock();
guard.contains(&from) || guard.contains(&to)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn no_partition_allows_all_traffic() {
let partition = PartitionController::new();
assert!(!partition.is_blocked(1, 2));
assert!(!partition.is_blocked(2, 1));
}
#[test]
fn isolate_blocks_outbound_and_inbound_for_the_node() {
let partition = PartitionController::new();
partition.isolate(2);
assert!(partition.is_blocked(2, 1));
assert!(partition.is_blocked(2, 3));
assert!(partition.is_blocked(1, 2));
assert!(partition.is_blocked(3, 2));
assert!(!partition.is_blocked(1, 3));
}
#[test]
fn restore_clears_a_single_node() {
let partition = PartitionController::new();
partition.isolate(2);
partition.isolate(3);
partition.restore(2);
assert!(!partition.is_blocked(2, 1));
assert!(!partition.is_blocked(1, 2));
assert!(partition.is_blocked(3, 1));
}
#[test]
fn heal_clears_all_isolation() {
let partition = PartitionController::new();
partition.isolate(2);
partition.isolate(3);
partition.heal();
assert!(!partition.is_blocked(2, 1));
assert!(!partition.is_blocked(3, 1));
}
}