mecha10_core/discovery/

registry.rs

//! Service registry implementation

use super::metadata::{current_timestamp, NodeMetadata};
use super::redis_client::{get_redis_connection, get_ttl, get_ttl_expire, RedisKeys};
use crate::context::Context;
use crate::error::{Mecha10Error, Result};
use anyhow::Context as _;

/// Service discovery extension trait for Context
pub trait DiscoveryExt {
    /// Register this node in the service registry
    ///
    /// # Arguments
    ///
    /// * `metadata` - Node metadata to register
    ///
    /// # Example
    ///
    /// ```rust
    /// use mecha10::prelude::*;
    /// use mecha10::discovery::{DiscoveryExt, NodeMetadata};
    ///
    /// # async fn example(ctx: &Context) -> Result<()> {
    /// ctx.register(NodeMetadata::new("my-node")
    ///     .with_type("perception")
    ///     .publishes("/perception/detections")
    ///     .subscribes("/sensors/camera/rgb")
    /// ).await?;
    /// # Ok(())
    /// # }
    /// ```
    fn register(&self, metadata: NodeMetadata) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Unregister this node from the service registry
    fn unregister(&self) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Send a heartbeat to update last seen timestamp
    fn heartbeat(&self) -> impl std::future::Future<Output = Result<()>> + Send;

    /// Discover all registered nodes
    fn discover_all(&self) -> impl std::future::Future<Output = Result<Vec<NodeMetadata>>> + Send;

    /// Discover nodes by type
    ///
    /// # Arguments
    ///
    /// * `node_type` - Type of nodes to find (e.g., "camera_driver")
    fn discover_by_type(&self, node_type: &str) -> impl std::future::Future<Output = Result<Vec<NodeMetadata>>> + Send;

    /// Discover who publishes to a specific topic
    ///
    /// # Arguments
    ///
    /// * `topic` - Topic path to query
    fn discover_publishers(&self, topic: &str) -> impl std::future::Future<Output = Result<Vec<NodeMetadata>>> + Send;

    /// Discover who subscribes to a specific topic
    ///
    /// # Arguments
    ///
    /// * `topic` - Topic path to query
    fn discover_subscribers(&self, topic: &str) -> impl std::future::Future<Output = Result<Vec<NodeMetadata>>> + Send;

    /// Get metadata for a specific node
    ///
    /// # Arguments
    ///
    /// * `node_id` - Node identifier to query
    fn get_node(&self, node_id: &str) -> impl std::future::Future<Output = Result<Option<NodeMetadata>>> + Send;
}

impl DiscoveryExt for Context {
    async fn register(&self, mut metadata: NodeMetadata) -> Result<()> {
        // Ensure node_id matches context
        if metadata.node_id.is_empty() {
            metadata.node_id = self.node_id().to_string();
        }

        // Update timestamps
        metadata.last_heartbeat = current_timestamp();
        if metadata.started_at == 0 {
            metadata.started_at = current_timestamp();
        }

        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("node registration").await?;
            let ttl = get_ttl();
            let ttl_expire = get_ttl_expire();

            // Serialize metadata
            let metadata_json = serde_json::to_string(&metadata)
                .with_context(|| format!("Failed to serialize metadata for node '{}'", metadata.node_id))
                .map_err(|e| Mecha10Error::Other(format!("{:#}", e)))?;

            // Store in Redis with TTL
            let key = RedisKeys::node(&metadata.node_id);
            conn.set_ex::<_, _, ()>(&key, &metadata_json, ttl)
                .await
                .with_context(|| format!("Failed to register node '{}' in Redis", metadata.node_id))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            // Add to type index
            if !metadata.node_type.is_empty() {
                let type_key = RedisKeys::by_type(&metadata.node_type);
                conn.sadd::<_, _, ()>(&type_key, &metadata.node_id)
                    .await
                    .with_context(|| {
                        format!(
                            "Failed to update type index for node '{}' of type '{}'",
                            metadata.node_id, metadata.node_type
                        )
                    })
                    .map_err(|e| Mecha10Error::MessagingError {
                        message: format!("{:#}", e),
                        suggestion: "Check Redis connection".to_string(),
                    })?;
                conn.expire::<_, ()>(&type_key, ttl_expire).await.ok();
            }

            // Add to publisher index
            for topic in &metadata.publishes {
                let pub_key = RedisKeys::publishers(topic);
                conn.sadd::<_, _, ()>(&pub_key, &metadata.node_id)
                    .await
                    .with_context(|| {
                        format!(
                            "Failed to update publisher index for topic '{}' on node '{}'",
                            topic, metadata.node_id
                        )
                    })
                    .map_err(|e| Mecha10Error::MessagingError {
                        message: format!("{:#}", e),
                        suggestion: "Check Redis connection".to_string(),
                    })?;
                conn.expire::<_, ()>(&pub_key, ttl_expire).await.ok();
            }

            // Add to subscriber index
            for topic in &metadata.subscribes {
                let sub_key = RedisKeys::subscribers(topic);
                conn.sadd::<_, _, ()>(&sub_key, &metadata.node_id)
                    .await
                    .with_context(|| {
                        format!(
                            "Failed to update subscriber index for topic '{}' on node '{}'",
                            topic, metadata.node_id
                        )
                    })
                    .map_err(|e| Mecha10Error::MessagingError {
                        message: format!("{:#}", e),
                        suggestion: "Check Redis connection".to_string(),
                    })?;
                conn.expire::<_, ()>(&sub_key, ttl_expire).await.ok();
            }

            Ok(())
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn unregister(&self) -> Result<()> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("node unregistration").await?;

            // Delete node metadata
            let key = RedisKeys::node(self.node_id());
            conn.del::<_, ()>(&key)
                .await
                .with_context(|| format!("Failed to unregister node '{}' from Redis", self.node_id()))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            Ok(())
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn heartbeat(&self) -> Result<()> {
        // Get current metadata and re-register to update TTL
        if let Some(mut metadata) = self.get_node(self.node_id()).await? {
            metadata.last_heartbeat = current_timestamp();
            self.register(metadata).await
        } else {
            // Node not registered, register with defaults
            self.register(NodeMetadata::new(self.node_id())).await
        }
    }

    async fn discover_all(&self) -> Result<Vec<NodeMetadata>> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("node discovery").await?;

            // Scan for all node keys
            let pattern = RedisKeys::all_nodes_pattern();
            let keys: Vec<String> = conn
                .keys(pattern)
                .await
                .with_context(|| format!("Failed to query all nodes with pattern '{}'", pattern))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            let mut nodes = Vec::new();
            for key in keys {
                let json: String = conn
                    .get(&key)
                    .await
                    .with_context(|| format!("Failed to get node metadata from key '{}'", key))
                    .map_err(|e| Mecha10Error::MessagingError {
                        message: format!("{:#}", e),
                        suggestion: "Check Redis connection".to_string(),
                    })?;

                if let Ok(metadata) = serde_json::from_str::<NodeMetadata>(&json) {
                    nodes.push(metadata);
                }
            }

            Ok(nodes)
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn discover_by_type(&self, node_type: &str) -> Result<Vec<NodeMetadata>> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("type discovery").await?;

            // Get node IDs from type index
            let type_key = RedisKeys::by_type(node_type);
            let node_ids: Vec<String> = conn
                .smembers(&type_key)
                .await
                .with_context(|| format!("Failed to query nodes of type '{}'", node_type))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            let mut nodes = Vec::new();
            for node_id in node_ids {
                if let Some(metadata) = self.get_node(&node_id).await? {
                    nodes.push(metadata);
                }
            }

            Ok(nodes)
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn discover_publishers(&self, topic: &str) -> Result<Vec<NodeMetadata>> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("publisher discovery").await?;

            let pub_key = RedisKeys::publishers(topic);
            let node_ids: Vec<String> = conn
                .smembers(&pub_key)
                .await
                .with_context(|| format!("Failed to query publishers for topic '{}'", topic))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            let mut nodes = Vec::new();
            for node_id in node_ids {
                if let Some(metadata) = self.get_node(&node_id).await? {
                    nodes.push(metadata);
                }
            }

            Ok(nodes)
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn discover_subscribers(&self, topic: &str) -> Result<Vec<NodeMetadata>> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("subscriber discovery").await?;

            let sub_key = RedisKeys::subscribers(topic);
            let node_ids: Vec<String> = conn
                .smembers(&sub_key)
                .await
                .with_context(|| format!("Failed to query subscribers for topic '{}'", topic))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            let mut nodes = Vec::new();
            for node_id in node_ids {
                if let Some(metadata) = self.get_node(&node_id).await? {
                    nodes.push(metadata);
                }
            }

            Ok(nodes)
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }

    async fn get_node(&self, node_id: &str) -> Result<Option<NodeMetadata>> {
        #[cfg(feature = "messaging")]
        {
            use redis::AsyncCommands;

            let mut conn = get_redis_connection("node lookup").await?;

            let key = RedisKeys::node(node_id);
            let json: Option<String> = conn
                .get(&key)
                .await
                .with_context(|| format!("Failed to get metadata for node '{}'", node_id))
                .map_err(|e| Mecha10Error::MessagingError {
                    message: format!("{:#}", e),
                    suggestion: "Check Redis connection".to_string(),
                })?;

            if let Some(json) = json {
                let metadata = serde_json::from_str::<NodeMetadata>(&json)
                    .with_context(|| format!("Failed to parse metadata JSON for node '{}'", node_id))
                    .map_err(|e| Mecha10Error::Other(format!("{:#}", e)))?;
                Ok(Some(metadata))
            } else {
                Ok(None)
            }
        }

        #[cfg(not(feature = "messaging"))]
        {
            Err(Mecha10Error::Other("Messaging feature not enabled".to_string()))
        }
    }
}