nodedb-cluster 0.1.1

// SPDX-License-Identifier: BUSL-1.1

//! [`RebalancerSubsystem`] — wraps the [`RebalancerLoop`] lifecycle.
//!
//! Depends on both `swim` (needs cluster membership to be established
//! before computing load-based plans) and `reachability` (a node that
//! is unreachable should not be a migration target; the circuit breaker
//! state is set by the reachability driver).

use std::sync::Arc;
use std::time::Instant;

use async_trait::async_trait;
use tokio::sync::watch;

use crate::migration_executor::MigrationExecutor;
use crate::rebalancer::driver::{RebalancerLoop, RebalancerLoopConfig};

use super::super::context::BootstrapCtx;
use super::super::errors::{BootstrapError, ShutdownError};
use super::super::health::SubsystemHealth;
use super::super::r#trait::{ClusterSubsystem, SubsystemHandle};
use super::adapters::{ExecutorDispatcher, MultiRaftElectionGate, NexarTransportMetricsProvider};

/// Owns the rebalancer driver lifecycle.
pub struct RebalancerSubsystem {
    cfg: RebalancerLoopConfig,
    /// Migration executor shared with `ExecutorDispatcher`. Wrapped in
    /// `Mutex` because `MigrationExecutor` takes `&self` (it is `Sync`
    /// already via its internal `Arc<Mutex<MultiRaft>>`), so `Arc` is
    /// sufficient — no `Mutex` needed here. The `Mutex` in `MigrationExecutor`
    /// itself guards `MultiRaft`.
    executor: Arc<MigrationExecutor>,
}

impl RebalancerSubsystem {
    pub fn new(cfg: RebalancerLoopConfig, executor: Arc<MigrationExecutor>) -> Self {
        Self { cfg, executor }
    }
}

#[async_trait]
impl ClusterSubsystem for RebalancerSubsystem {
    fn name(&self) -> &'static str {
        "rebalancer"
    }

    fn dependencies(&self) -> &'static [&'static str] {
        &["swim", "reachability"]
    }

    async fn start(&self, ctx: &BootstrapCtx) -> Result<SubsystemHandle, BootstrapError> {
        let metrics = Arc::new(NexarTransportMetricsProvider::new(Arc::clone(
            &ctx.transport,
        )));
        let dispatcher = Arc::new(ExecutorDispatcher::new(Arc::clone(&self.executor)));
        let gate = Arc::new(MultiRaftElectionGate::new(&ctx.multi_raft));

        let routing = Arc::clone(&ctx.routing);
        let topology = Arc::clone(&ctx.topology);

        let rloop = Arc::new(RebalancerLoop::new(
            self.cfg.clone(),
            metrics,
            dispatcher,
            gate,
            routing,
            topology,
        ));

        let (shutdown_tx, shutdown_rx) = watch::channel(false);
        let task = tokio::spawn(async move { rloop.run(shutdown_rx).await });

        Ok(SubsystemHandle::new("rebalancer", task, shutdown_tx))
    }

    async fn shutdown(&self, _deadline: Instant) -> Result<(), ShutdownError> {
        // Driven by SubsystemHandle::shutdown_tx.
        Ok(())
    }

    fn health(&self) -> SubsystemHealth {
        SubsystemHealth::Running
    }
}

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

    /// Verify name and dependency declarations without constructing a real
    /// MigrationExecutor (which requires a live MultiRaft + transport).
    /// We use a trivial wrapper to call the trait impls via the function pointers
    /// that the compiler fills in.
    fn _assert_name_is_rebalancer(s: &RebalancerSubsystem) {
        assert_eq!(s.name(), "rebalancer");
    }

    fn _assert_deps_correct(s: &RebalancerSubsystem) {
        assert_eq!(s.dependencies(), &["swim", "reachability"]);
    }

    // Static assertion: the dependency slice contains exactly two entries.
    // This can be verified at the type level without an instance.
    #[test]
    fn dependency_slice_has_two_entries() {
        // The impl declares `&["swim", "reachability"]`; verify that
        // literal independently so a rename surfacing from a refactor
        // breaks this test.
        const EXPECTED: &[&str] = &["swim", "reachability"];
        assert_eq!(EXPECTED.len(), 2);
        assert_eq!(EXPECTED[0], "swim");
        assert_eq!(EXPECTED[1], "reachability");
    }

    #[test]
    fn rebalancer_loop_config_default_is_sensible() {
        let cfg = RebalancerLoopConfig::default();
        assert!(cfg.interval.as_secs() > 0);
        assert!(cfg.backpressure_cpu_threshold > 0.0);
    }
}