NxMesh/apps/nxmesh-agent/src/service/master_handler/mod.rs

use std::sync::Arc;

use nxmesh_proto::AgentMessage;
use tokio::sync::mpsc;
use tokio_stream::wrappers::ReceiverStream;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};

use crate::{
    connector::master::{MasterConnector, MasterConnectorTrait},
    service::master_handler::handlers::MasterMessageHandler,
};

pub mod handlers;

#[derive(Debug)]
pub enum MasterHandlerError {
    ConnectionError(String),
    // TODO: should be protobuf error to transmit the error to master
    MessageHandlingError(String),
    RetryLimitExceeded(String),
    SendMessageError(String),
}

pub type MessageResult<T> = std::result::Result<T, MasterHandlerError>;

#[async_trait::async_trait]
pub trait MasterHandler: Send + Sync + 'static {
    // Create a new routine to handle incoming messages from the master
    // This method will auto-reconnect if the connection is lost, so it should run indefinitely until the agent is shut down
    async fn start_handle_master_message(&self) -> MessageResult<()>;
    async fn stop_handle_master_message(&self) -> MessageResult<()>;

    // Send a message to the master, response should be handled by the agent message handler registered
    async fn send_message_to_master(&self, message: AgentMessage) -> MessageResult<()>;
}

struct MessageHandleInfo {
    tx: mpsc::Sender<AgentMessage>,
    // used to signal the running handler/connection to stop
    cancel: CancellationToken,
}

pub struct MasterHandlerImpl<MMH>
where
    MMH: MasterMessageHandler + ?Sized,
{
    agent_id: String,
    connector: Arc<MasterConnector>,
    message_handler: Arc<MMH>,
    message_handle_lock: tokio::sync::RwLock<Option<MessageHandleInfo>>,
}

impl<MMH> MasterHandlerImpl<MMH>
where
    MMH: MasterMessageHandler + ?Sized,
{
    pub fn new(agent_id: &str, connector: Arc<MasterConnector>, message_handler: Arc<MMH>) -> Self {
        Self {
            agent_id: agent_id.to_string(),
            connector,
            message_handler,
            message_handle_lock: tokio::sync::RwLock::new(None),
        }
    }
}

#[async_trait::async_trait]
impl<MMH> MasterHandler for MasterHandlerImpl<MMH>
where
    MMH: MasterMessageHandler + ?Sized,
{
    async fn start_handle_master_message(&self) -> MessageResult<()> {
        info!("Starting master message handler...");
        let mut client = self.connector.get_client();

        // ensure only one caller can start the handler
        // create the cancel token for the lifetime of this handler invocation
        let cancel_token = CancellationToken::new();
        {
            let mut guard = self.message_handle_lock.write().await;
            if guard.is_some() {
                warn!("Master message handler is already running");
                return Ok(());
            }
            // placeholder tx; will be replaced per-connection
            let (tx, _rx) = mpsc::channel(1);
            *guard = Some(MessageHandleInfo {
                tx,
                cancel: cancel_token.clone(),
            });
        }

        'connection_loop: loop {
            // fresh outbound channel per connection
            let (tx, rx) = mpsc::channel(32);
            let outbound_stream = ReceiverStream::new(rx);

            // try to connect
            let mut stream = match client.stream(outbound_stream).await {
                Ok(s) => s.into_inner(),
                Err(e) => {
                    error!(
                        "Failed to connect to master: {}. Retrying in 5 seconds...",
                        e
                    );
                    // update stored sender so any callers see the current tx
                    {
                        let mut guard = self.message_handle_lock.write().await;
                        if let Some(info) = guard.as_mut() {
                            info.tx = tx.clone();
                        }
                    }
                    let conn_token = cancel_token.child_token();
                    tokio::select! {
                        _ = conn_token.cancelled() => break 'connection_loop,
                        _ = tokio::time::sleep(std::time::Duration::from_secs(5)) => continue 'connection_loop,
                    }
                }
            };
            // store current tx so senders can use it
            {
                let mut guard = self.message_handle_lock.write().await;
                if let Some(info) = guard.as_mut() {
                    info.tx = tx.clone();
                }
            }
            // connection-level token to observe stop requests
            let conn_token = cancel_token.child_token();
            info!("Connected to master, starting to receive messages...");
            // process messages inline so we can clear the slot on exit
            'message_processing: loop {
                tokio::select! {
                    _ = conn_token.cancelled() => {
                        info!("Stop requested for master handler");
                        break 'connection_loop;
                    }
                    message = stream.message() => {
                        match message {
                            Ok(Some(msg)) => {
                                if let Err(e) = self.message_handler.handle_master_message(&self.agent_id, msg).await {
                                    error!("Failed to handle master message: {:?}", e);
                                }
                                continue;
                            }
                            Ok(None) => {
                                warn!("Master closed the connection");
                                break 'message_processing;
                            }
                            Err(e) => {
                                error!("Error receiving message from master: {:?}", e);
                                break 'message_processing;
                            }
                        }
                    }
                }
            }

            // connection ended — clear stored info
            {
                let mut guard = self.message_handle_lock.write().await;
                guard.take();
            }

            // if stop requested, exit
            if cancel_token.is_cancelled() {
                break 'connection_loop;
            }

            // otherwise reconnect after backoff
            tokio::time::sleep(std::time::Duration::from_secs(5)).await;
        }

        // final cleanup
        let mut guard = self.message_handle_lock.write().await;
        guard.take();
        Ok(())
    }

    async fn stop_handle_master_message(&self) -> MessageResult<()> {
        // Signal the running handler to stop and wait for it to clear
        let mut maybe_cancel = None;
        {
            let mut guard = self.message_handle_lock.write().await;
            if let Some(info) = guard.take() {
                maybe_cancel = Some(info.cancel);
            }
        }

        if let Some(cancel) = maybe_cancel {
            cancel.cancel();

            // wait for the handler to clear (with timeout)
            for _ in 0..50 {
                if self.message_handle_lock.read().await.is_none() {
                    info!("Master message handler task stopped successfully");
                    return Ok(());
                }
                tokio::time::sleep(std::time::Duration::from_millis(100)).await;
            }
            warn!("Timed out waiting for master message handler to stop");
        } else {
            warn!("Master message handler is not running");
        }
        Ok(())
    }

    async fn send_message_to_master(&self, message: AgentMessage) -> MessageResult<()> {
        let guard = self.message_handle_lock.read().await;
        if let Some(handle_info) = guard.as_ref() {
            handle_info.tx.send(message).await.map_err(|e| {
                MasterHandlerError::SendMessageError(format!(
                    "Failed to send message to master: {}",
                    e
                ))
            })?;
        } else {
            return Err(MasterHandlerError::SendMessageError(
                "Master message handler is not running".to_string(),
            ));
        }
        Ok(())
    }
}