use std::{collections::BTreeMap, net::SocketAddr, pin::Pin}; use anyhow::Context; use futures::Stream; use notmad::{Component, MadError}; use tonic::transport::Server; use crate::{agent::models::Commands, grpc::*}; #[derive(Clone)] pub struct GrpcServer { grpc_host: SocketAddr, } impl GrpcServer { pub fn new(grpc_host: SocketAddr) -> Self { Self { grpc_host } } } #[async_trait::async_trait] impl Component for GrpcServer { async fn run( &self, cancellation_token: tokio_util::sync::CancellationToken, ) -> Result<(), notmad::MadError> { let task = Server::builder() .add_service(crate::grpc::churn_server::ChurnServer::new(self.clone())) .serve(self.grpc_host); tokio::select! { _ = cancellation_token.cancelled() => {}, res = task => { res.context("failed to run grpc server").map_err(MadError::Inner)?; } } Ok(()) } } #[async_trait::async_trait] impl crate::grpc::churn_server::Churn for GrpcServer { async fn get_key( &self, request: tonic::Request, ) -> std::result::Result, tonic::Status> { todo!() } async fn set_key( &self, request: tonic::Request, ) -> std::result::Result, tonic::Status> { todo!() } #[doc = " Server streaming response type for the ListenEvents method."] type ListenEventsStream = Pin> + Send>>; async fn listen_events( &self, request: tonic::Request, ) -> std::result::Result, tonic::Status> { let (tx, rx) = tokio::sync::mpsc::channel(128); tokio::spawn(async move { let mut interval = tokio::time::interval(std::time::Duration::from_secs(60 * 10)); loop { interval.tick().await; let Ok(schedule_task) = serde_json::to_string(&Commands::ScheduleTask { task: "refresh".into(), properties: BTreeMap::default(), }) else { tracing::warn!("failed to serialize event"); continue; }; if let Err(e) = tx .send(Ok(ListenEventsResponse { id: uuid::Uuid::new_v4().to_string(), value: schedule_task, })) .await { tracing::warn!("failed to send response: {}", e); break; } } }); let stream = futures::stream::unfold(rx, |mut msg| async move { let next = msg.recv().await?; Some((next, msg)) }); Ok(tonic::Response::new(Box::pin(stream))) } }