forge_runtime/cluster/

registry.rs

use std::net::IpAddr;
use std::time::Duration;

use sqlx::Row;
use uuid::Uuid;

use forge_core::cluster::{NodeId, NodeInfo, NodeRole, NodeStatus};
use forge_core::{ForgeError, Result};

/// Node registry for cluster membership.
pub struct NodeRegistry {
    pool: sqlx::PgPool,
    local_node: NodeInfo,
}

impl NodeRegistry {
    /// Create a new node registry.
    pub fn new(pool: sqlx::PgPool, local_node: NodeInfo) -> Self {
        Self { pool, local_node }
    }

    /// Get the local node info.
    pub fn local_node(&self) -> &NodeInfo {
        &self.local_node
    }

    /// Get the local node ID.
    pub fn local_id(&self) -> NodeId {
        self.local_node.id
    }

    /// Register the local node in the cluster.
    pub async fn register(&self) -> Result<()> {
        let roles: Vec<&str> = self.local_node.roles.iter().map(|r| r.as_str()).collect();

        sqlx::query(
            r#"
            INSERT INTO forge_nodes (
                id, hostname, ip_address, http_port, grpc_port,
                roles, worker_capabilities, status, version, started_at, last_heartbeat
            ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, NOW())
            ON CONFLICT (id) DO UPDATE SET
                hostname = EXCLUDED.hostname,
                ip_address = EXCLUDED.ip_address,
                http_port = EXCLUDED.http_port,
                grpc_port = EXCLUDED.grpc_port,
                roles = EXCLUDED.roles,
                worker_capabilities = EXCLUDED.worker_capabilities,
                status = EXCLUDED.status,
                version = EXCLUDED.version,
                last_heartbeat = NOW()
            "#,
        )
        .bind(self.local_node.id.as_uuid())
        .bind(&self.local_node.hostname)
        .bind(self.local_node.ip_address.to_string())
        .bind(self.local_node.http_port as i32)
        .bind(self.local_node.grpc_port as i32)
        .bind(&roles)
        .bind(&self.local_node.worker_capabilities)
        .bind(self.local_node.status.as_str())
        .bind(&self.local_node.version)
        .bind(self.local_node.started_at)
        .execute(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(())
    }

    /// Update node status.
    pub async fn set_status(&self, status: NodeStatus) -> Result<()> {
        sqlx::query(
            r#"
            UPDATE forge_nodes
            SET status = $2
            WHERE id = $1
            "#,
        )
        .bind(self.local_node.id.as_uuid())
        .bind(status.as_str())
        .execute(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(())
    }

    /// Deregister the local node from the cluster.
    pub async fn deregister(&self) -> Result<()> {
        sqlx::query(
            r#"
            DELETE FROM forge_nodes WHERE id = $1
            "#,
        )
        .bind(self.local_node.id.as_uuid())
        .execute(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(())
    }

    /// Get all active nodes in the cluster.
    pub async fn get_active_nodes(&self) -> Result<Vec<NodeInfo>> {
        self.get_nodes_by_status(NodeStatus::Active).await
    }

    /// Get nodes by status.
    pub async fn get_nodes_by_status(&self, status: NodeStatus) -> Result<Vec<NodeInfo>> {
        let rows = sqlx::query(
            r#"
            SELECT id, hostname, ip_address, http_port, grpc_port,
                   roles, worker_capabilities, status, version,
                   started_at, last_heartbeat, current_connections,
                   current_jobs, cpu_usage, memory_usage
            FROM forge_nodes
            WHERE status = $1
            ORDER BY started_at
            "#,
        )
        .bind(status.as_str())
        .fetch_all(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        rows.into_iter().map(|row| parse_node_row(&row)).collect()
    }

    /// Get a specific node by ID.
    pub async fn get_node(&self, node_id: NodeId) -> Result<Option<NodeInfo>> {
        let row = sqlx::query(
            r#"
            SELECT id, hostname, ip_address, http_port, grpc_port,
                   roles, worker_capabilities, status, version,
                   started_at, last_heartbeat, current_connections,
                   current_jobs, cpu_usage, memory_usage
            FROM forge_nodes
            WHERE id = $1
            "#,
        )
        .bind(node_id.as_uuid())
        .fetch_optional(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        row.map(|r| parse_node_row(&r)).transpose()
    }

    /// Count nodes by status.
    pub async fn count_by_status(&self) -> Result<NodeCounts> {
        let row = sqlx::query(
            r#"
            SELECT
                COUNT(*) FILTER (WHERE status = 'active') as active,
                COUNT(*) FILTER (WHERE status = 'draining') as draining,
                COUNT(*) FILTER (WHERE status = 'dead') as dead,
                COUNT(*) FILTER (WHERE status = 'joining') as joining,
                COUNT(*) as total
            FROM forge_nodes
            "#,
        )
        .fetch_one(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(NodeCounts {
            active: row.get::<i64, _>("active") as usize,
            draining: row.get::<i64, _>("draining") as usize,
            dead: row.get::<i64, _>("dead") as usize,
            joining: row.get::<i64, _>("joining") as usize,
            total: row.get::<i64, _>("total") as usize,
        })
    }

    /// Mark stale nodes as dead.
    pub async fn mark_dead_nodes(&self, threshold: Duration) -> Result<u64> {
        let threshold_secs = threshold.as_secs() as i64;

        let result = sqlx::query(
            r#"
            UPDATE forge_nodes
            SET status = 'dead'
            WHERE status = 'active'
              AND last_heartbeat < NOW() - make_interval(secs => $1)
            "#,
        )
        .bind(threshold_secs as f64)
        .execute(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(result.rows_affected())
    }

    /// Clean up old dead nodes.
    pub async fn cleanup_dead_nodes(&self, older_than: Duration) -> Result<u64> {
        let threshold_secs = older_than.as_secs() as i64;

        let result = sqlx::query(
            r#"
            DELETE FROM forge_nodes
            WHERE status = 'dead'
              AND last_heartbeat < NOW() - make_interval(secs => $1)
            "#,
        )
        .bind(threshold_secs as f64)
        .execute(&self.pool)
        .await
        .map_err(|e| ForgeError::Database(e.to_string()))?;

        Ok(result.rows_affected())
    }
}

fn parse_node_row(row: &sqlx::postgres::PgRow) -> Result<NodeInfo> {
    let id: Uuid = row.get("id");

    let ip_str: String = row.get("ip_address");
    let ip_address: IpAddr = ip_str
        .parse()
        .map_err(|e| ForgeError::Validation(format!("invalid IP address '{}': {}", ip_str, e)))?;

    let roles_str: Vec<String> = row.get("roles");
    let roles: Vec<NodeRole> = roles_str
        .iter()
        .map(|s| {
            s.parse()
                .map_err(|e| ForgeError::Validation(format!("invalid role '{}': {}", s, e)))
        })
        .collect::<Result<Vec<_>>>()?;

    let status_str: String = row.get("status");
    let status: NodeStatus = status_str
        .parse()
        .map_err(|e| ForgeError::Validation(format!("invalid status '{}': {}", status_str, e)))?;

    Ok(NodeInfo {
        id: NodeId::from_uuid(id),
        hostname: row.get("hostname"),
        ip_address,
        http_port: row.get::<i32, _>("http_port") as u16,
        grpc_port: row.get::<i32, _>("grpc_port") as u16,
        roles,
        worker_capabilities: row.get("worker_capabilities"),
        status,
        version: row.get("version"),
        started_at: row.get("started_at"),
        last_heartbeat: row.get("last_heartbeat"),
        current_connections: row.get::<i32, _>("current_connections") as u32,
        current_jobs: row.get::<i32, _>("current_jobs") as u32,
        cpu_usage: row.get::<f32, _>("cpu_usage"),
        memory_usage: row.get::<f32, _>("memory_usage"),
    })
}

/// Node count statistics.
#[derive(Debug, Clone, Default)]
pub struct NodeCounts {
    /// Active nodes.
    pub active: usize,
    /// Draining nodes.
    pub draining: usize,
    /// Dead nodes.
    pub dead: usize,
    /// Joining nodes.
    pub joining: usize,
    /// Total nodes.
    pub total: usize,
}

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

    #[test]
    fn test_node_counts_default() {
        let counts = NodeCounts::default();
        assert_eq!(counts.active, 0);
        assert_eq!(counts.total, 0);
    }
}