refactor: streamline master message handling by removing unnecessary Arc and Mutex usage, and reorganizing handler structure

apps/nxmesh-agent/src/service/master_handler/handlers.rs

@@ -0,0 +1,62 @@
+use std::sync::Arc;
+
+use nxmesh_proto::{AgentMessage, ConfigUpdate, MasterMessage, master_message::Payload};
+
+use crate::service::master_handler::{MasterHandlerError, MessageResult};
+
+#[async_trait::async_trait]
+pub trait MasterMessageHandler: Send + Sync + 'static {
+    async fn handle_master_message(&self, message: MasterMessage) -> MessageResult<()>;
+}
+
+#[async_trait::async_trait]
+pub trait OnConfigUpdateHandler: Send + Sync + 'static {
+    // Handle the config update message from master, write the config content to files, validate the new config and reload nginx
+    async fn on_config_update(&self, config_info: ConfigUpdate) -> MessageResult<()>;
+}
+
+pub struct HandlerImpl<OCH>
+where
+    OCH: OnConfigUpdateHandler,
+{
+    on_config_update_handler: Arc<OCH>,
+}
+
+impl<OCH> HandlerImpl<OCH>
+where
+    OCH: OnConfigUpdateHandler,
+{
+    pub fn new(on_config_update_handler: Arc<OCH>) -> Self {
+        Self {
+            on_config_update_handler,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<OCH> MasterMessageHandler for HandlerImpl<OCH>
+where
+    OCH: OnConfigUpdateHandler,
+{
+    async fn handle_master_message(&self, message: MasterMessage) -> MessageResult<()> {
+        match message.payload {
+            Some(Payload::ConfigUpdate(config_info)) => {
+                self.on_config_update_handler
+                    .on_config_update(config_info)
+                    .await
+            }
+            Some(_) => {
+                // We should never receive other types of messages from the master, but we should handle it anyway
+                Err(MasterHandlerError::MessageHandlingError(
+                    "Received unsupported master message type".to_string(),
+                ))
+            }
+            None => {
+                // This should never happen as the master should always send a valid message, but we should handle it anyway
+                return Err(MasterHandlerError::MessageHandlingError(
+                    "Received master message with empty payload".to_string(),
+                ));
+            }
+        }
+    }
+}

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

@@ -0,0 +1,163 @@
+use std::sync::Arc;
+
+use nxmesh_proto::AgentMessage;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+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 {
+    handle: tokio::task::JoinHandle<()>,
+    tx: mpsc::Sender<AgentMessage>,
+}
+
+pub struct MasterHandlerImpl<MMH>
+where
+    MMH: MasterMessageHandler,
+{
+    connector: Arc<MasterConnector>,
+    message_handler: Arc<MMH>,
+    message_handle_lock: tokio::sync::RwLock<Option<MessageHandleInfo>>,
+}
+
+impl<MMH> MasterHandlerImpl<MMH>
+where
+    MMH: MasterMessageHandler,
+{
+    pub fn new(connector: Arc<MasterConnector>, message_handler: Arc<MMH>) -> Self {
+        Self {
+            connector,
+            message_handler,
+            message_handle_lock: tokio::sync::RwLock::new(None),
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<MMH> MasterHandler for MasterHandlerImpl<MMH>
+where
+    MMH: MasterMessageHandler,
+{
+    async fn start_handle_master_message(&self) -> MessageResult<()> {
+        let mut client = self.connector.get_client();
+
+        'connection_loop: loop {
+            let guard_result = self.message_handle_lock.try_write();
+            let mut guard = match guard_result {
+                Ok(g) if g.is_none() => g,
+                Ok(_) => {
+                    warn!("Master message handler is already running");
+                    return Ok(());
+                }
+                Err(e) => {
+                    return Err(MasterHandlerError::MessageHandlingError(format!(
+                        "Failed to acquire lock for message handler: {}",
+                        e
+                    )));
+                }
+            };
+            let (tx, rx) = mpsc::channel(32);
+            let outbound_stream = ReceiverStream::new(rx);
+            // 2. Connect to the master and start the bi-directional streaming RPC
+            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
+                    );
+                    tokio::time::sleep(std::time::Duration::from_secs(5)).await;
+                    continue 'connection_loop;
+                }
+            };
+
+            // 3. Spawn a task to handle incoming messages from the master
+            let message_handler = self.message_handler.clone();
+            let handle = tokio::spawn(async move {
+                'message_loop: loop {
+                    let message = stream.message().await;
+                    match message {
+                        Ok(Some(msg)) => {
+                            if let Err(e) = message_handler.handle_master_message(msg).await {
+                                error!("Failed to handle master message: {:?}", e);
+                            }
+                            continue 'message_loop;
+                        }
+                        Ok(None) => {
+                            warn!("Master closed the connection");
+                            return;
+                        }
+                        Err(e) => {
+                            error!("Error receiving message from master: {:?}", e);
+                            return;
+                        }
+                    }
+                }
+            });
+            *guard = Some(MessageHandleInfo { handle, tx });
+        }
+    }
+
+    async fn stop_handle_master_message(&self) -> MessageResult<()> {
+        // 1. signal the task to stop (e.g. using a cancellation token or channel)
+        // 2. wait for the task to finish and handle any errors
+        // 3. set handle_master_message_task to None
+
+        let mut guard = self.message_handle_lock.write().await;
+        if let Some(handle_info) = guard.take() {
+            handle_info.handle.abort();
+            match handle_info.handle.await {
+                Ok(_) => info!("Master message handler task stopped successfully"),
+                Err(e) => error!("Failed to stop master message handler task: {:?}", e),
+            }
+        } 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(())
+    }
+}