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

use agent_state::AgentState;
use refresh::AgentRefresh;
mod agent_state;
mod refresh;
mod grpc_client {
    use tonic::transport::{Channel, ClientTlsConfig};

    use crate::grpc::{churn_client::ChurnClient, *};

    pub struct GrpcClient {
        host: String,
    }

    impl GrpcClient {
        pub fn new(host: impl Into<String>) -> Self {
            Self { host: host.into() }
        }

        pub async fn get_key(
            &self,
            namespace: &str,
            id: Option<impl Into<String>>,
            key: &str,
        ) -> anyhow::Result<Option<String>> {
            let mut client = self.client().await?;

            let resp = client
                .get_key(GetKeyRequest {
                    key: key.into(),
                    namespace: namespace.into(),
                    id: id.map(|i| i.into()),
                })
                .await?;
            let resp = resp.into_inner();

            Ok(resp.value)
        }

        pub async fn set_key(
            &self,
            namespace: &str,
            id: Option<impl Into<String>>,
            key: &str,
            value: &str,
        ) -> anyhow::Result<()> {
            let mut client = self.client().await?;

            client
                .set_key(SetKeyRequest {
                    key: key.into(),
                    value: value.into(),
                    namespace: namespace.into(),
                    id: id.map(|i| i.into()),
                })
                .await?;

            Ok(())
        }

        async fn client(&self) -> anyhow::Result<ChurnClient<tonic::transport::Channel>> {
            let channel = Channel::from_shared(self.host.to_owned())?
                .tls_config(ClientTlsConfig::new().with_native_roots())?
                .connect()
                .await?;

            let client = ChurnClient::new(channel);

            Ok(client)
        }
    }
}

pub async fn execute(host: impl Into<String>) -> anyhow::Result<()> {
    let state = AgentState::new().await?;

    notmad::Mad::builder()
        .add(AgentRefresh::new(&state, host))
        .cancellation(Some(std::time::Duration::from_secs(2)))
        .run()
        .await?;

    Ok(())
}