photoncloud-monorepo/chainfire/crates/chainfire-gossip/src/membership.rs

//! Membership state management

use crate::identity::GossipId;
use chainfire_types::NodeId;
use dashmap::DashMap;
use std::net::SocketAddr;
use tracing::debug;

/// Membership change event
#[derive(Debug, Clone)]
pub enum MembershipChange {
    /// A member joined or became reachable
    MemberUp(GossipId),
    /// A member left or became unreachable
    MemberDown(GossipId),
}

/// Manages cluster membership state
pub struct MembershipState {
    /// Known members
    members: DashMap<NodeId, GossipId>,
}

impl MembershipState {
    /// Create a new membership state
    pub fn new() -> Self {
        Self {
            members: DashMap::new(),
        }
    }

    /// Handle a membership change
    pub fn handle_change(&self, change: MembershipChange) {
        match change {
            MembershipChange::MemberUp(id) => {
                debug!(node_id = id.node_id, addr = %id.addr, "Adding member");
                self.members.insert(id.node_id, id);
            }
            MembershipChange::MemberDown(id) => {
                debug!(node_id = id.node_id, "Removing member");
                self.members.remove(&id.node_id);
            }
        }
    }

    /// Get a member by node ID
    pub fn get(&self, node_id: NodeId) -> Option<GossipId> {
        self.members.get(&node_id).map(|r| r.clone())
    }

    /// Get all members
    pub fn all(&self) -> Vec<GossipId> {
        self.members.iter().map(|r| r.clone()).collect()
    }

    /// Get member count
    pub fn count(&self) -> usize {
        self.members.len()
    }

    /// Check if a node is a member
    pub fn contains(&self, node_id: NodeId) -> bool {
        self.members.contains_key(&node_id)
    }

    /// Get all member addresses
    pub fn addresses(&self) -> Vec<SocketAddr> {
        self.members.iter().map(|r| r.addr).collect()
    }

    /// Get all control plane members
    pub fn control_plane_members(&self) -> Vec<GossipId> {
        self.members
            .iter()
            .filter(|r| r.is_control_plane())
            .map(|r| r.clone())
            .collect()
    }

    /// Get all worker members
    pub fn worker_members(&self) -> Vec<GossipId> {
        self.members
            .iter()
            .filter(|r| r.is_worker())
            .map(|r| r.clone())
            .collect()
    }
}

impl Default for MembershipState {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use chainfire_types::node::NodeRole;

    fn create_id(node_id: NodeId, role: NodeRole) -> GossipId {
        GossipId::new(
            node_id,
            format!("127.0.0.1:{}", 5000 + node_id).parse().unwrap(),
            role,
        )
    }

    #[test]
    fn test_membership_changes() {
        let state = MembershipState::new();

        let id1 = create_id(1, NodeRole::ControlPlane);
        let id2 = create_id(2, NodeRole::Worker);

        state.handle_change(MembershipChange::MemberUp(id1.clone()));
        state.handle_change(MembershipChange::MemberUp(id2.clone()));

        assert_eq!(state.count(), 2);
        assert!(state.contains(1));
        assert!(state.contains(2));

        state.handle_change(MembershipChange::MemberDown(id1));
        assert_eq!(state.count(), 1);
        assert!(!state.contains(1));
    }

    #[test]
    fn test_role_filtering() {
        let state = MembershipState::new();

        state.handle_change(MembershipChange::MemberUp(create_id(1, NodeRole::ControlPlane)));
        state.handle_change(MembershipChange::MemberUp(create_id(2, NodeRole::ControlPlane)));
        state.handle_change(MembershipChange::MemberUp(create_id(3, NodeRole::Worker)));
        state.handle_change(MembershipChange::MemberUp(create_id(4, NodeRole::Worker)));
        state.handle_change(MembershipChange::MemberUp(create_id(5, NodeRole::Worker)));

        assert_eq!(state.control_plane_members().len(), 2);
        assert_eq!(state.worker_members().len(), 3);
    }
}