mecha10-core 0.8.0

//! 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 _;

pub trait DiscoveryExt {
    /// Register this node in the service registry
    ///
    /// # Arguments
    ///
    /// * `metadata` - Node metadata to register
    ///
    /// # Example
    ///
    /// ```rust
    /// use mecha10_core::prelude::*;
    /// use mecha10_core::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 Context {
    /// Resolve a `Vec<NodeMetadata>` from the node IDs stored in a Redis set at `key`.
    ///
    /// Shared by `discover_by_type`, `discover_publishers`, and `discover_subscribers`, which
    /// only differ in the Redis key they query and the connection/error-context strings used.
    #[cfg(feature = "messaging")]
    async fn resolve_nodes_by_key(&self, key: &str, conn_label: &str, err_context: &str) -> Result<Vec<NodeMetadata>> {
        use redis::AsyncCommands;

        let mut conn = get_redis_connection(conn_label).await?;

        let node_ids: Vec<String> = conn
            .smembers(key)
            .await
            .with_context(|| err_context.to_string())
            .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)
    }
}

macro_rules! discovery_lookup {
    ($fn_name:ident, $arg:ident, $key_fn:expr, $context:literal, $err_fmt:literal) => {
        async fn $fn_name(&self, $arg: &str) -> Result<Vec<NodeMetadata>> {
            #[cfg(feature = "messaging")]
            {
                let key = $key_fn($arg);
                self.resolve_nodes_by_key(&key, $context, &format!($err_fmt, $arg))
                    .await
            }

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

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

        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();

            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)))?;

            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(),
                })?;

            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();
            }

            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();
            }

            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?;

            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<()> {
        if let Some(mut metadata) = self.get_node(self.node_id()).await? {
            metadata.last_heartbeat = current_timestamp();
            self.register(metadata).await
        } else {
            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?;

            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()))
        }
    }

    discovery_lookup!(
        discover_by_type,
        node_type,
        RedisKeys::by_type,
        "type discovery",
        "Failed to query nodes of type '{}'"
    );

    discovery_lookup!(
        discover_publishers,
        topic,
        RedisKeys::publishers,
        "publisher discovery",
        "Failed to query publishers for topic '{}'"
    );

    discovery_lookup!(
        discover_subscribers,
        topic,
        RedisKeys::subscribers,
        "subscriber discovery",
        "Failed to query subscribers for topic '{}'"
    );

    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()))
        }
    }
}