use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::ipc::GlobalPid;
use crate::process::Pid;
pub struct ConsistentHashRing {
ring: Vec<(u64, String)>,
virtual_nodes: usize,
}
impl ConsistentHashRing {
pub fn new(virtual_nodes: usize) -> Self {
Self {
ring: Vec::new(),
virtual_nodes,
}
}
pub fn add_node(&mut self, node_id: &str) {
for i in 0..self.virtual_nodes {
let key = format!("{node_id}:{i}");
let hash = Self::hash(&key);
self.ring.push((hash, node_id.to_string()));
}
self.ring.sort_by_key(|(h, _)| *h);
}
pub fn remove_node(&mut self, node_id: &str) {
self.ring.retain(|(_, id)| id != node_id);
}
pub fn assign(&self, pid: u64) -> Option<&str> {
if self.ring.is_empty() {
return None;
}
let hash = Self::hash(&pid.to_string());
let pos = self.ring.partition_point(|(h, _)| *h < hash);
let idx = if pos >= self.ring.len() { 0 } else { pos };
Some(&self.ring[idx].1)
}
pub fn node_count(&self) -> usize {
let mut nodes: std::collections::HashSet<&str> = std::collections::HashSet::new();
for (_, id) in &self.ring {
nodes.insert(id);
}
nodes.len()
}
fn hash(key: &str) -> u64 {
use std::hash::{Hash, Hasher};
let mut hasher = std::collections::hash_map::DefaultHasher::new();
key.hash(&mut hasher);
hasher.finish()
}
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ConsensusRole {
Follower,
Candidate,
Leader,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ConsensusEntry {
pub term: u64,
pub index: u64,
pub operation: ConsensusOp,
}
#[non_exhaustive]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ConsensusOp {
RegisterService { name: String, node_id: String },
DeregisterService { name: String, node_id: String },
RegisterProcess { global_pid: String, agent_type: String },
DeregisterProcess { global_pid: String },
}
pub struct MetadataConsensus {
role: ConsensusRole,
term: u64,
commit_index: u64,
log: Vec<ConsensusEntry>,
}
impl MetadataConsensus {
pub fn new() -> Self {
Self {
role: ConsensusRole::Follower,
term: 0,
commit_index: 0,
log: Vec::new(),
}
}
pub fn role(&self) -> ConsensusRole {
self.role
}
pub fn term(&self) -> u64 {
self.term
}
pub fn commit_index(&self) -> u64 {
self.commit_index
}
pub fn set_role(&mut self, role: ConsensusRole) {
self.role = role;
}
pub fn advance_term(&mut self, new_term: u64) {
if new_term > self.term {
self.term = new_term;
self.role = ConsensusRole::Follower;
}
}
pub fn append(&mut self, operation: ConsensusOp) -> u64 {
let index = self.log.len() as u64 + 1;
self.log.push(ConsensusEntry {
term: self.term,
index,
operation,
});
index
}
pub fn commit(&mut self, up_to: u64) -> Vec<&ConsensusEntry> {
let old = self.commit_index;
self.commit_index = up_to.min(self.log.len() as u64);
self.log
.iter()
.filter(|e| e.index > old && e.index <= self.commit_index)
.collect()
}
pub fn entries(&self) -> &[ConsensusEntry] {
&self.log
}
pub fn log_len(&self) -> usize {
self.log.len()
}
}
impl Default for MetadataConsensus {
fn default() -> Self {
Self::new()
}
}
pub struct CrdtGossipState {
entries: HashMap<String, (serde_json::Value, u64)>,
clock: u64,
}
impl CrdtGossipState {
pub fn new() -> Self {
Self {
entries: HashMap::new(),
clock: 0,
}
}
pub fn set(&mut self, key: String, value: serde_json::Value) {
self.clock += 1;
self.entries.insert(key, (value, self.clock));
}
pub fn get(&self, key: &str) -> Option<&serde_json::Value> {
self.entries.get(key).map(|(v, _)| v)
}
pub fn merge(&mut self, remote_entries: &[(String, serde_json::Value, u64)]) {
for (key, value, remote_clock) in remote_entries {
match self.entries.get(key) {
Some((_, local_clock)) if local_clock >= remote_clock => {
}
_ => {
self.entries
.insert(key.clone(), (value.clone(), *remote_clock));
if *remote_clock > self.clock {
self.clock = *remote_clock;
}
}
}
}
}
pub fn delta_since(&self, since_clock: u64) -> Vec<(String, serde_json::Value, u64)> {
self.entries
.iter()
.filter(|(_, (_, clock))| *clock > since_clock)
.map(|(k, (v, c))| (k.clone(), v.clone(), *c))
.collect()
}
pub fn len(&self) -> usize {
self.entries.len()
}
pub fn is_empty(&self) -> bool {
self.entries.is_empty()
}
pub fn clock(&self) -> u64 {
self.clock
}
}
impl Default for CrdtGossipState {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProcessAdvertisement {
pub global_pid: GlobalPid,
pub agent_type: String,
pub capabilities: Vec<String>,
pub services: Vec<String>,
pub status: ProcessStatus,
pub last_updated: u64,
pub resource_summary: Option<ResourceSummary>,
}
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum ProcessStatus {
Running,
Suspended,
Stopping,
Unreachable,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceSummary {
pub memory_bytes: u64,
pub cpu_time_us: u64,
pub inbox_depth: u32,
}
pub struct DistributedProcessTable {
processes: HashMap<String, ProcessAdvertisement>,
}
impl DistributedProcessTable {
pub fn new() -> Self {
Self {
processes: HashMap::new(),
}
}
pub fn merge(&mut self, advert: ProcessAdvertisement) {
let key = advert.global_pid.to_string();
match self.processes.get(&key) {
Some(existing) if existing.last_updated >= advert.last_updated => {
}
_ => {
self.processes.insert(key, advert);
}
}
}
pub fn locate(&self, pid: Pid, node_id: &str) -> Option<&ProcessAdvertisement> {
let key = format!("{node_id}:{pid}");
self.processes.get(&key)
}
pub fn find_by_type(&self, agent_type: &str) -> Vec<&ProcessAdvertisement> {
self.processes
.values()
.filter(|p| p.agent_type == agent_type && p.status == ProcessStatus::Running)
.collect()
}
pub fn find_by_capability(&self, capability: &str) -> Vec<&ProcessAdvertisement> {
self.processes
.values()
.filter(|p| {
p.capabilities.contains(&capability.to_string())
&& p.status == ProcessStatus::Running
})
.collect()
}
pub fn least_loaded_node(&self) -> Option<String> {
let mut node_loads: HashMap<String, u32> = HashMap::new();
for advert in self.processes.values() {
if advert.status == ProcessStatus::Running {
*node_loads
.entry(advert.global_pid.node_id.clone())
.or_default() += 1;
}
}
node_loads
.into_iter()
.min_by_key(|(_, count)| *count)
.map(|(node, _)| node)
}
pub fn remove_node(&mut self, node_id: &str) {
self.processes
.retain(|_, p| p.global_pid.node_id != node_id);
}
pub fn mark_node_unreachable(&mut self, node_id: &str) {
for advert in self.processes.values_mut() {
if advert.global_pid.node_id == node_id {
advert.status = ProcessStatus::Unreachable;
}
}
}
pub fn all_advertisements(&self) -> Vec<&ProcessAdvertisement> {
self.processes.values().collect()
}
pub fn len(&self) -> usize {
self.processes.len()
}
pub fn is_empty(&self) -> bool {
self.processes.is_empty()
}
pub fn assign_pid(&self, pid: u64, ring: &ConsistentHashRing) -> Option<String> {
ring.assign(pid).map(|s| s.to_string())
}
}
impl Default for DistributedProcessTable {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
fn make_advert(node: &str, pid: Pid, agent_type: &str, ts: u64) -> ProcessAdvertisement {
ProcessAdvertisement {
global_pid: GlobalPid::local(pid, node),
agent_type: agent_type.to_string(),
capabilities: vec!["code".to_string()],
services: vec![],
status: ProcessStatus::Running,
last_updated: ts,
resource_summary: None,
}
}
#[test]
fn process_advertisement_serde_roundtrip() {
let advert = make_advert("node-a", 1, "coder", 100);
let json = serde_json::to_string(&advert).unwrap();
let restored: ProcessAdvertisement = serde_json::from_str(&json).unwrap();
assert_eq!(restored.global_pid.node_id, "node-a");
assert_eq!(restored.global_pid.pid, 1);
assert_eq!(restored.agent_type, "coder");
assert_eq!(restored.last_updated, 100);
}
#[test]
fn merge_lww_newer_wins() {
let mut table = DistributedProcessTable::new();
let old = make_advert("node-a", 1, "coder", 100);
let new = ProcessAdvertisement {
agent_type: "reviewer".to_string(),
last_updated: 200,
..make_advert("node-a", 1, "coder", 200)
};
table.merge(old);
assert_eq!(table.len(), 1);
table.merge(new);
let found = table.locate(1, "node-a").unwrap();
assert_eq!(found.agent_type, "reviewer");
assert_eq!(found.last_updated, 200);
}
#[test]
fn merge_lww_older_ignored() {
let mut table = DistributedProcessTable::new();
let newer = make_advert("node-a", 1, "reviewer", 200);
let older = make_advert("node-a", 1, "coder", 100);
table.merge(newer);
table.merge(older);
let found = table.locate(1, "node-a").unwrap();
assert_eq!(found.agent_type, "reviewer");
}
#[test]
fn find_by_type_returns_matching() {
let mut table = DistributedProcessTable::new();
table.merge(make_advert("node-a", 1, "coder", 100));
table.merge(make_advert("node-b", 2, "reviewer", 100));
table.merge(make_advert("node-c", 3, "coder", 100));
let coders = table.find_by_type("coder");
assert_eq!(coders.len(), 2);
for c in &coders {
assert_eq!(c.agent_type, "coder");
}
}
#[test]
fn find_by_capability() {
let mut table = DistributedProcessTable::new();
let mut advert = make_advert("node-a", 1, "coder", 100);
advert.capabilities = vec!["code".to_string(), "review".to_string()];
table.merge(advert);
let mut advert2 = make_advert("node-b", 2, "tester", 100);
advert2.capabilities = vec!["test".to_string()];
table.merge(advert2);
let reviewers = table.find_by_capability("review");
assert_eq!(reviewers.len(), 1);
assert_eq!(reviewers[0].global_pid.node_id, "node-a");
}
#[test]
fn least_loaded_node() {
let mut table = DistributedProcessTable::new();
table.merge(make_advert("node-a", 1, "coder", 100));
table.merge(make_advert("node-a", 2, "coder", 100));
table.merge(make_advert("node-a", 3, "coder", 100));
table.merge(make_advert("node-b", 1, "coder", 100));
let node = table.least_loaded_node().unwrap();
assert_eq!(node, "node-b");
}
#[test]
fn remove_node() {
let mut table = DistributedProcessTable::new();
table.merge(make_advert("node-a", 1, "coder", 100));
table.merge(make_advert("node-a", 2, "coder", 100));
table.merge(make_advert("node-b", 1, "coder", 100));
table.remove_node("node-a");
assert_eq!(table.len(), 1);
assert!(table.locate(1, "node-a").is_none());
assert!(table.locate(1, "node-b").is_some());
}
#[test]
fn mark_node_unreachable() {
let mut table = DistributedProcessTable::new();
table.merge(make_advert("node-a", 1, "coder", 100));
table.merge(make_advert("node-a", 2, "coder", 100));
table.merge(make_advert("node-b", 1, "coder", 100));
table.mark_node_unreachable("node-a");
let a1 = table.locate(1, "node-a").unwrap();
assert_eq!(a1.status, ProcessStatus::Unreachable);
let a2 = table.locate(2, "node-a").unwrap();
assert_eq!(a2.status, ProcessStatus::Unreachable);
let b1 = table.locate(1, "node-b").unwrap();
assert_eq!(b1.status, ProcessStatus::Running);
}
#[test]
fn consistent_hash_ring_add_remove() {
let mut ring = ConsistentHashRing::new(64);
assert_eq!(ring.node_count(), 0);
ring.add_node("node-a");
assert_eq!(ring.node_count(), 1);
ring.add_node("node-b");
assert_eq!(ring.node_count(), 2);
ring.remove_node("node-a");
assert_eq!(ring.node_count(), 1);
assert_eq!(ring.assign(42), Some("node-b"));
}
#[test]
fn consistent_hash_ring_assignment_stable() {
let mut ring = ConsistentHashRing::new(64);
ring.add_node("node-a");
ring.add_node("node-b");
ring.add_node("node-c");
let first = ring.assign(12345).unwrap().to_string();
for _ in 0..100 {
assert_eq!(ring.assign(12345).unwrap(), first);
}
}
#[test]
fn consistent_hash_ring_redistribution() {
let mut ring = ConsistentHashRing::new(64);
ring.add_node("node-a");
ring.add_node("node-b");
ring.add_node("node-c");
let before: Vec<String> = (0..1000)
.map(|pid| ring.assign(pid).unwrap().to_string())
.collect();
ring.remove_node("node-b");
let after: Vec<String> = (0..1000)
.map(|pid| ring.assign(pid).unwrap().to_string())
.collect();
let mut moved = 0;
for i in 0..1000 {
if before[i] != after[i] {
moved += 1;
assert_eq!(before[i], "node-b");
}
}
assert!(moved > 0, "at least some PIDs should be redistributed");
assert!(moved < 1000, "not all PIDs should move");
}
#[test]
fn consistent_hash_ring_empty_returns_none() {
let ring = ConsistentHashRing::new(64);
assert!(ring.assign(1).is_none());
}
#[test]
fn pid_to_node_assignment() {
let mut ring = ConsistentHashRing::new(64);
ring.add_node("node-a");
ring.add_node("node-b");
let table = DistributedProcessTable::new();
let node = table.assign_pid(42, &ring);
assert!(node.is_some());
let n = node.unwrap();
assert!(n == "node-a" || n == "node-b");
}
#[test]
fn consensus_append_commit_cycle() {
let mut mc = MetadataConsensus::new();
assert_eq!(mc.role(), ConsensusRole::Follower);
assert_eq!(mc.term(), 0);
assert_eq!(mc.commit_index(), 0);
mc.set_role(ConsensusRole::Leader);
assert_eq!(mc.role(), ConsensusRole::Leader);
let idx1 = mc.append(ConsensusOp::RegisterService {
name: "cache".into(),
node_id: "node-a".into(),
});
assert_eq!(idx1, 1);
let idx2 = mc.append(ConsensusOp::RegisterProcess {
global_pid: "node-a:1".into(),
agent_type: "coder".into(),
});
assert_eq!(idx2, 2);
assert_eq!(mc.log_len(), 2);
let committed = mc.commit(2);
assert_eq!(committed.len(), 2);
assert_eq!(mc.commit_index(), 2);
let committed2 = mc.commit(2);
assert_eq!(committed2.len(), 0);
}
#[test]
fn consensus_advance_term() {
let mut mc = MetadataConsensus::new();
mc.set_role(ConsensusRole::Leader);
mc.advance_term(5);
assert_eq!(mc.term(), 5);
assert_eq!(mc.role(), ConsensusRole::Follower);
mc.advance_term(3);
assert_eq!(mc.term(), 5);
}
#[test]
fn consensus_commit_clamps_to_log_len() {
let mut mc = MetadataConsensus::new();
mc.append(ConsensusOp::DeregisterProcess {
global_pid: "node-a:1".into(),
});
mc.commit(999);
assert_eq!(mc.commit_index(), 1);
}
#[test]
fn crdt_set_get() {
let mut state = CrdtGossipState::new();
assert!(state.is_empty());
state.set("key1".into(), serde_json::json!("value1"));
assert_eq!(state.len(), 1);
assert_eq!(state.get("key1"), Some(&serde_json::json!("value1")));
assert!(state.get("missing").is_none());
}
#[test]
fn crdt_merge_lww() {
let mut local = CrdtGossipState::new();
local.set("key1".into(), serde_json::json!("local_v1"));
let remote = vec![("key1".into(), serde_json::json!("remote_v2"), 5u64)];
local.merge(&remote);
assert_eq!(local.get("key1"), Some(&serde_json::json!("remote_v2")));
assert_eq!(local.clock(), 5);
let old_remote = vec![("key1".into(), serde_json::json!("old_value"), 2u64)];
local.merge(&old_remote);
assert_eq!(local.get("key1"), Some(&serde_json::json!("remote_v2")));
}
#[test]
fn crdt_delta_since() {
let mut state = CrdtGossipState::new();
state.set("a".into(), serde_json::json!(1));
state.set("b".into(), serde_json::json!(2));
state.set("c".into(), serde_json::json!(3));
let delta = state.delta_since(1);
assert_eq!(delta.len(), 2);
let all = state.delta_since(0);
assert_eq!(all.len(), 3);
let none = state.delta_since(state.clock());
assert_eq!(none.len(), 0);
}
#[test]
fn crdt_concurrent_updates_converge() {
let mut node_a = CrdtGossipState::new();
let mut node_b = CrdtGossipState::new();
node_a.set("key1".into(), serde_json::json!("A"));
node_b.set("warmup".into(), serde_json::json!("x"));
node_b.set("key1".into(), serde_json::json!("B"));
let delta_a = node_a.delta_since(0);
let delta_b = node_b.delta_since(0);
node_a.merge(&delta_b);
node_b.merge(&delta_a);
assert_eq!(node_a.get("key1"), Some(&serde_json::json!("B")));
assert_eq!(node_b.get("key1"), Some(&serde_json::json!("B")));
}
}