photoncloud-monorepo/chainfire/crates/chainfire-api/src/watch_service.rs

//! 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<WatchRegistry>,
    /// Cluster ID
    cluster_id: u64,
    /// Member ID
    member_id: u64,
}

impl WatchServiceImpl {
    /// Create a new watch service
    pub fn new(registry: Arc<WatchRegistry>, 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<Box<dyn tokio_stream::Stream<Item = Result<WatchResponse, Status>> + Send>>;

    async fn watch(
        &self,
        request: Request<Streaming<WatchRequest>>,
    ) -> Result<Response<Self::WatchStream>, 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::<crate::proto::WatchResponse>(128);

        // Spawn task to handle the bidirectional stream
        tokio::spawn(async move {
            let mut stream = WatchStream::new(Arc::clone(&registry), {
                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(),
    }
}