NxMesh/apps/nxmesh-agent/src/service/nginx_handler/message_handler.rs

use std::sync::Arc;

use dashmap::DashMap;
use nxmesh_proto::{
    ConfigUpdate, ConfigUpdateResult,
    agent_message::Payload::ConfigUpdateResult as ConfigUpdateResultPayload, command::Command,
    command_result,
};
use tracing::{info, warn};

use crate::{
    config::settings::NginxSettings,
    service::{
        master_handler::{
            MasterHandler, MessageResult,
            handlers::{OnCommandHandler, OnConfigUpdateHandler},
        },
        nginx_handler::{command_handler::CommandHandler, fs_handler::FsHandler},
    },
};

const DEFAULT_CONFIG_PATH: &str = "nginx.conf";
const DEFAULT_NGINX_CONFIG_CONTENT: &str = r#"
events {}
"#;

pub trait NginxMasterMessageHandler: Send + Sync +  'static
//
+ OnConfigUpdateHandler
+ OnCommandHandler
{}

pub struct NginxMasterMessageHandlerImpl {
    settings: Arc<NginxSettings>,
    command_handler: Arc<dyn CommandHandler>,
    fs_handler: Arc<dyn FsHandler>,
    master_handler: Arc<dyn MasterHandler>,
    //
    // dash_map for for storing the on-going config updates, with the key as deployment_id, and the value as a tuple of (version_id, timestamp). On-going update must lock the deployment_id, and the new update with newer timestamp will wait until the lock is released. This is to ensure the config updates are applied in order.
    // When the current timestamp is older than the timestamp in the map, the current update must be rejected, and the master should be informed to resend the update with the latest timestamp.
    ongoing_updates: DashMap<String, (String, i64)>,
}

impl NginxMasterMessageHandlerImpl {
    pub fn new(
        settings: Arc<NginxSettings>,
        command_handler: Arc<dyn CommandHandler>,
        fs_handler: Arc<dyn FsHandler>,
        master_handler: Arc<dyn MasterHandler>,
    ) -> Self {
        Self {
            settings,
            command_handler,
            fs_handler,
            master_handler,
            ongoing_updates: DashMap::new(),
        }
    }
}

impl NginxMasterMessageHandler for NginxMasterMessageHandlerImpl {}

#[async_trait::async_trait]
impl OnConfigUpdateHandler for NginxMasterMessageHandlerImpl {
    async fn on_config_update(
        &self,
        agent_id: &str,
        timestamp: i64,
        message_id: &str,
        config_info: ConfigUpdate,
    ) -> MessageResult<()> {
        // TODO: handle concurrency, expect only the latest version with latest timestamp is applied
        // when a newer config update comes in, and the older config update is still being processed. The new config will wait until the old config is applied.
        let deployment_id = format!("{}-{}", config_info.config_id, config_info.version);
        // write the configs
        let root_config_path = match config_info.root_config {
            Some(config_content) => {
                self.fs_handler
                    .write_config(
                        &deployment_id,
                        &config_content.content,
                        &config_content.path,
                    )
                    .await?
            }
            None => {
                // If the config content is not provided, write a default config to ensure the deployment folder is created and can be used for later updates.
                warn!(
                    "Config content is not provided for config update, writing a default minimal config for deployment_id: {}",
                    deployment_id
                );
                self.fs_handler
                    .write_config(
                        &deployment_id,
                        DEFAULT_NGINX_CONFIG_CONTENT,
                        DEFAULT_CONFIG_PATH,
                    )
                    .await?
            }
        };
        //
        for config in config_info.configs {
            self.fs_handler
                .write_config(&deployment_id, &config.content, &config.path)
                .await?;
        }
        // apply reload on the root config
        self.command_handler.reload(Some(&root_config_path)).await?;
        // persist deployment path so Reload/Test commands survive agent restarts
        self.fs_handler.save_last_deployment(&root_config_path).await?;
        info!("Persisted last deployment path: {}", root_config_path);
        // Reply the master to confirm the config update is successful
        self.master_handler
            .send_message_to_master(nxmesh_proto::AgentMessage {
                agent_id: agent_id.to_string(),
                timestamp,
                message_id: message_id.to_string(),
                payload: Some(ConfigUpdateResultPayload(ConfigUpdateResult {
                    success: true,
                    error_message: None,
                    config_id: config_info.config_id,
                    version: config_info.version,
                })),
            })
            .await?;
        //
        Ok(())
    }
}

#[async_trait::async_trait]
impl OnCommandHandler for NginxMasterMessageHandlerImpl {
    async fn on_command(
        &self,
        agent_id: &str,
        timestamp: i64,
        message_id: &str,
        command: Command,
    ) -> MessageResult<()> {
        // execute the command
        let mut agent_message = nxmesh_proto::AgentMessage {
            agent_id: agent_id.to_string(),
            timestamp,
            message_id: message_id.to_string(),
            payload: None,
        };
        // load the last known deployment path for use with Reload/Test commands
        let last_config_path = self.fs_handler.load_last_deployment().await?;

        let result: command_result::Result = match command {
            Command::Reload(_) => {
                let result = self
                    .command_handler
                    .reload(last_config_path.as_deref())
                    .await;
                command_result::Result::ReloadResult(nxmesh_proto::ReloadResult {
                    success: result.is_ok(),
                    error_message: result.err().map(|e| e.to_string()).unwrap_or_default(),
                })
            }
            Command::Test(_) => {
                let result = self
                    .command_handler
                    .validate(last_config_path.as_deref())
                    .await;
                command_result::Result::TestResult(nxmesh_proto::TestResult {
                    success: result.is_ok(),
                    error_message: result.err().map(|e| e.to_string()).unwrap_or_default(),
                })
            }
        };
        // Reply the master to confirm the command execution is successful, and return the command output
        agent_message.payload = Some(nxmesh_proto::agent_message::Payload::CommandResult(
            nxmesh_proto::CommandResult {
                result: Some(result),
            },
        ));

        self.master_handler
            .send_message_to_master(agent_message)
            .await?;
        //
        Ok(())
    }
}