use crate::prelude::now_micros;
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{debug, info, warn};
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopicInfo {
pub topic: String,
pub message_type: String,
pub publishers: Vec<String>,
pub subscribers: Vec<String>,
pub created_at: u64,
pub last_publisher_added: Option<u64>,
pub last_subscriber_added: Option<u64>,
pub description: Option<String>,
pub metadata: HashMap<String, String>,
}
impl TopicInfo {
pub fn new(topic: &str, message_type: &str) -> Self {
Self {
topic: topic.to_string(),
message_type: message_type.to_string(),
publishers: Vec::new(),
subscribers: Vec::new(),
created_at: now_micros(),
last_publisher_added: None,
last_subscriber_added: None,
description: None,
metadata: HashMap::new(),
}
}
pub fn add_publisher(&mut self, node_id: &str) {
if !self.publishers.contains(&node_id.to_string()) {
self.publishers.push(node_id.to_string());
self.last_publisher_added = Some(now_micros());
}
}
pub fn add_subscriber(&mut self, node_id: &str) {
if !self.subscribers.contains(&node_id.to_string()) {
self.subscribers.push(node_id.to_string());
self.last_subscriber_added = Some(now_micros());
}
}
pub fn remove_publisher(&mut self, node_id: &str) {
self.publishers.retain(|p| p != node_id);
}
pub fn remove_subscriber(&mut self, node_id: &str) {
self.subscribers.retain(|s| s != node_id);
}
pub fn is_orphaned(&self) -> bool {
self.publishers.is_empty() || self.subscribers.is_empty()
}
pub fn publisher_count(&self) -> usize {
self.publishers.len()
}
pub fn subscriber_count(&self) -> usize {
self.subscribers.len()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Role {
Publisher,
Subscriber,
}
impl Role {
fn label(&self) -> &'static str {
match self {
Role::Publisher => "publisher",
Role::Subscriber => "subscriber",
}
}
}
#[derive(Clone)]
pub struct TopicRegistry {
topics: Arc<RwLock<HashMap<String, TopicInfo>>>,
node_publishers: Arc<RwLock<HashMap<String, HashSet<String>>>>,
node_subscribers: Arc<RwLock<HashMap<String, HashSet<String>>>>,
}
impl TopicRegistry {
pub fn new() -> Self {
info!("Creating new topic registry");
Self {
topics: Arc::new(RwLock::new(HashMap::new())),
node_publishers: Arc::new(RwLock::new(HashMap::new())),
node_subscribers: Arc::new(RwLock::new(HashMap::new())),
}
}
pub async fn register_publisher(&self, topic: &str, node_id: &str, message_type: &str) {
self.register(Role::Publisher, topic, node_id, message_type).await;
}
pub async fn register_subscriber(&self, topic: &str, node_id: &str, message_type: &str) {
self.register(Role::Subscriber, topic, node_id, message_type).await;
}
pub async fn unregister_publisher(&self, topic: &str, node_id: &str) {
self.unregister(Role::Publisher, topic, node_id).await;
}
pub async fn unregister_subscriber(&self, topic: &str, node_id: &str) {
self.unregister(Role::Subscriber, topic, node_id).await;
}
async fn register(&self, role: Role, topic: &str, node_id: &str, message_type: &str) {
let mut topics = self.topics.write().await;
let topic_info = topics
.entry(topic.to_string())
.or_insert_with(|| TopicInfo::new(topic, message_type));
match role {
Role::Publisher => topic_info.add_publisher(node_id),
Role::Subscriber => topic_info.add_subscriber(node_id),
}
let node_map = match role {
Role::Publisher => &self.node_publishers,
Role::Subscriber => &self.node_subscribers,
};
let mut node_entries = node_map.write().await;
node_entries
.entry(node_id.to_string())
.or_default()
.insert(topic.to_string());
debug!(
"Registered {}: node='{}', topic='{}', type='{}'",
role.label(),
node_id,
topic,
message_type
);
}
async fn unregister(&self, role: Role, topic: &str, node_id: &str) {
let mut topics = self.topics.write().await;
if let Some(info) = topics.get_mut(topic) {
match role {
Role::Publisher => info.remove_publisher(node_id),
Role::Subscriber => info.remove_subscriber(node_id),
}
if info.publishers.is_empty() && info.subscribers.is_empty() {
topics.remove(topic);
debug!("Removed topic '{}' (no publishers or subscribers)", topic);
}
}
let node_map = match role {
Role::Publisher => &self.node_publishers,
Role::Subscriber => &self.node_subscribers,
};
let mut node_entries = node_map.write().await;
if let Some(topics_set) = node_entries.get_mut(node_id) {
topics_set.remove(topic);
if topics_set.is_empty() {
node_entries.remove(node_id);
}
}
debug!("Unregistered {}: node='{}', topic='{}'", role.label(), node_id, topic);
}
pub async fn unregister_node(&self, node_id: &str) {
let topics_published = {
let node_pubs = self.node_publishers.read().await;
node_pubs.get(node_id).cloned().unwrap_or_default()
};
for topic in topics_published {
self.unregister_publisher(&topic, node_id).await;
}
let topics_subscribed = {
let node_subs = self.node_subscribers.read().await;
node_subs.get(node_id).cloned().unwrap_or_default()
};
for topic in topics_subscribed {
self.unregister_subscriber(&topic, node_id).await;
}
info!("Unregistered all topics for node '{}'", node_id);
}
pub async fn get_topic(&self, topic: &str) -> Option<TopicInfo> {
let topics = self.topics.read().await;
topics.get(topic).cloned()
}
pub async fn list_topics(&self) -> Vec<String> {
let topics = self.topics.read().await;
topics.keys().cloned().collect()
}
pub async fn get_all_topics(&self) -> Vec<TopicInfo> {
let topics = self.topics.read().await;
topics.values().cloned().collect()
}
pub async fn find_by_message_type(&self, message_type: &str) -> Vec<TopicInfo> {
let topics = self.topics.read().await;
topics
.values()
.filter(|info| info.message_type == message_type)
.cloned()
.collect()
}
pub async fn topics_published_by(&self, node_id: &str) -> Vec<String> {
let node_pubs = self.node_publishers.read().await;
node_pubs
.get(node_id)
.cloned()
.unwrap_or_default()
.into_iter()
.collect()
}
pub async fn topics_subscribed_by(&self, node_id: &str) -> Vec<String> {
let node_subs = self.node_subscribers.read().await;
node_subs
.get(node_id)
.cloned()
.unwrap_or_default()
.into_iter()
.collect()
}
pub async fn topic_count(&self) -> usize {
let topics = self.topics.read().await;
topics.len()
}
pub async fn validate(&self) -> Vec<String> {
let mut issues = Vec::new();
let topics = self.topics.read().await;
for (topic, info) in topics.iter() {
if info.publishers.is_empty() {
issues.push(format!(
"Topic '{}' has no publishers (subscribed by: {:?})",
topic, info.subscribers
));
}
if info.subscribers.is_empty() {
issues.push(format!(
"Topic '{}' has no subscribers (published by: {:?})",
topic, info.publishers
));
}
}
if !issues.is_empty() {
warn!("Found {} topology issues", issues.len());
}
issues
}
pub async fn stats(&self) -> TopologyStats {
let topics = self.topics.read().await;
let total_topics = topics.len();
let total_publishers: usize = topics.values().map(|t| t.publishers.len()).sum();
let total_subscribers: usize = topics.values().map(|t| t.subscribers.len()).sum();
let orphaned_topics = topics.values().filter(|t| t.is_orphaned()).count();
let topics_without_publishers = topics.values().filter(|t| t.publishers.is_empty()).count();
let topics_without_subscribers = topics.values().filter(|t| t.subscribers.is_empty()).count();
TopologyStats {
total_topics,
total_publishers,
total_subscribers,
orphaned_topics,
topics_without_publishers,
topics_without_subscribers,
}
}
pub async fn export_dot(&self) -> String {
let mut dot = String::from("digraph TopicTopology {\n");
dot.push_str(" rankdir=LR;\n");
dot.push_str(" node [shape=box];\n\n");
let topics = self.topics.read().await;
for (topic, info) in topics.iter() {
let label = format!("{}\\n({})", topic, info.message_type);
dot.push_str(&format!(
" \"{}\" [label=\"{}\", shape=ellipse, style=filled, fillcolor=lightblue];\n",
topic, label
));
}
dot.push('\n');
for (topic, info) in topics.iter() {
for pub_node in &info.publishers {
dot.push_str(&format!(
" \"{}\" -> \"{}\" [label=\"pub\", color=green];\n",
pub_node, topic
));
}
for sub_node in &info.subscribers {
dot.push_str(&format!(
" \"{}\" -> \"{}\" [label=\"sub\", color=blue];\n",
topic, sub_node
));
}
}
dot.push_str("}\n");
dot
}
pub async fn clear(&self) {
let mut topics = self.topics.write().await;
topics.clear();
let mut node_pubs = self.node_publishers.write().await;
node_pubs.clear();
let mut node_subs = self.node_subscribers.write().await;
node_subs.clear();
info!("Cleared topic registry");
}
}
impl Default for TopicRegistry {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TopologyStats {
pub total_topics: usize,
pub total_publishers: usize,
pub total_subscribers: usize,
pub orphaned_topics: usize,
pub topics_without_publishers: usize,
pub topics_without_subscribers: usize,
}