//! Watch service implementation use crate::conversions::make_header; use crate::proto::{ watch_server::Watch, WatchRequest, WatchResponse, }; use chainfire_watch::{WatchRegistry, WatchStream}; use std::pin::Pin; use std::sync::Arc; use tokio::sync::mpsc; use tokio_stream::{wrappers::ReceiverStream, StreamExt}; use tonic::{Request, Response, Status, Streaming}; use tracing::{debug, warn}; /// Watch service implementation pub struct WatchServiceImpl { /// Watch registry registry: Arc, /// Cluster ID cluster_id: u64, /// Member ID member_id: u64, } impl WatchServiceImpl { /// Create a new watch service pub fn new(registry: Arc, cluster_id: u64, member_id: u64) -> Self { Self { registry, cluster_id, member_id, } } fn make_header(&self, revision: u64) -> crate::proto::ResponseHeader { make_header(self.cluster_id, self.member_id, revision, 0) } } #[tonic::async_trait] impl Watch for WatchServiceImpl { type WatchStream = Pin> + Send>>; async fn watch( &self, request: Request>, ) -> Result, Status> { let mut in_stream = request.into_inner(); let registry = Arc::clone(&self.registry); let cluster_id = self.cluster_id; let member_id = self.member_id; // Channel for sending responses back to client let (tx, rx) = mpsc::channel(128); let tx_for_events = tx.clone(); // Channel for watch events let (event_tx, mut event_rx) = mpsc::channel::(128); // Spawn task to handle the bidirectional stream tokio::spawn(async move { let mut stream = WatchStream::new(Arc::clone(®istry), { let event_tx = event_tx.clone(); let (watch_tx, mut watch_rx) = mpsc::channel(64); // Forward internal watch responses to proto responses tokio::spawn(async move { while let Some(resp) = watch_rx.recv().await { let proto_resp = internal_to_proto_response(resp, cluster_id, member_id); if event_tx.send(proto_resp).await.is_err() { break; } } }); watch_tx }); while let Some(result) = in_stream.next().await { match result { Ok(req) => { if let Some(request_union) = req.request_union { let response = match request_union { crate::proto::watch_request::RequestUnion::CreateRequest(create) => { let internal_req: chainfire_types::watch::WatchRequest = create.into(); let resp = stream.create_watch(internal_req); internal_to_proto_response(resp, cluster_id, member_id) } crate::proto::watch_request::RequestUnion::CancelRequest(cancel) => { let resp = stream.cancel_watch(cancel.watch_id); internal_to_proto_response(resp, cluster_id, member_id) } crate::proto::watch_request::RequestUnion::ProgressRequest(_) => { // Send progress notification WatchResponse { header: Some(make_header( cluster_id, member_id, registry.current_revision(), 0, )), watch_id: 0, created: false, canceled: false, compact_revision: 0, cancel_reason: String::new(), events: vec![], } } }; if tx.send(Ok(response)).await.is_err() { break; } } } Err(e) => { warn!(error = %e, "Watch stream error"); break; } } } debug!(watches = stream.watch_count(), "Watch stream closed"); // Stream cleanup happens in WatchStream::drop }); // Spawn task to forward watch events tokio::spawn(async move { while let Some(response) = event_rx.recv().await { if tx_for_events.send(Ok(response)).await.is_err() { break; } } }); let output_stream = ReceiverStream::new(rx); Ok(Response::new(Box::pin(output_stream))) } } fn internal_to_proto_response( resp: chainfire_types::watch::WatchResponse, cluster_id: u64, member_id: u64, ) -> WatchResponse { WatchResponse { header: Some(make_header(cluster_id, member_id, resp.compact_revision, 0)), watch_id: resp.watch_id, created: resp.created, canceled: resp.canceled, compact_revision: resp.compact_revision as i64, cancel_reason: String::new(), events: resp.events.into_iter().map(Into::into).collect(), } }