use parking_lot::RwLock;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::broadcast;
use uuid::Uuid;
use super::role_manager::NodeRole;
use crate::Result;
static TOPOLOGY_MANAGER: once_cell::sync::Lazy<TopologyManager> =
once_cell::sync::Lazy::new(|| TopologyManager::new(None));
pub fn topology_manager() -> &'static TopologyManager {
&TOPOLOGY_MANAGER
}
#[derive(Debug, Clone)]
pub struct NodeInfo {
pub node_id: Uuid,
pub alias: Option<String>,
pub role: NodeRole,
pub client_addr: String,
pub replication_addr: String,
pub last_lsn: u64,
pub replication_lag_ms: u64,
pub last_seen: Instant,
pub is_healthy: bool,
pub priority: u32,
pub weight: u32,
pub tags: HashMap<String, String>,
pub health_message: Option<String>,
pub health_failures: u32,
}
impl NodeInfo {
pub fn new(node_id: Uuid, role: NodeRole, client_addr: String, replication_addr: String) -> Self {
Self {
node_id,
alias: None,
role,
client_addr,
replication_addr,
last_lsn: 0,
replication_lag_ms: 0,
last_seen: Instant::now(),
is_healthy: true,
priority: 100,
weight: 100,
tags: HashMap::new(),
health_message: None,
health_failures: 0,
}
}
pub fn with_alias(mut self, alias: impl Into<String>) -> Self {
self.alias = Some(alias.into());
self
}
pub fn with_priority(mut self, priority: u32) -> Self {
self.priority = priority;
self
}
pub fn with_weight(mut self, weight: u32) -> Self {
self.weight = weight;
self
}
pub fn display_name(&self) -> String {
self.alias.clone().unwrap_or_else(|| {
let uuid_str = self.node_id.to_string();
format!("{}...", &uuid_str[..8])
})
}
pub fn can_read(&self) -> bool {
self.is_healthy && self.role.can_read()
}
pub fn can_write(&self) -> bool {
self.is_healthy && self.role.can_write()
}
pub fn time_since_seen(&self) -> Duration {
self.last_seen.elapsed()
}
}
#[derive(Debug, Clone)]
pub enum TopologyEvent {
NodeJoined(NodeInfo),
NodeLeft { node_id: Uuid },
RoleChanged {
node_id: Uuid,
old_role: NodeRole,
new_role: NodeRole,
},
AliasChanged {
node_id: Uuid,
old_alias: Option<String>,
new_alias: Option<String>,
},
HealthChanged { node_id: Uuid, is_healthy: bool },
PrimaryChanged {
old_primary: Option<Uuid>,
new_primary: Uuid,
},
Refreshed,
}
pub struct TopologyManager {
nodes: RwLock<HashMap<Uuid, NodeInfo>>,
aliases: RwLock<HashMap<String, Uuid>>,
primary_node: RwLock<Option<Uuid>>,
event_tx: broadcast::Sender<TopologyEvent>,
health_timeout: Duration,
local_node_id: Option<Uuid>,
}
impl TopologyManager {
pub fn new(local_node_id: Option<Uuid>) -> Self {
let (event_tx, _) = broadcast::channel(64);
Self {
nodes: RwLock::new(HashMap::new()),
aliases: RwLock::new(HashMap::new()),
primary_node: RwLock::new(None),
event_tx,
health_timeout: Duration::from_secs(10),
local_node_id,
}
}
pub fn with_health_timeout(mut self, timeout: Duration) -> Self {
self.health_timeout = timeout;
self
}
pub fn subscribe(&self) -> broadcast::Receiver<TopologyEvent> {
self.event_tx.subscribe()
}
pub fn register_node(&self, info: NodeInfo) {
let is_new;
let old_role;
let old_alias;
let is_primary = info.role == NodeRole::Primary;
let node_id = info.node_id;
{
let mut nodes = self.nodes.write();
is_new = !nodes.contains_key(&info.node_id);
old_role = nodes.get(&info.node_id).map(|n| n.role);
old_alias = nodes.get(&info.node_id).and_then(|n| n.alias.clone());
nodes.insert(info.node_id, info.clone());
}
if info.alias != old_alias {
let mut aliases = self.aliases.write();
if let Some(ref old) = old_alias {
aliases.remove(old);
}
if let Some(ref new_alias) = info.alias {
aliases.insert(new_alias.clone(), node_id);
}
if !is_new {
let _ = self.event_tx.send(TopologyEvent::AliasChanged {
node_id,
old_alias,
new_alias: info.alias.clone(),
});
}
}
if is_new {
let _ = self.event_tx.send(TopologyEvent::NodeJoined(info.clone()));
} else if let Some(old) = old_role {
if old != info.role {
let _ = self.event_tx.send(TopologyEvent::RoleChanged {
node_id,
old_role: old,
new_role: info.role,
});
}
}
if is_primary {
let old_primary = *self.primary_node.read();
if old_primary != Some(node_id) {
*self.primary_node.write() = Some(node_id);
let _ = self.event_tx.send(TopologyEvent::PrimaryChanged {
old_primary,
new_primary: node_id,
});
}
}
}
pub fn remove_node(&self, node_id: Uuid) {
let removed = self.nodes.write().remove(&node_id);
if let Some(ref node) = removed {
if let Some(ref alias) = node.alias {
self.aliases.write().remove(alias);
}
let _ = self.event_tx.send(TopologyEvent::NodeLeft { node_id });
let mut primary = self.primary_node.write();
if *primary == Some(node_id) {
*primary = None;
}
}
}
pub fn update_health(&self, node_id: Uuid, is_healthy: bool) {
let changed;
{
let mut nodes = self.nodes.write();
if let Some(node) = nodes.get_mut(&node_id) {
changed = node.is_healthy != is_healthy;
node.is_healthy = is_healthy;
node.last_seen = Instant::now();
} else {
return;
}
}
if changed {
let _ = self.event_tx.send(TopologyEvent::HealthChanged { node_id, is_healthy });
}
}
pub fn update_lsn(&self, node_id: Uuid, lsn: u64, lag_ms: u64) {
let mut nodes = self.nodes.write();
if let Some(node) = nodes.get_mut(&node_id) {
node.last_lsn = lsn;
node.replication_lag_ms = lag_ms;
node.last_seen = Instant::now();
}
}
pub fn get_primary(&self) -> Option<NodeInfo> {
let primary_id = *self.primary_node.read();
primary_id.and_then(|id| self.nodes.read().get(&id).cloned())
}
pub fn get_primary_id(&self) -> Option<Uuid> {
*self.primary_node.read()
}
pub fn get_node(&self, node_id: Uuid) -> Option<NodeInfo> {
self.nodes.read().get(&node_id).cloned()
}
pub fn get_node_by_alias(&self, alias: &str) -> Option<NodeInfo> {
let node_id = self.aliases.read().get(alias).copied()?;
self.nodes.read().get(&node_id).cloned()
}
pub fn resolve_node_id(&self, identifier: &str) -> Option<Uuid> {
if let Some(node_id) = self.aliases.read().get(identifier).copied() {
return Some(node_id);
}
if let Ok(uuid) = Uuid::parse_str(identifier) {
if self.nodes.read().contains_key(&uuid) {
return Some(uuid);
}
}
None
}
pub fn set_alias(&self, node_id: Uuid, alias: Option<String>) -> bool {
let mut nodes = self.nodes.write();
if let Some(node) = nodes.get_mut(&node_id) {
let old_alias = node.alias.clone();
{
let mut aliases = self.aliases.write();
if let Some(ref old) = old_alias {
aliases.remove(old);
}
if let Some(ref new_alias) = alias {
if let Some(&existing_id) = aliases.get(new_alias) {
if existing_id != node_id {
return false; }
}
aliases.insert(new_alias.clone(), node_id);
}
}
node.alias = alias.clone();
let _ = self.event_tx.send(TopologyEvent::AliasChanged {
node_id,
old_alias,
new_alias: alias,
});
true
} else {
false
}
}
pub fn get_all_aliases(&self) -> HashMap<String, Uuid> {
self.aliases.read().clone()
}
pub fn get_healthy_standbys(&self) -> Vec<NodeInfo> {
self.nodes
.read()
.values()
.filter(|n| n.is_healthy && n.role == NodeRole::Standby)
.cloned()
.collect()
}
pub fn get_all_nodes(&self) -> Vec<NodeInfo> {
self.nodes.read().values().cloned().collect()
}
pub fn get_read_nodes(&self) -> Vec<NodeInfo> {
self.nodes.read().values().filter(|n| n.can_read()).cloned().collect()
}
pub fn get_best_promotion_candidate(&self) -> Option<NodeInfo> {
self.nodes
.read()
.values()
.filter(|n| n.is_healthy && n.role == NodeRole::Standby)
.min_by(|a, b| {
a.priority
.cmp(&b.priority)
.then(a.replication_lag_ms.cmp(&b.replication_lag_ms))
})
.cloned()
}
pub fn check_health_timeouts(&self) -> Vec<Uuid> {
let mut timed_out = Vec::new();
{
let mut nodes = self.nodes.write();
for node in nodes.values_mut() {
if node.is_healthy && node.time_since_seen() > self.health_timeout {
node.is_healthy = false;
timed_out.push(node.node_id);
}
}
}
for node_id in &timed_out {
let _ = self.event_tx.send(TopologyEvent::HealthChanged {
node_id: *node_id,
is_healthy: false,
});
}
timed_out
}
pub fn select_read_standby(&self) -> Option<NodeInfo> {
let standbys = self.get_healthy_standbys();
if standbys.is_empty() {
return None;
}
let total_weight: u32 = standbys.iter().map(|s| s.weight).sum();
if total_weight == 0 {
return standbys.first().cloned();
}
let random_point = rand::random::<u32>() % total_weight;
let mut cumulative = 0;
for standby in standbys {
cumulative += standby.weight;
if random_point < cumulative {
return Some(standby);
}
}
None
}
pub fn get_cluster_summary(&self) -> ClusterSummary {
let nodes = self.nodes.read();
let mut summary = ClusterSummary::default();
for node in nodes.values() {
summary.total_nodes += 1;
if node.is_healthy {
summary.healthy_nodes += 1;
}
match node.role {
NodeRole::Primary => summary.primary_count += 1,
NodeRole::Standby => summary.standby_count += 1,
_ => summary.transitioning_count += 1,
}
if node.role == NodeRole::Standby {
summary.max_lag_ms = summary.max_lag_ms.max(node.replication_lag_ms);
}
}
summary.primary_id = *self.primary_node.read();
summary
}
}
#[derive(Debug, Default, Clone)]
pub struct ClusterSummary {
pub total_nodes: usize,
pub healthy_nodes: usize,
pub primary_count: usize,
pub standby_count: usize,
pub transitioning_count: usize,
pub primary_id: Option<Uuid>,
pub max_lag_ms: u64,
}
impl ClusterSummary {
pub fn is_healthy(&self) -> bool {
self.primary_count == 1 && self.healthy_nodes > 1
}
}
#[derive(Debug, Clone)]
pub struct NodeStatus {
pub node_id: Uuid,
pub alias: Option<String>,
pub display_name: String,
pub role: String,
pub client_addr: String,
pub replication_addr: String,
pub is_healthy: bool,
pub health_message: Option<String>,
pub health_failures: u32,
pub last_seen_secs: u64,
pub lsn: u64,
pub lag_ms: u64,
pub priority: u32,
pub weight: u32,
pub tags: HashMap<String, String>,
}
#[derive(Debug, Clone)]
pub struct TopologyDescription {
pub nodes: Vec<NodeStatus>,
pub primary_id: Option<Uuid>,
pub primary_name: Option<String>,
pub cluster_healthy: bool,
pub health_summary: String,
pub total_nodes: usize,
pub healthy_nodes: usize,
pub max_lag_ms: u64,
}
impl TopologyManager {
pub fn get_topology_description(&self) -> TopologyDescription {
let nodes_map = self.nodes.read();
let primary_id = *self.primary_node.read();
let mut nodes: Vec<NodeStatus> = nodes_map
.values()
.map(|n| NodeStatus {
node_id: n.node_id,
alias: n.alias.clone(),
display_name: n.display_name(),
role: format!("{:?}", n.role),
client_addr: n.client_addr.clone(),
replication_addr: n.replication_addr.clone(),
is_healthy: n.is_healthy,
health_message: n.health_message.clone(),
health_failures: n.health_failures,
last_seen_secs: n.time_since_seen().as_secs(),
lsn: n.last_lsn,
lag_ms: n.replication_lag_ms,
priority: n.priority,
weight: n.weight,
tags: n.tags.clone(),
})
.collect();
nodes.sort_by(|a, b| {
let a_is_primary = Some(a.node_id) == primary_id;
let b_is_primary = Some(b.node_id) == primary_id;
b_is_primary
.cmp(&a_is_primary)
.then(a.priority.cmp(&b.priority))
.then(a.display_name.cmp(&b.display_name))
});
let total_nodes = nodes.len();
let healthy_nodes = nodes.iter().filter(|n| n.is_healthy).count();
let max_lag_ms = nodes
.iter()
.filter(|n| n.role == "Standby")
.map(|n| n.lag_ms)
.max()
.unwrap_or(0);
let primary_count = nodes.iter().filter(|n| n.role == "Primary").count();
let standby_count = nodes.iter().filter(|n| n.role == "Standby").count();
let cluster_healthy = primary_count == 1 && healthy_nodes > 1;
let health_summary = if cluster_healthy {
format!(
"Healthy: 1 primary, {} standbys, {} nodes total",
standby_count, total_nodes
)
} else if primary_count == 0 {
"CRITICAL: No primary node".to_string()
} else if primary_count > 1 {
format!("CRITICAL: Multiple primaries detected ({})", primary_count)
} else if healthy_nodes <= 1 {
format!(
"WARNING: Insufficient healthy nodes ({}/{})",
healthy_nodes, total_nodes
)
} else {
format!("WARNING: {} healthy nodes out of {}", healthy_nodes, total_nodes)
};
let primary_name = primary_id.and_then(|id| nodes_map.get(&id).map(|n| n.display_name()));
TopologyDescription {
nodes,
primary_id,
primary_name,
cluster_healthy,
health_summary,
total_nodes,
healthy_nodes,
max_lag_ms,
}
}
}
#[derive(Debug, Clone)]
pub struct DiscoveryConfig {
pub seed_nodes: Vec<String>,
pub refresh_interval: Duration,
pub dns_discovery: bool,
pub dns_hostname: Option<String>,
}
impl Default for DiscoveryConfig {
fn default() -> Self {
Self {
seed_nodes: vec![],
refresh_interval: Duration::from_secs(5),
dns_discovery: false,
dns_hostname: None,
}
}
}
pub struct TopologyDiscovery {
topology: Arc<TopologyManager>,
config: DiscoveryConfig,
}
impl TopologyDiscovery {
pub fn new(topology: Arc<TopologyManager>, config: DiscoveryConfig) -> Self {
Self { topology, config }
}
pub async fn run(&self) {
let mut interval = tokio::time::interval(self.config.refresh_interval);
loop {
interval.tick().await;
let timed_out = self.topology.check_health_timeouts();
for node_id in timed_out {
tracing::warn!("Node {} health timeout", node_id);
}
if let Err(e) = self.refresh_from_seeds().await {
tracing::debug!("Topology refresh from seeds failed: {}", e);
}
if self.config.dns_discovery {
if let Err(e) = self.discover_from_dns().await {
tracing::debug!("DNS discovery failed: {}", e);
}
}
}
}
async fn refresh_from_seeds(&self) -> Result<()> {
use tokio::net::TcpStream;
use tokio::time::timeout;
for seed in &self.config.seed_nodes {
let addr = match seed.parse::<std::net::SocketAddr>() {
Ok(a) => a,
Err(_) => {
match tokio::net::lookup_host(seed).await {
Ok(mut addrs) => match addrs.next() {
Some(a) => a,
None => continue,
},
Err(_) => continue,
}
}
};
let connect_timeout = Duration::from_secs(2);
match timeout(connect_timeout, TcpStream::connect(addr)).await {
Ok(Ok(_stream)) => {
tracing::trace!("Seed node {} is reachable", seed);
let nodes = self.topology.get_all_nodes();
for node in nodes {
if node.client_addr.contains(&addr.ip().to_string())
|| node.replication_addr.contains(&addr.ip().to_string())
{
self.topology.update_health(node.node_id, true);
}
}
}
Ok(Err(e)) => {
tracing::trace!("Seed node {} connection failed: {}", seed, e);
}
Err(_) => {
tracing::trace!("Seed node {} connection timeout", seed);
}
}
}
Ok(())
}
async fn discover_from_dns(&self) -> Result<()> {
use std::net::ToSocketAddrs;
if let Some(hostname) = &self.config.dns_hostname {
match tokio::task::spawn_blocking({
let hostname = hostname.clone();
move || {
let lookup_addr = format!("{}:5433", hostname); lookup_addr.to_socket_addrs()
}
})
.await
{
Ok(Ok(addrs)) => {
for addr in addrs {
tracing::trace!("DNS discovery found address: {}", addr);
}
}
Ok(Err(e)) => {
tracing::trace!("DNS lookup for {} failed: {}", hostname, e);
}
Err(e) => {
tracing::trace!("DNS lookup task failed: {}", e);
}
}
tracing::trace!(
"SRV record lookup for _heliosdb._tcp.{} would be performed with DNS resolver",
hostname
);
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_topology_manager() {
let manager = TopologyManager::new(None);
let primary = NodeInfo::new(
Uuid::new_v4(),
NodeRole::Primary,
"primary:5432".to_string(),
"primary:5433".to_string(),
);
let standby = NodeInfo::new(
Uuid::new_v4(),
NodeRole::Standby,
"standby:5432".to_string(),
"standby:5433".to_string(),
);
manager.register_node(primary.clone());
manager.register_node(standby.clone());
assert_eq!(manager.get_primary_id(), Some(primary.node_id));
assert_eq!(manager.get_healthy_standbys().len(), 1);
assert_eq!(manager.get_read_nodes().len(), 2);
let summary = manager.get_cluster_summary();
assert_eq!(summary.total_nodes, 2);
assert_eq!(summary.primary_count, 1);
assert_eq!(summary.standby_count, 1);
assert!(summary.is_healthy());
}
#[test]
fn test_health_timeout() {
let manager = TopologyManager::new(None).with_health_timeout(Duration::from_millis(10));
let node = NodeInfo::new(
Uuid::new_v4(),
NodeRole::Standby,
"standby:5432".to_string(),
"standby:5433".to_string(),
);
manager.register_node(node.clone());
assert!(manager.get_node(node.node_id).unwrap().is_healthy);
std::thread::sleep(Duration::from_millis(20));
let timed_out = manager.check_health_timeouts();
assert!(timed_out.contains(&node.node_id));
assert!(!manager.get_node(node.node_id).unwrap().is_healthy);
}
}