photoncloud-monorepo/chainfire/chainfire-client/src/node.rs

//! Node metadata helpers for Chainfire KVS
//!
//! This module provides helpers for storing and retrieving node metadata
//! in the Chainfire distributed KVS.
//!
//! # KVS Key Schema
//!
//! Node metadata is stored with the following key structure:
//! - `/nodes/<id>/info` - JSON-encoded NodeMetadata
//! - `/nodes/<id>/roles` - JSON-encoded roles (raft_role, gossip_role)
//! - `/nodes/<id>/capacity/cpu` - CPU cores (u32)
//! - `/nodes/<id>/capacity/memory_gb` - Memory in GB (u32)
//! - `/nodes/<id>/labels/<key>` - Custom labels (string)
//! - `/nodes/<id>/api_addr` - API address (string)

use crate::error::Result;
use crate::Client;
use chainfire_types::node::NodeRole;
use chainfire_types::RaftRole;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// Node metadata stored in KVS
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeMetadata {
    /// Unique node identifier
    pub id: u64,
    /// Human-readable node name
    pub name: String,
    /// Raft participation role
    pub raft_role: RaftRole,
    /// Gossip/cluster role
    pub gossip_role: NodeRole,
    /// API address for client connections
    pub api_addr: String,
    /// Raft address for inter-node communication (optional for workers)
    pub raft_addr: Option<String>,
    /// Gossip address for membership protocol
    pub gossip_addr: String,
    /// Node capacity information
    pub capacity: NodeCapacity,
    /// Custom labels for node selection
    pub labels: HashMap<String, String>,
}

/// Node capacity information
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct NodeCapacity {
    /// Number of CPU cores
    pub cpu_cores: u32,
    /// Memory in gigabytes
    pub memory_gb: u32,
    /// Disk space in gigabytes (optional)
    pub disk_gb: Option<u32>,
}

/// Filter for listing nodes
#[derive(Debug, Clone, Default)]
pub struct NodeFilter {
    /// Filter by Raft role
    pub raft_role: Option<RaftRole>,
    /// Filter by gossip role
    pub gossip_role: Option<NodeRole>,
    /// Filter by labels (all must match)
    pub labels: HashMap<String, String>,
}

impl NodeMetadata {
    /// Create a new NodeMetadata for a control-plane node
    pub fn control_plane(
        id: u64,
        name: impl Into<String>,
        api_addr: impl Into<String>,
        raft_addr: impl Into<String>,
        gossip_addr: impl Into<String>,
    ) -> Self {
        Self {
            id,
            name: name.into(),
            raft_role: RaftRole::Voter,
            gossip_role: NodeRole::ControlPlane,
            api_addr: api_addr.into(),
            raft_addr: Some(raft_addr.into()),
            gossip_addr: gossip_addr.into(),
            capacity: NodeCapacity::default(),
            labels: HashMap::new(),
        }
    }

    /// Create a new NodeMetadata for a worker node
    pub fn worker(
        id: u64,
        name: impl Into<String>,
        api_addr: impl Into<String>,
        gossip_addr: impl Into<String>,
    ) -> Self {
        Self {
            id,
            name: name.into(),
            raft_role: RaftRole::None,
            gossip_role: NodeRole::Worker,
            api_addr: api_addr.into(),
            raft_addr: None,
            gossip_addr: gossip_addr.into(),
            capacity: NodeCapacity::default(),
            labels: HashMap::new(),
        }
    }

    /// Set capacity information
    pub fn with_capacity(mut self, cpu_cores: u32, memory_gb: u32) -> Self {
        self.capacity.cpu_cores = cpu_cores;
        self.capacity.memory_gb = memory_gb;
        self
    }

    /// Add a label
    pub fn with_label(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
        self.labels.insert(key.into(), value.into());
        self
    }
}

/// Key prefix for all node metadata
const NODE_PREFIX: &str = "/nodes/";

/// Generate the key for node info
fn node_info_key(id: u64) -> String {
    format!("{}{}/info", NODE_PREFIX, id)
}

/// Generate the key for a node label
fn node_label_key(id: u64, label: &str) -> String {
    format!("{}{}/labels/{}", NODE_PREFIX, id, label)
}

/// Register a node in the cluster by storing its metadata in KVS
///
/// # Arguments
///
/// * `client` - The Chainfire client
/// * `meta` - Node metadata to register
///
/// # Returns
///
/// The revision number of the write operation
pub async fn register_node(client: &mut Client, meta: &NodeMetadata) -> Result<u64> {
    let key = node_info_key(meta.id);
    let value = serde_json::to_string(meta)
        .map_err(|e| crate::error::ClientError::Internal(e.to_string()))?;

    client.put_str(&key, &value).await
}

/// Update a specific node attribute
pub async fn update_node_label(
    client: &mut Client,
    node_id: u64,
    label: &str,
    value: &str,
) -> Result<u64> {
    let key = node_label_key(node_id, label);
    client.put_str(&key, value).await
}

/// Get a node's metadata by ID
///
/// # Arguments
///
/// * `client` - The Chainfire client
/// * `node_id` - The node ID to look up
///
/// # Returns
///
/// The node metadata if found, None otherwise
pub async fn get_node(client: &mut Client, node_id: u64) -> Result<Option<NodeMetadata>> {
    let key = node_info_key(node_id);
    let value = client.get_str(&key).await?;

    match value {
        Some(json) => {
            let meta: NodeMetadata = serde_json::from_str(&json)
                .map_err(|e| crate::error::ClientError::Internal(e.to_string()))?;
            Ok(Some(meta))
        }
        None => Ok(None),
    }
}

/// List all registered nodes
///
/// # Arguments
///
/// * `client` - The Chainfire client
/// * `filter` - Optional filter criteria
///
/// # Returns
///
/// A list of node metadata matching the filter
pub async fn list_nodes(client: &mut Client, filter: &NodeFilter) -> Result<Vec<NodeMetadata>> {
    let prefix = NODE_PREFIX.to_string();
    let entries = client.get_prefix(&prefix).await?;

    let mut nodes = Vec::new();

    for (key, value) in entries {
        let key_str = String::from_utf8_lossy(&key);

        // Only process /nodes/<id>/info keys
        if !key_str.ends_with("/info") {
            continue;
        }

        let json = String::from_utf8_lossy(&value);
        if let Ok(meta) = serde_json::from_str::<NodeMetadata>(&json) {
            // Apply filters
            if let Some(ref raft_role) = filter.raft_role {
                if meta.raft_role != *raft_role {
                    continue;
                }
            }

            if let Some(ref gossip_role) = filter.gossip_role {
                if meta.gossip_role != *gossip_role {
                    continue;
                }
            }

            // Check label filters
            let mut labels_match = true;
            for (k, v) in &filter.labels {
                match meta.labels.get(k) {
                    Some(node_v) if node_v == v => {}
                    _ => {
                        labels_match = false;
                        break;
                    }
                }
            }

            if labels_match {
                nodes.push(meta);
            }
        }
    }

    // Sort by node ID for consistent ordering
    nodes.sort_by_key(|n| n.id);

    Ok(nodes)
}

/// Unregister a node from the cluster
///
/// # Arguments
///
/// * `client` - The Chainfire client
/// * `node_id` - The node ID to unregister
///
/// # Returns
///
/// True if the node was found and deleted
pub async fn unregister_node(client: &mut Client, node_id: u64) -> Result<bool> {
    let key = node_info_key(node_id);
    client.delete(&key).await
}

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

    #[test]
    fn test_node_info_key() {
        assert_eq!(node_info_key(1), "/nodes/1/info");
        assert_eq!(node_info_key(123), "/nodes/123/info");
    }

    #[test]
    fn test_node_label_key() {
        assert_eq!(node_label_key(1, "zone"), "/nodes/1/labels/zone");
    }

    #[test]
    fn test_control_plane_metadata() {
        let meta = NodeMetadata::control_plane(
            1,
            "cp-1",
            "127.0.0.1:2379",
            "127.0.0.1:2380",
            "127.0.0.1:2381",
        );

        assert_eq!(meta.id, 1);
        assert_eq!(meta.raft_role, RaftRole::Voter);
        assert_eq!(meta.gossip_role, NodeRole::ControlPlane);
        assert!(meta.raft_addr.is_some());
    }

    #[test]
    fn test_worker_metadata() {
        let meta = NodeMetadata::worker(100, "worker-1", "127.0.0.1:3379", "127.0.0.1:3381");

        assert_eq!(meta.id, 100);
        assert_eq!(meta.raft_role, RaftRole::None);
        assert_eq!(meta.gossip_role, NodeRole::Worker);
        assert!(meta.raft_addr.is_none());
    }

    #[test]
    fn test_metadata_with_capacity() {
        let meta = NodeMetadata::worker(1, "worker", "addr", "gossip")
            .with_capacity(8, 32)
            .with_label("zone", "us-west-1");

        assert_eq!(meta.capacity.cpu_cores, 8);
        assert_eq!(meta.capacity.memory_gb, 32);
        assert_eq!(meta.labels.get("zone"), Some(&"us-west-1".to_string()));
    }

    #[test]
    fn test_metadata_serialization() {
        let meta = NodeMetadata::control_plane(1, "test", "api", "raft", "gossip")
            .with_capacity(4, 16)
            .with_label("env", "prod");

        let json = serde_json::to_string(&meta).unwrap();
        let deserialized: NodeMetadata = serde_json::from_str(&json).unwrap();

        assert_eq!(meta.id, deserialized.id);
        assert_eq!(meta.raft_role, deserialized.raft_role);
        assert_eq!(meta.capacity.cpu_cores, deserialized.capacity.cpu_cores);
    }
}