use crate::infrastructure::high_availability::config::HighAvailabilityConfig;
use crate::infrastructure::high_availability::{ClusterStatus, FailoverPhase, HaError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};
use uuid::Uuid;
pub struct ClusterManager {
config: Arc<HighAvailabilityConfig>,
nodes: Arc<RwLock<HashMap<String, NodeInfo>>>,
node_id: String,
current_leader: Arc<RwLock<Option<String>>>,
status: Arc<RwLock<ClusterStatus>>,
discovery_service: Box<dyn NodeDiscovery + Send + Sync>,
membership_service: Box<dyn ClusterMembership + Send + Sync>,
}
impl ClusterManager {
pub fn new(config: &HighAvailabilityConfig) -> Self {
let node_id = Uuid::new_v4().to_string();
let discovery_service = create_discovery_service(config);
let membership_service = create_membership_service(config);
Self {
config: Arc::new(config.clone()),
nodes: Arc::new(RwLock::new(HashMap::new())),
node_id,
current_leader: Arc::new(RwLock::new(None)),
status: Arc::new(RwLock::new(ClusterStatus::Initializing)),
discovery_service,
membership_service,
}
}
pub async fn initialize(&mut self) -> Result<(), HaError> {
info!("Initializing cluster manager");
let discovered_nodes = self
.discovery_service
.discover_nodes()
.await
.map_err(|e| HaError::ClusterError(format!("Failed to discover nodes: {e}")))?;
let mut nodes = self.nodes.write().await;
for node_addr in discovered_nodes {
let node_info = NodeInfo {
id: node_addr.clone(), address: node_addr,
status: NodeStatus::Unknown,
role: NodeRole::Unknown,
last_heartbeat: None,
metadata: HashMap::new(),
};
nodes.insert(node_info.id.clone(), node_info);
}
let self_node = NodeInfo {
id: self.node_id.clone(),
address: self
.membership_service
.get_local_address()
.map_err(|e| HaError::ClusterError(format!("Failed to get local address: {e}")))?,
status: NodeStatus::Starting,
role: NodeRole::Follower, last_heartbeat: Some(chrono::Utc::now()),
metadata: HashMap::new(),
};
nodes.insert(self.node_id.clone(), self_node);
self.membership_service.initialize().await.map_err(|e| {
HaError::ClusterError(format!("Failed to initialize membership service: {e}"))
})?;
*self.status.write().await = ClusterStatus::Initializing;
debug!("Cluster manager initialized with {} nodes", nodes.len());
Ok(())
}
pub async fn join_cluster(&mut self) -> Result<(), HaError> {
info!("Joining cluster {}", self.config.cluster.cluster_name);
self.membership_service
.join(self.node_id.clone())
.await
.map_err(|e| HaError::ClusterError(format!("Failed to join cluster: {e}")))?;
let mut nodes = self.nodes.write().await;
if let Some(node) = nodes.get_mut(&self.node_id) {
node.status = NodeStatus::Active;
}
let leader = self.current_leader.read().await;
if leader.is_none() {
drop(leader);
drop(nodes);
self.elect_leader().await?;
}
let mut status = self.status.write().await;
*status = ClusterStatus::Healthy;
info!(
"Successfully joined cluster {}",
self.config.cluster.cluster_name
);
Ok(())
}
pub async fn leave_cluster(&mut self) -> Result<(), HaError> {
info!("Leaving cluster {}", self.config.cluster.cluster_name);
self.membership_service
.leave(self.node_id.clone())
.await
.map_err(|e| HaError::ClusterError(format!("Failed to leave cluster: {e}")))?;
let mut nodes = self.nodes.write().await;
if let Some(node) = nodes.get_mut(&self.node_id) {
node.status = NodeStatus::Leaving;
}
let mut status = self.status.write().await;
*status = ClusterStatus::Down {
reason: "Node left the cluster".to_string(),
};
info!(
"Successfully left cluster {}",
self.config.cluster.cluster_name
);
Ok(())
}
pub async fn get_status(&self) -> Result<ClusterStatus, HaError> {
let status = self.status.read().await;
Ok(status.clone())
}
pub async fn update_config(&mut self, config: &HighAvailabilityConfig) -> Result<(), HaError> {
info!("Updating cluster configuration");
self.config = Arc::new(config.clone());
if self.discovery_service.needs_update(config) {
self.discovery_service = create_discovery_service(config);
}
if self.membership_service.needs_update(config) {
let old_membership = std::mem::replace(
&mut self.membership_service,
create_membership_service(config),
);
old_membership
.leave(self.node_id.clone())
.await
.map_err(|e| {
HaError::ClusterError(format!("Failed to leave old membership: {e}"))
})?;
self.membership_service.initialize().await.map_err(|e| {
HaError::ClusterError(format!("Failed to initialize new membership: {e}"))
})?;
self.membership_service
.join(self.node_id.clone())
.await
.map_err(|e| {
HaError::ClusterError(format!("Failed to join with new membership: {e}"))
})?;
}
Ok(())
}
async fn elect_leader(&mut self) -> Result<(), HaError> {
info!("Starting leader election");
let nodes = self.nodes.read().await;
let active_nodes: Vec<&NodeInfo> = nodes
.values()
.filter(|n| n.status == NodeStatus::Active)
.collect();
if active_nodes.is_empty() {
return Err(HaError::ClusterError(
"No active nodes for leader election".to_string(),
));
}
let new_leader = active_nodes
.iter()
.min_by_key(|n| &n.id)
.map(|n| n.id.clone())
.ok_or_else(|| {
HaError::ClusterError("No active nodes available for leader election".to_string())
})?;
let mut leader = self.current_leader.write().await;
*leader = Some(new_leader.clone());
drop(leader);
drop(nodes);
let mut nodes = self.nodes.write().await;
for (id, node) in nodes.iter_mut() {
if node.status == NodeStatus::Active {
if *id == new_leader {
node.role = NodeRole::Leader;
} else {
node.role = NodeRole::Follower;
}
}
}
info!("Leader election completed, new leader: {}", new_leader);
Ok(())
}
pub async fn process_heartbeat(&mut self, node_id: &str) -> Result<(), HaError> {
let mut nodes = self.nodes.write().await;
if let Some(node) = nodes.get_mut(node_id) {
node.last_heartbeat = Some(chrono::Utc::now());
if node.status == NodeStatus::Unknown || node.status == NodeStatus::Suspected {
node.status = NodeStatus::Active;
}
} else {
let node_info = NodeInfo {
id: node_id.to_string(),
address: node_id.to_string(), status: NodeStatus::Active,
role: NodeRole::Follower,
last_heartbeat: Some(chrono::Utc::now()),
metadata: HashMap::new(),
};
nodes.insert(node_id.to_string(), node_info);
}
Ok(())
}
pub async fn check_node_health(&mut self) -> Result<(), HaError> {
debug!("Checking node health");
let now = chrono::Utc::now();
let node_timeout = chrono::Duration::from_std(self.config.cluster.node_timeout)
.map_err(|_| HaError::ClusterError("Invalid node timeout duration".to_string()))?;
let mut nodes = self.nodes.write().await;
let mut active_count = 0;
let mut total_count = 0;
let mut failures = Vec::new();
for (id, node) in nodes.iter_mut() {
if id == &self.node_id {
continue;
}
total_count += 1;
if let Some(last_heartbeat) = node.last_heartbeat {
let elapsed = now - last_heartbeat;
if elapsed > node_timeout {
if node.status == NodeStatus::Active {
warn!(
"Node {} has not sent heartbeat in {:?}, marking as suspected",
id, elapsed
);
node.status = NodeStatus::Suspected;
} else if node.status == NodeStatus::Suspected {
error!("Node {} has not recovered, marking as failed", id);
node.status = NodeStatus::Failed;
failures.push(id.clone());
}
} else if node.status == NodeStatus::Active {
active_count += 1;
}
}
}
drop(nodes);
let mut status = self.status.write().await;
if !failures.is_empty() {
if active_count < total_count && active_count > 0 {
*status = ClusterStatus::Degraded {
active_nodes: active_count,
total_nodes: total_count,
details: format!("Nodes failed: {}", failures.join(", ")),
};
} else if active_count == 0 {
*status = ClusterStatus::Down {
reason: "All nodes failed".to_string(),
};
}
let leader = self.current_leader.read().await;
if let Some(leader_id) = leader.as_ref() {
if failures.contains(leader_id) {
let leader_id_clone = leader_id.clone();
drop(leader);
drop(status);
let mut status = self.status.write().await;
*status = ClusterStatus::Failover {
primary_node: leader_id_clone.clone(),
failing_node: Some(leader_id_clone.clone()),
failover_phase: FailoverPhase::Detection,
};
info!(
"Leader node {} failed, cluster entering failover state",
leader_id_clone
);
}
}
} else if active_count == total_count {
*status = ClusterStatus::Healthy;
}
Ok(())
}
pub async fn get_nodes(&self) -> Result<Vec<NodeInfo>, HaError> {
let nodes = self.nodes.read().await;
Ok(nodes.values().cloned().collect())
}
pub async fn get_leader(&self) -> Result<Option<String>, HaError> {
let leader = self.current_leader.read().await;
Ok(leader.clone())
}
pub async fn is_leader(&self) -> Result<bool, HaError> {
let leader = self.current_leader.read().await;
Ok(leader.as_ref() == Some(&self.node_id))
}
}
fn create_discovery_service(
config: &HighAvailabilityConfig,
) -> Box<dyn NodeDiscovery + Send + Sync> {
match config.cluster.discovery_method {
crate::infrastructure::high_availability::config::DiscoveryMethod::Static => {
Box::new(StaticDiscovery::new(config.cluster.static_nodes.clone()))
}
crate::infrastructure::high_availability::config::DiscoveryMethod::Dns => {
if let Some(ref dns_url) = config.cluster.dns_discovery_url {
Box::new(DnsDiscovery::new(dns_url.clone()))
} else {
Box::new(StaticDiscovery::new(config.cluster.static_nodes.clone()))
}
}
crate::infrastructure::high_availability::config::DiscoveryMethod::Kubernetes => {
if let Some(ref service_name) = config.cluster.k8s_service_name {
Box::new(KubernetesDiscovery::new(service_name.clone()))
} else {
Box::new(StaticDiscovery::new(config.cluster.static_nodes.clone()))
}
}
_ => {
Box::new(StaticDiscovery::new(config.cluster.static_nodes.clone()))
}
}
}
fn create_membership_service(
config: &HighAvailabilityConfig,
) -> Box<dyn ClusterMembership + Send + Sync> {
Box::new(BasicMembership::new(config))
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeInfo {
pub id: String,
pub address: String,
pub status: NodeStatus,
pub role: NodeRole,
pub last_heartbeat: Option<chrono::DateTime<chrono::Utc>>,
pub metadata: HashMap<String, String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeStatus {
Unknown,
Starting,
Active,
Suspected,
Failed,
Leaving,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum NodeRole {
Unknown,
Leader,
Follower,
Observer,
}
#[async_trait::async_trait]
pub trait NodeDiscovery {
async fn discover_nodes(&self) -> Result<Vec<String>, String>;
fn needs_update(&self, config: &HighAvailabilityConfig) -> bool;
}
pub struct StaticDiscovery {
nodes: Vec<String>,
}
impl StaticDiscovery {
pub fn new(nodes: Vec<String>) -> Self {
Self { nodes }
}
}
#[async_trait::async_trait]
impl NodeDiscovery for StaticDiscovery {
async fn discover_nodes(&self) -> Result<Vec<String>, String> {
Ok(self.nodes.clone())
}
fn needs_update(&self, config: &HighAvailabilityConfig) -> bool {
self.nodes != config.cluster.static_nodes
}
}
pub struct DnsDiscovery {
dns_url: String,
}
impl DnsDiscovery {
pub fn new(dns_url: String) -> Self {
Self { dns_url }
}
}
#[async_trait::async_trait]
impl NodeDiscovery for DnsDiscovery {
async fn discover_nodes(&self) -> Result<Vec<String>, String> {
Ok(vec![
format!("node1.{}", self.dns_url),
format!("node2.{}", self.dns_url),
format!("node3.{}", self.dns_url),
])
}
fn needs_update(&self, config: &HighAvailabilityConfig) -> bool {
if let Some(ref dns_url) = config.cluster.dns_discovery_url {
self.dns_url != *dns_url
} else {
true
}
}
}
pub struct KubernetesDiscovery {
service_name: String,
}
impl KubernetesDiscovery {
pub fn new(service_name: String) -> Self {
Self { service_name }
}
}
#[async_trait::async_trait]
impl NodeDiscovery for KubernetesDiscovery {
async fn discover_nodes(&self) -> Result<Vec<String>, String> {
Ok(vec![
format!("{}-0.{}", self.service_name, self.service_name),
format!("{}-1.{}", self.service_name, self.service_name),
format!("{}-2.{}", self.service_name, self.service_name),
])
}
fn needs_update(&self, config: &HighAvailabilityConfig) -> bool {
if let Some(ref service_name) = config.cluster.k8s_service_name {
self.service_name != *service_name
} else {
true
}
}
}
#[async_trait::async_trait]
pub trait ClusterMembership {
async fn initialize(&mut self) -> Result<(), String>;
async fn join(&self, node_id: String) -> Result<(), String>;
async fn leave(&self, node_id: String) -> Result<(), String>;
fn get_local_address(&self) -> Result<String, String>;
fn needs_update(&self, _config: &HighAvailabilityConfig) -> bool {
false
}
}
pub struct BasicMembership {
config: HighAvailabilityConfig,
local_address: String,
}
impl BasicMembership {
pub fn new(config: &HighAvailabilityConfig) -> Self {
let local_address = "127.0.0.1:5001".to_string();
Self {
config: config.clone(),
local_address,
}
}
}
#[async_trait::async_trait]
impl ClusterMembership for BasicMembership {
async fn initialize(&mut self) -> Result<(), String> {
Ok(())
}
async fn join(&self, node_id: String) -> Result<(), String> {
info!(
"Node {} joining cluster {}",
node_id, self.config.cluster.cluster_name
);
Ok(())
}
async fn leave(&self, node_id: String) -> Result<(), String> {
info!(
"Node {} leaving cluster {}",
node_id, self.config.cluster.cluster_name
);
Ok(())
}
fn get_local_address(&self) -> Result<String, String> {
Ok(self.local_address.clone())
}
}