photon_etcd_cluster/cluster_node/

mod.rs

//! Cluster membership and leader election.
//!
//! The [`ClusterNode`] manages cluster membership, leader election, and health
//! monitoring via etcd. Each worker process creates one `ClusterNode` instance.
//!
//! # Example
//!
//! ```ignore
//! use photon_etcd_cluster::ClusterNode;
//! use tokio::sync::broadcast;
//!
//! let node = ClusterNode::new(
//!     vec!["http://localhost:2379".to_string()],
//!     "worker-1".to_string(),
//!     "192.168.1.10".parse()?,
//!     "my-service".to_string(),
//!     Some(5),
//! );
//!
//! // Run node in background
//! let n = node.clone();
//! let (shutdown_tx, _) = broadcast::channel(1);
//! let mut shutdown_rx = shutdown_tx.subscribe();
//! tokio::spawn(async move {
//!     n.run(&mut shutdown_rx).await
//! });
//!
//! // Check leadership status
//! if node.is_leader() {
//!     println!("I am the leader!");
//! }
//! ```
//!
//! # With Metrics Collection
//!
//! ```ignore
//! use photon_etcd_cluster::{ClusterNodeBuilder, SystemMetricsCollector};
//!
//! let node = ClusterNodeBuilder::new(
//!     vec!["http://localhost:2379".to_string()],
//!     "worker-1".to_string(),
//!     "192.168.1.10".parse()?,
//!     "my-service".to_string(),
//! )
//! .ttl(5)
//! .metrics_collector(SystemMetricsCollector::new())
//! .metrics_update_interval(5)  // Update every 5 seconds
//! .build();
//! ```

mod election;
mod lease;
mod metrics_task;
mod state;

use std::net::IpAddr;
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;

use etcd_client::{Client, LeaderKey, ResignOptions};
use tokio::select;
use tokio::sync::{broadcast, watch};
use tracing::{error, info};

use crate::connection::{connect_with_retry, node_key};
use crate::error::ClusterError;
use crate::metrics::{MetricsCollector, NoopMetricsCollector};
use crate::types::HealthStatus;

use state::ClusterNodeInner;

/// Manages cluster membership and leader election via etcd.
///
/// Each worker process creates one `ClusterNode` instance that:
/// - Registers the worker in the cluster node registry
/// - Participates in leader election
/// - Maintains health status via lease keep-alive
#[derive(Clone)]
pub struct ClusterNode {
    inner: Arc<ClusterNodeInner>,
}

impl ClusterNode {
    /// Creates a new ClusterNode.
    ///
    /// # Arguments
    /// * `etcd_endpoints` - List of etcd endpoints (e.g., `["http://localhost:2379"]`)
    /// * `node_id` - Unique identifier for this node
    /// * `node_ip` - IP address of this node (used for service discovery)
    /// * `group_name` - Logical group name for cluster membership
    /// * `ttl` - Lease TTL in seconds (default: 5)
    ///
    /// # Panics
    /// Panics if `etcd_endpoints` is empty or contains invalid URLs.
    pub fn new(
        etcd_endpoints: Vec<String>,
        node_id: String,
        node_ip: IpAddr,
        group_name: String,
        ttl: Option<i64>,
    ) -> Self {
        assert!(!etcd_endpoints.is_empty(), "etcd_endpoints cannot be empty");
        for endpoint in &etcd_endpoints {
            assert!(
                endpoint.starts_with("http://") || endpoint.starts_with("https://"),
                "Invalid etcd endpoint '{}': must start with http:// or https://",
                endpoint
            );
        }

        ClusterNode {
            inner: Arc::new(ClusterNodeInner::new(
                etcd_endpoints,
                node_id,
                node_ip,
                group_name,
                ttl.unwrap_or(5),
            )),
        }
    }

    /// Returns the current health status of the node's etcd connection.
    pub fn current_health(&self) -> HealthStatus {
        self.inner.node_health.load()
    }

    /// Returns true if this node is currently the leader.
    pub fn is_leader(&self) -> bool {
        self.inner.is_leader.load(Ordering::Acquire)
    }

    pub(in crate::cluster_node) fn set_unhealthy(&self) {
        self.inner.node_health.store(HealthStatus::Unhealthy);
    }

    pub(in crate::cluster_node) fn set_healthy(&self) {
        self.inner.node_health.store(HealthStatus::Healthy);
    }

    // Accessors for submodules (visibility limited to cluster_node module)
    #[inline]
    pub(in crate::cluster_node) fn inner(&self) -> &ClusterNodeInner {
        &self.inner
    }

    #[inline]
    pub(in crate::cluster_node) fn set_leader(&self, is_leader: bool) {
        self.inner.is_leader.store(is_leader, Ordering::Release);
    }

    /// Starts the node's main loop.
    ///
    /// This method runs indefinitely, maintaining cluster membership and
    /// participating in leader election. It automatically reconnects on
    /// connection failures.
    ///
    /// If a metrics collector and update interval are configured, a background
    /// task will periodically update the node's metrics in etcd.
    ///
    /// # Arguments
    /// * `shutdown` - Shutdown signal receiver. When a message is received,
    ///   the node will gracefully resign leadership and exit.
    ///
    /// # Returns
    /// Returns `Ok(())` on graceful shutdown, `Err` if unable to connect to etcd.
    pub async fn run(&self, shutdown: &mut broadcast::Receiver<()>) -> Result<(), ClusterError> {
        let (leader_tx, leader_rx) = watch::channel(None);
        let (lease_healthy_tx, mut lease_healthy_rx) = watch::channel(true);
        let (metrics_healthy_tx, mut metrics_healthy_rx) = watch::channel(true);

        loop {
            info!("🔵 Starting main loop ({})", self.inner.node_id);
            let _ = leader_tx.send(None);
            self.set_leader(false); // Reset leader state at start of each loop

            // Reset lease health for new iteration (only if it was false)
            // Use send_if_modified to avoid triggering change notification if already true
            lease_healthy_tx.send_if_modified(|val| {
                if !*val {
                    *val = true;
                    true
                } else {
                    false
                }
            });
            metrics_healthy_tx.send_if_modified(|val| {
                if !*val {
                    *val = true;
                    true
                } else {
                    false
                }
            });

            let client_watcher = shutdown.resubscribe();
            let client = connect_with_retry(&self.inner.etcd_endpoints, Some(client_watcher)).await;
            let Some(mut client) = client else {
                return Err(ClusterError::Internal("Failed to connect to etcd".into()));
            };

            let (lease_id, node_lease_ka) =
                match lease::create_lease(self, &mut client, self.inner.ttl, lease_healthy_tx.clone())
                    .await
                {
                    Ok((lease_id, keepalive)) => (lease_id, keepalive),
                    Err(e) => {
                        error!("Failed to create node lease: {}", e);
                        self.set_unhealthy();
                        tokio::time::sleep(Duration::from_secs(1)).await;
                        continue;
                    }
                };

            match lease::register_local_node(self, &mut client, lease_id).await {
                Ok(_) => {
                    info!("Node registered: {}", self.inner.node_id);
                    self.set_healthy();
                }
                Err(e) => {
                    error!("Failed to register local node: {}", e);
                    self.set_unhealthy();
                    node_lease_ka.abort();
                    // Await the aborted task to ensure cleanup completes
                    let _ = node_lease_ka.await;
                    tokio::time::sleep(Duration::from_secs(1)).await;
                    continue;
                }
            }

            // Spawn metrics update task if configured
            let metrics_task =
                metrics_task::spawn_metrics_update_task(self, lease_id, metrics_healthy_tx.clone());

            let self_clone = self.clone();
            let leader_tx_copy = leader_tx.clone();
            let shutdown_watcher = shutdown.resubscribe();

            // Mark the initial values as seen so changed() doesn't trigger immediately
            // borrow_and_update() marks the current value as seen
            let _ = lease_healthy_rx.borrow_and_update();
            let _ = metrics_healthy_rx.borrow_and_update();

            select! {
                _ = shutdown.recv() => {
                    node_lease_ka.abort();
                    metrics_task::abort_metrics_task(metrics_task).await;
                    // Await the aborted task to ensure cleanup completes
                    let _ = node_lease_ka.await;

                    self.leave(&mut client, lease_id).await.ok();

                    let leader = leader_rx.borrow().clone();
                    if let Some(leader) = leader {
                        self.resign(&mut client, leader.clone()).await.ok();
                    }
                    self.set_leader(false); // Clear leader state on shutdown
                    info!("ClusterNode shutdown complete ({})", self.inner.node_id);
                    return Ok(());
                },
                // React to lease health changes
                _ = lease_healthy_rx.changed() => {
                    if !*lease_healthy_rx.borrow() {
                        error!("Lease keep-alive failed - triggering reconnection");
                        node_lease_ka.abort();
                        metrics_task::abort_metrics_task(metrics_task).await;
                        let _ = node_lease_ka.await;

                        let leader = leader_rx.borrow().clone();
                        if let Some(leader) = leader {
                            self.resign(&mut client, leader.clone()).await.ok();
                        }

                        self.set_unhealthy();
                        self.set_leader(false);
                        self.leave(&mut client, lease_id).await.ok();
                        continue;
                    }
                },
                // React to metrics update health changes
                _ = metrics_healthy_rx.changed() => {
                    if !*metrics_healthy_rx.borrow() {
                        error!("Metrics update failed repeatedly - triggering reconnection");
                        node_lease_ka.abort();
                        metrics_task::abort_metrics_task(metrics_task).await;
                        let _ = node_lease_ka.await;

                        let leader = leader_rx.borrow().clone();
                        if let Some(leader) = leader {
                            self.resign(&mut client, leader.clone()).await.ok();
                        }

                        self.set_unhealthy();
                        self.set_leader(false);
                        self.leave(&mut client, lease_id).await.ok();
                        continue;
                    }
                },
                res = election::run_election_loop(&self_clone, lease_id, leader_tx_copy, shutdown_watcher) => {
                    if let Err(e) = res {
                        error!("Election loop error: {}", e);

                        let leader = leader_rx.borrow().clone();
                        if let Some(leader) = leader {
                            self.resign(&mut client, leader.clone()).await.ok();
                        }

                        self.set_unhealthy();
                        self.set_leader(false); // Clear leader state on error
                        node_lease_ka.abort();
                        metrics_task::abort_metrics_task(metrics_task).await;
                        // Await the aborted task to ensure cleanup completes
                        let _ = node_lease_ka.await;
                        self.leave(&mut client, lease_id).await.ok();
                        continue;
                    }
                },
            }
        }
    }

    /// Removes this node from the cluster registry.
    pub async fn leave(&self, client: &mut Client, lease_id: i64) -> Result<(), ClusterError> {
        // Delete node registration
        let key = node_key(&self.inner.group_name, &self.inner.node_id);
        if let Err(e) = client.delete(key, None).await {
            error!("Failed to delete node registration: {}", e);
        }
        self.set_unhealthy();
        self.set_leader(false); // Clear leader state when leaving
        if let Err(e) = client.lease_revoke(lease_id).await {
            error!("Failed to revoke lease: {}", e);
        };

        Ok(())
    }

    /// Resigns from leadership.
    pub async fn resign(&self, client: &mut Client, leader: LeaderKey) -> Result<(), ClusterError> {
        let resign_option = ResignOptions::new().with_leader(leader);
        client.resign(Some(resign_option)).await.map_err(|e| {
            ClusterError::Internal(format!("Failed to resign from election: {}", e))
        })?;
        self.set_leader(false); // Clear leader state when resigning
        Ok(())
    }
}

/// Builder for creating [`ClusterNode`] instances with custom configuration.
///
/// This builder provides a fluent API for configuring cluster nodes with
/// optional features like metrics collection.
///
/// # Example
///
/// ```ignore
/// use photon_etcd_cluster::{ClusterNodeBuilder, SystemMetricsCollector};
///
/// // Basic usage (no metrics)
/// let node = ClusterNodeBuilder::new(
///     vec!["http://localhost:2379".to_string()],
///     "worker-1".to_string(),
///     "192.168.1.10".parse()?,
///     "my-service".to_string(),
/// )
/// .ttl(5)
/// .build();
///
/// // With system metrics
/// let node = ClusterNodeBuilder::new(
///     vec!["http://localhost:2379".to_string()],
///     "worker-1".to_string(),
///     "192.168.1.10".parse()?,
///     "my-service".to_string(),
/// )
/// .ttl(5)
/// .metrics_collector(SystemMetricsCollector::new())
/// .metrics_update_interval(5)  // Update every 5 seconds
/// .build();
/// ```
pub struct ClusterNodeBuilder {
    etcd_endpoints: Vec<String>,
    node_id: String,
    node_ip: IpAddr,
    group_name: String,
    ttl: i64,
    metrics_collector: Arc<dyn MetricsCollector>,
    metrics_update_interval: u64,
}

impl ClusterNodeBuilder {
    /// Creates a new builder with required parameters.
    ///
    /// # Arguments
    ///
    /// * `etcd_endpoints` - List of etcd endpoints (e.g., `["http://localhost:2379"]`)
    /// * `node_id` - Unique identifier for this node
    /// * `node_ip` - IP address of this node (used for service discovery)
    /// * `group_name` - Logical group name for cluster membership
    ///
    /// # Panics
    ///
    /// Panics if `etcd_endpoints` is empty or contains invalid URLs.
    pub fn new(
        etcd_endpoints: Vec<String>,
        node_id: String,
        node_ip: IpAddr,
        group_name: String,
    ) -> Self {
        assert!(!etcd_endpoints.is_empty(), "etcd_endpoints cannot be empty");
        for endpoint in &etcd_endpoints {
            assert!(
                endpoint.starts_with("http://") || endpoint.starts_with("https://"),
                "Invalid etcd endpoint '{}': must start with http:// or https://",
                endpoint
            );
        }

        Self {
            etcd_endpoints,
            node_id,
            node_ip,
            group_name,
            ttl: 5, // Default TTL
            metrics_collector: Arc::new(NoopMetricsCollector::new()),
            metrics_update_interval: 0, // Disabled by default
        }
    }

    /// Sets the lease TTL in seconds.
    ///
    /// The TTL determines how long etcd will wait before considering the node
    /// as failed if heartbeats stop. Default is 5 seconds.
    ///
    /// # Arguments
    ///
    /// * `ttl` - Lease TTL in seconds (must be > 0)
    pub fn ttl(mut self, ttl: i64) -> Self {
        assert!(ttl > 0, "TTL must be positive");
        self.ttl = ttl;
        self
    }

    /// Sets a custom metrics collector.
    ///
    /// The collector will be called periodically to gather metrics that are
    /// stored in etcd alongside the node registration. These metrics can be
    /// used by load balancers for weighted routing decisions.
    ///
    /// By default, a [`NoopMetricsCollector`] is used which produces no metrics.
    ///
    /// # Example
    ///
    /// ```ignore
    /// use photon_etcd_cluster::{ClusterNodeBuilder, SystemMetricsCollector};
    ///
    /// let node = ClusterNodeBuilder::new(endpoints, id, ip, group)
    ///     .metrics_collector(SystemMetricsCollector::new())
    ///     .metrics_update_interval(5)
    ///     .build();
    /// ```
    pub fn metrics_collector<C: MetricsCollector + 'static>(mut self, collector: C) -> Self {
        self.metrics_collector = Arc::new(collector);
        self
    }

    /// Sets the interval in seconds for metrics updates.
    ///
    /// When set to a value > 0, a background task will periodically update
    /// the node's metrics in etcd at the specified interval.
    ///
    /// Set to 0 (default) to disable periodic metrics updates. In this case,
    /// metrics will only be set during initial registration.
    ///
    /// Recommended value: Same as or slightly less than the TTL to ensure
    /// fresh metrics are available for load balancing decisions.
    ///
    /// # Arguments
    ///
    /// * `interval` - Update interval in seconds (0 = disabled)
    pub fn metrics_update_interval(mut self, interval: u64) -> Self {
        self.metrics_update_interval = interval;
        self
    }

    /// Builds the [`ClusterNode`] with the configured options.
    pub fn build(self) -> ClusterNode {
        if self.metrics_update_interval > 0 {
            info!(
                "ClusterNode configured with metrics collector (update_interval={}s)",
                self.metrics_update_interval
            );
        }

        ClusterNode {
            inner: Arc::new(ClusterNodeInner::with_metrics(
                self.etcd_endpoints,
                self.node_id,
                self.node_ip,
                self.group_name,
                self.ttl,
                self.metrics_collector,
                self.metrics_update_interval,
            )),
        }
    }
}

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

    #[tokio::test]
    async fn test_leader_state_tracking() {
        let node = ClusterNode::new(
            vec!["http://localhost:2379".to_string()],
            "test-node".to_string(),
            "127.0.0.1".parse().unwrap(),
            "test-group".to_string(),
            Some(5),
        );

        // Initially not a leader
        assert!(!node.is_leader());

        // Simulate becoming a leader
        node.set_leader(true);
        assert!(node.is_leader());

        // Simulate stepping down
        node.set_leader(false);
        assert!(!node.is_leader());
    }

    #[test]
    #[should_panic(expected = "etcd_endpoints cannot be empty")]
    fn test_empty_endpoints_panics() {
        let _ = ClusterNode::new(
            vec![],
            "test".to_string(),
            "127.0.0.1".parse().unwrap(),
            "test".to_string(),
            None,
        );
    }

    #[test]
    #[should_panic(expected = "Invalid etcd endpoint")]
    fn test_invalid_endpoint_panics() {
        let _ = ClusterNode::new(
            vec!["localhost:2379".to_string()],
            "test".to_string(),
            "127.0.0.1".parse().unwrap(),
            "test".to_string(),
            None,
        );
    }
}