churn-v2/crates/churn/src/agent/refresh.rs

use std::collections::BTreeMap;

use crate::agent::models::Commands;

use super::{agent_state::AgentState, queue::AgentQueue};

#[derive(Clone)]
pub struct AgentRefresh {
    process_host: String,
    queue: AgentQueue,
}

impl AgentRefresh {
    pub fn new(state: impl Into<AgentState>) -> Self {
        let state: AgentState = state.into();
        Self {
            process_host: state.discovery.process_host.clone(),
            queue: state.queue.clone(),
        }
    }
}

#[async_trait::async_trait]
impl notmad::Component for AgentRefresh {
    fn name(&self) -> Option<String> {
        Some("agent_refresh".into())
    }

    async fn run(
        &self,
        cancellation_token: tokio_util::sync::CancellationToken,
    ) -> Result<(), notmad::MadError> {
        // let cancel =
        //     nodrift::schedule_drifter(std::time::Duration::from_secs(60 * 10), self.clone());
        let cancel = nodrift::schedule_drifter(std::time::Duration::from_secs(5), self.clone());
        tokio::select! {
            _ = cancel.cancelled() => {},
            _ = cancellation_token.cancelled() => {
                tracing::debug!("cancelling agent refresh");
                cancel.cancel();
            }
        }

        Ok(())
    }
}

#[async_trait::async_trait]
impl nodrift::Drifter for AgentRefresh {
    async fn execute(&self, _token: tokio_util::sync::CancellationToken) -> anyhow::Result<()> {
        tracing::info!(process_host = self.process_host, "refreshing agent");

        self.queue
            .publish(Commands::ScheduleTask {
                task: "update".into(),
                properties: BTreeMap::default(),
            })
            .await?;

        Ok(())
    }
}