churn/crates/churn-server/src/agent.rs

use std::collections::HashMap;

use std::sync::Arc;

use axum::async_trait;

use churn_domain::ServerEnrollReq;
use tokio::sync::Mutex;

use crate::Agent;

#[derive(Clone)]
pub struct AgentService(Arc<dyn AgentServiceTrait + Send + Sync + 'static>);

impl std::ops::Deref for AgentService {
    type Target = Arc<dyn AgentServiceTrait + Send + Sync + 'static>;

    fn deref(&self) -> &Self::Target {
        &self.0
    }
}

impl Default for AgentService {
    fn default() -> Self {
        Self(Arc::new(DefaultAgentService::default()))
    }
}

#[derive(Default)]
struct DefaultAgentService {
    agents: Arc<Mutex<HashMap<String, Agent>>>,
}

#[async_trait]
pub trait AgentServiceTrait {
    async fn enroll(&self, req: ServerEnrollReq) -> anyhow::Result<String>;
}

#[async_trait]
impl AgentServiceTrait for DefaultAgentService {
    async fn enroll(&self, req: ServerEnrollReq) -> anyhow::Result<String> {
        let agent_name = req.agent_name;

        let mut agents = self.agents.lock().await;

        match agents.insert(
            agent_name.clone(),
            Agent {
                name: agent_name.clone(),
            },
        ) {
            Some(_) => {
                tracing::debug!("agents store already contained agent, replaced existing");

                Ok(agent_name)
            }
            None => {
                tracing::debug!("agents store didn't contain agent, inserted");

                Ok(agent_name)
            }
        }
    }
}