use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use tracing::{debug, instrument};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RoutingTable {
node_to_cell: HashMap<String, (String, u64)>,
cell_to_zone: HashMap<String, (String, u64)>,
cell_leaders: HashMap<String, (String, u64)>,
}
impl RoutingTable {
pub fn new() -> Self {
Self {
node_to_cell: HashMap::new(),
cell_to_zone: HashMap::new(),
cell_leaders: HashMap::new(),
}
}
#[instrument(skip(self))]
pub fn assign_node(&mut self, node_id: &str, cell_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.node_to_cell.get(node_id) {
if timestamp <= *existing_ts {
debug!(
"Rejected node assignment (old timestamp): {} → {} @ {}",
node_id, cell_id, timestamp
);
return false;
}
}
debug!(
"Assigning node to cell: {} → {} @ {}",
node_id, cell_id, timestamp
);
self.node_to_cell
.insert(node_id.to_string(), (cell_id.to_string(), timestamp));
true
}
#[instrument(skip(self))]
pub fn remove_node(&mut self, node_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.node_to_cell.get(node_id) {
if timestamp < *existing_ts {
debug!(
"Rejected node removal (old timestamp): {} @ {}",
node_id, timestamp
);
return false;
}
}
debug!("Removing node from cell: {} @ {}", node_id, timestamp);
self.node_to_cell.remove(node_id);
true
}
#[instrument(skip(self))]
pub fn assign_cell(&mut self, cell_id: &str, zone_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.cell_to_zone.get(cell_id) {
if timestamp <= *existing_ts {
debug!(
"Rejected cell assignment (old timestamp): {} → {} @ {}",
cell_id, zone_id, timestamp
);
return false;
}
}
debug!(
"Assigning cell to zone: {} → {} @ {}",
cell_id, zone_id, timestamp
);
self.cell_to_zone
.insert(cell_id.to_string(), (zone_id.to_string(), timestamp));
true
}
#[instrument(skip(self))]
pub fn remove_cell(&mut self, cell_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.cell_to_zone.get(cell_id) {
if timestamp < *existing_ts {
debug!(
"Rejected cell removal (old timestamp): {} @ {}",
cell_id, timestamp
);
return false;
}
}
debug!("Removing cell from zone: {} @ {}", cell_id, timestamp);
self.cell_to_zone.remove(cell_id);
true
}
#[instrument(skip(self))]
pub fn set_cell_leader(&mut self, cell_id: &str, leader_node_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.cell_leaders.get(cell_id) {
if timestamp <= *existing_ts {
debug!(
"Rejected leader assignment (old timestamp): {} → {} @ {}",
cell_id, leader_node_id, timestamp
);
return false;
}
}
debug!(
"Setting cell leader: {} → {} @ {}",
cell_id, leader_node_id, timestamp
);
self.cell_leaders
.insert(cell_id.to_string(), (leader_node_id.to_string(), timestamp));
true
}
#[instrument(skip(self))]
pub fn remove_cell_leader(&mut self, cell_id: &str, timestamp: u64) -> bool {
if let Some((_, existing_ts)) = self.cell_leaders.get(cell_id) {
if timestamp < *existing_ts {
debug!(
"Rejected leader removal (old timestamp): {} @ {}",
cell_id, timestamp
);
return false;
}
}
debug!("Removing cell leader: {} @ {}", cell_id, timestamp);
self.cell_leaders.remove(cell_id);
true
}
pub fn get_node_cell(&self, node_id: &str) -> Option<&str> {
self.node_to_cell
.get(node_id)
.map(|(cell_id, _)| cell_id.as_str())
}
pub fn get_cell_zone(&self, cell_id: &str) -> Option<&str> {
self.cell_to_zone
.get(cell_id)
.map(|(zone_id, _)| zone_id.as_str())
}
pub fn get_node_zone(&self, node_id: &str) -> Option<&str> {
let cell_id = self.get_node_cell(node_id)?;
self.get_cell_zone(cell_id)
}
pub fn is_cell_leader(&self, node_id: &str, cell_id: &str) -> bool {
self.cell_leaders
.get(cell_id)
.map(|(leader_id, _)| leader_id == node_id)
.unwrap_or(false)
}
pub fn get_cell_leader(&self, cell_id: &str) -> Option<&str> {
self.cell_leaders
.get(cell_id)
.map(|(leader_id, _)| leader_id.as_str())
}
#[instrument(skip(self, other))]
pub fn merge(&mut self, other: &RoutingTable) {
debug!("Merging routing tables");
for (node_id, (cell_id, timestamp)) in &other.node_to_cell {
self.assign_node(node_id, cell_id, *timestamp);
}
for (cell_id, (zone_id, timestamp)) in &other.cell_to_zone {
self.assign_cell(cell_id, zone_id, *timestamp);
}
for (cell_id, (leader_id, timestamp)) in &other.cell_leaders {
self.set_cell_leader(cell_id, leader_id, *timestamp);
}
}
pub fn stats(&self) -> RoutingTableStats {
RoutingTableStats {
node_count: self.node_to_cell.len(),
cell_count: self.cell_to_zone.len(),
leader_count: self.cell_leaders.len(),
}
}
pub fn get_cell_nodes(&self, cell_id: &str) -> Vec<&str> {
self.node_to_cell
.iter()
.filter(|(_, (cid, _))| cid == cell_id)
.map(|(node_id, _)| node_id.as_str())
.collect()
}
pub fn get_zone_cells(&self, zone_id: &str) -> Vec<&str> {
self.cell_to_zone
.iter()
.filter(|(_, (zid, _))| zid == zone_id)
.map(|(cell_id, _)| cell_id.as_str())
.collect()
}
#[instrument(skip(self))]
pub fn merge_cells(
&mut self,
source_cell_ids: &[&str],
merged_cell_id: &str,
zone_id: Option<&str>,
timestamp: u64,
) -> usize {
debug!(
"Merging cells {:?} into {} @ {}",
source_cell_ids, merged_cell_id, timestamp
);
let mut reassigned_count = 0;
let nodes_to_reassign: Vec<String> = source_cell_ids
.iter()
.flat_map(|cell_id| {
self.node_to_cell
.iter()
.filter(|(_, (cid, _))| cid == *cell_id)
.map(|(node_id, _)| node_id.clone())
.collect::<Vec<_>>()
})
.collect();
for node_id in nodes_to_reassign {
if self.assign_node(&node_id, merged_cell_id, timestamp) {
reassigned_count += 1;
}
}
if let Some(zid) = zone_id {
self.assign_cell(merged_cell_id, zid, timestamp);
}
for cell_id in source_cell_ids {
self.remove_cell(cell_id, timestamp);
self.remove_cell_leader(cell_id, timestamp);
}
debug!(
"Merge complete: {} nodes reassigned to {}",
reassigned_count, merged_cell_id
);
reassigned_count
}
#[instrument(skip(self, nodes_for_a, nodes_for_b))]
#[allow(clippy::too_many_arguments)]
pub fn split_cell(
&mut self,
source_cell_id: &str,
cell_a_id: &str,
cell_b_id: &str,
nodes_for_a: &[&str],
nodes_for_b: &[&str],
zone_id: Option<&str>,
timestamp: u64,
) -> (usize, usize) {
debug!(
"Splitting cell {} into {} and {} @ {}",
source_cell_id, cell_a_id, cell_b_id, timestamp
);
let mut count_a = 0;
let mut count_b = 0;
for node_id in nodes_for_a {
if self.assign_node(node_id, cell_a_id, timestamp) {
count_a += 1;
}
}
for node_id in nodes_for_b {
if self.assign_node(node_id, cell_b_id, timestamp) {
count_b += 1;
}
}
if let Some(zid) = zone_id {
self.assign_cell(cell_a_id, zid, timestamp);
self.assign_cell(cell_b_id, zid, timestamp);
}
self.remove_cell(source_cell_id, timestamp);
self.remove_cell_leader(source_cell_id, timestamp);
debug!(
"Split complete: {} → {} nodes in {}, {} nodes in {}",
source_cell_id, count_a, cell_a_id, count_b, cell_b_id
);
(count_a, count_b)
}
}
impl Default for RoutingTable {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone)]
pub struct RoutingTableStats {
pub node_count: usize,
pub cell_count: usize,
pub leader_count: usize,
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_routing_table_creation() {
let table = RoutingTable::new();
let stats = table.stats();
assert_eq!(stats.node_count, 0);
assert_eq!(stats.cell_count, 0);
assert_eq!(stats.leader_count, 0);
}
#[test]
fn test_node_assignment() {
let mut table = RoutingTable::new();
assert!(table.assign_node("node1", "cell_alpha", 100));
assert_eq!(table.get_node_cell("node1"), Some("cell_alpha"));
assert!(table.assign_node("node2", "cell_alpha", 101));
assert_eq!(table.get_node_cell("node2"), Some("cell_alpha"));
let stats = table.stats();
assert_eq!(stats.node_count, 2);
}
#[test]
fn test_lww_semantics_node_assignment() {
let mut table = RoutingTable::new();
assert!(table.assign_node("node1", "cell_alpha", 100));
assert_eq!(table.get_node_cell("node1"), Some("cell_alpha"));
assert!(table.assign_node("node1", "cell_beta", 200));
assert_eq!(table.get_node_cell("node1"), Some("cell_beta"));
assert!(!table.assign_node("node1", "cell_gamma", 150));
assert_eq!(table.get_node_cell("node1"), Some("cell_beta"));
assert!(!table.assign_node("node1", "cell_delta", 200));
assert_eq!(table.get_node_cell("node1"), Some("cell_beta"));
}
#[test]
fn test_cell_assignment() {
let mut table = RoutingTable::new();
assert!(table.assign_cell("cell_alpha", "zone_north", 100));
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_north"));
let stats = table.stats();
assert_eq!(stats.cell_count, 1);
}
#[test]
fn test_lww_semantics_cell_assignment() {
let mut table = RoutingTable::new();
assert!(table.assign_cell("cell_alpha", "zone_north", 100));
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_north"));
assert!(table.assign_cell("cell_alpha", "zone_south", 200));
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_south"));
assert!(!table.assign_cell("cell_alpha", "zone_east", 150));
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_south"));
}
#[test]
fn test_cell_leader() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_alpha", 101);
assert!(table.set_cell_leader("cell_alpha", "node1", 102));
assert!(table.is_cell_leader("node1", "cell_alpha"));
assert!(!table.is_cell_leader("node2", "cell_alpha"));
assert_eq!(table.get_cell_leader("cell_alpha"), Some("node1"));
let stats = table.stats();
assert_eq!(stats.leader_count, 1);
}
#[test]
fn test_lww_semantics_leader() {
let mut table = RoutingTable::new();
assert!(table.set_cell_leader("cell_alpha", "node1", 100));
assert!(table.is_cell_leader("node1", "cell_alpha"));
assert!(table.set_cell_leader("cell_alpha", "node2", 200));
assert!(!table.is_cell_leader("node1", "cell_alpha"));
assert!(table.is_cell_leader("node2", "cell_alpha"));
assert!(!table.set_cell_leader("cell_alpha", "node3", 150));
assert!(table.is_cell_leader("node2", "cell_alpha"));
}
#[test]
fn test_node_zone_lookup() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_cell("cell_alpha", "zone_north", 101);
assert_eq!(table.get_node_zone("node1"), Some("zone_north"));
assert_eq!(table.get_node_zone("node2"), None);
table.assign_node("node3", "cell_beta", 102);
assert_eq!(table.get_node_zone("node3"), None);
}
#[test]
fn test_get_cell_nodes() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_alpha", 101);
table.assign_node("node3", "cell_beta", 102);
let mut alpha_nodes = table.get_cell_nodes("cell_alpha");
alpha_nodes.sort();
assert_eq!(alpha_nodes, vec!["node1", "node2"]);
let beta_nodes = table.get_cell_nodes("cell_beta");
assert_eq!(beta_nodes, vec!["node3"]);
let empty_nodes = table.get_cell_nodes("cell_gamma");
assert_eq!(empty_nodes.len(), 0);
}
#[test]
fn test_get_zone_cells() {
let mut table = RoutingTable::new();
table.assign_cell("cell_alpha", "zone_north", 100);
table.assign_cell("cell_beta", "zone_north", 101);
table.assign_cell("cell_gamma", "zone_south", 102);
let mut north_cells = table.get_zone_cells("zone_north");
north_cells.sort();
assert_eq!(north_cells, vec!["cell_alpha", "cell_beta"]);
let south_cells = table.get_zone_cells("zone_south");
assert_eq!(south_cells, vec!["cell_gamma"]);
let empty_cells = table.get_zone_cells("zone_east");
assert_eq!(empty_cells.len(), 0);
}
#[test]
fn test_node_removal() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
assert_eq!(table.get_node_cell("node1"), Some("cell_alpha"));
assert!(!table.remove_node("node1", 50));
assert_eq!(table.get_node_cell("node1"), Some("cell_alpha"));
assert!(table.remove_node("node1", 200));
assert_eq!(table.get_node_cell("node1"), None);
}
#[test]
fn test_cell_removal() {
let mut table = RoutingTable::new();
table.assign_cell("cell_alpha", "zone_north", 100);
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_north"));
assert!(!table.remove_cell("cell_alpha", 50));
assert_eq!(table.get_cell_zone("cell_alpha"), Some("zone_north"));
assert!(table.remove_cell("cell_alpha", 200));
assert_eq!(table.get_cell_zone("cell_alpha"), None);
}
#[test]
fn test_leader_removal() {
let mut table = RoutingTable::new();
table.set_cell_leader("cell_alpha", "node1", 100);
assert_eq!(table.get_cell_leader("cell_alpha"), Some("node1"));
assert!(!table.remove_cell_leader("cell_alpha", 50));
assert_eq!(table.get_cell_leader("cell_alpha"), Some("node1"));
assert!(table.remove_cell_leader("cell_alpha", 200));
assert_eq!(table.get_cell_leader("cell_alpha"), None);
}
#[test]
fn test_merge() {
let mut table1 = RoutingTable::new();
let mut table2 = RoutingTable::new();
table1.assign_node("node1", "cell_alpha", 100);
table1.assign_cell("cell_alpha", "zone_north", 100);
table2.assign_node("node1", "cell_beta", 200);
table2.assign_node("node2", "cell_gamma", 150);
table2.assign_cell("cell_beta", "zone_south", 200);
table1.merge(&table2);
assert_eq!(table1.get_node_cell("node1"), Some("cell_beta"));
assert_eq!(table1.get_node_cell("node2"), Some("cell_gamma"));
assert_eq!(table1.get_cell_zone("cell_beta"), Some("zone_south"));
assert_eq!(table1.get_cell_zone("cell_alpha"), Some("zone_north"));
}
#[test]
fn test_merge_leaders() {
let mut table1 = RoutingTable::new();
let mut table2 = RoutingTable::new();
table1.set_cell_leader("cell_alpha", "node1", 100);
table2.set_cell_leader("cell_alpha", "node2", 200);
table1.merge(&table2);
assert!(table1.is_cell_leader("node2", "cell_alpha"));
assert!(!table1.is_cell_leader("node1", "cell_alpha"));
}
#[test]
fn test_merge_cells() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_alpha", 101);
table.assign_node("node3", "cell_beta", 102);
table.assign_node("node4", "cell_beta", 103);
table.assign_cell("cell_alpha", "zone_north", 100);
table.assign_cell("cell_beta", "zone_north", 101);
table.set_cell_leader("cell_alpha", "node1", 100);
table.set_cell_leader("cell_beta", "node3", 101);
let count = table.merge_cells(
&["cell_alpha", "cell_beta"],
"cell_merged",
Some("zone_north"),
200,
);
assert_eq!(count, 4);
assert_eq!(table.get_node_cell("node1"), Some("cell_merged"));
assert_eq!(table.get_node_cell("node2"), Some("cell_merged"));
assert_eq!(table.get_node_cell("node3"), Some("cell_merged"));
assert_eq!(table.get_node_cell("node4"), Some("cell_merged"));
assert_eq!(table.get_cell_zone("cell_merged"), Some("zone_north"));
assert_eq!(table.get_cell_zone("cell_alpha"), None);
assert_eq!(table.get_cell_zone("cell_beta"), None);
assert_eq!(table.get_cell_leader("cell_alpha"), None);
assert_eq!(table.get_cell_leader("cell_beta"), None);
assert_eq!(table.get_cell_leader("cell_merged"), None);
}
#[test]
fn test_split_cell() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_alpha", 101);
table.assign_node("node3", "cell_alpha", 102);
table.assign_node("node4", "cell_alpha", 103);
table.assign_cell("cell_alpha", "zone_north", 100);
table.set_cell_leader("cell_alpha", "node1", 100);
let (count_a, count_b) = table.split_cell(
"cell_alpha",
"cell_a",
"cell_b",
&["node1", "node2"],
&["node3", "node4"],
Some("zone_north"),
200,
);
assert_eq!(count_a, 2);
assert_eq!(count_b, 2);
assert_eq!(table.get_node_cell("node1"), Some("cell_a"));
assert_eq!(table.get_node_cell("node2"), Some("cell_a"));
assert_eq!(table.get_node_cell("node3"), Some("cell_b"));
assert_eq!(table.get_node_cell("node4"), Some("cell_b"));
assert_eq!(table.get_cell_zone("cell_a"), Some("zone_north"));
assert_eq!(table.get_cell_zone("cell_b"), Some("zone_north"));
assert_eq!(table.get_cell_zone("cell_alpha"), None);
assert_eq!(table.get_cell_leader("cell_alpha"), None);
assert_eq!(table.get_cell_leader("cell_a"), None);
assert_eq!(table.get_cell_leader("cell_b"), None);
}
#[test]
fn test_merge_cells_without_zone() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_beta", 101);
let count = table.merge_cells(&["cell_alpha", "cell_beta"], "cell_merged", None, 200);
assert_eq!(count, 2);
assert_eq!(table.get_node_cell("node1"), Some("cell_merged"));
assert_eq!(table.get_node_cell("node2"), Some("cell_merged"));
assert_eq!(table.get_cell_zone("cell_merged"), None);
}
#[test]
fn test_split_cell_without_zone() {
let mut table = RoutingTable::new();
table.assign_node("node1", "cell_alpha", 100);
table.assign_node("node2", "cell_alpha", 101);
let (count_a, count_b) = table.split_cell(
"cell_alpha",
"cell_a",
"cell_b",
&["node1"],
&["node2"],
None,
200,
);
assert_eq!(count_a, 1);
assert_eq!(count_b, 1);
assert_eq!(table.get_cell_zone("cell_a"), None);
assert_eq!(table.get_cell_zone("cell_b"), None);
}
}