zelos_trace_grpc/subscribe/
service.rs

1use std::{pin::Pin, sync::Arc, time::Duration};
2
3use tokio_stream::{Stream, StreamExt};
4use tonic::{Request, Response, Status, Streaming};
5use zelos_proto::trace::{
6    subscribe_request::Cmd,
7    trace_subscribe_server::{TraceSubscribe, TraceSubscribeServer},
8    SubscribeRequest, SubscribeResponse,
9};
10use zelos_trace::{filter::Filter, TraceRouter};
11
12const CHUNK_SIZE: usize = 1024;
13const CHUNK_TIMEOUT: Duration = Duration::from_millis(10);
14
15pub struct TraceSubscribeService {
16    router: Arc<TraceRouter>,
17}
18
19impl TraceSubscribeService {
20    pub fn new(router: Arc<TraceRouter>) -> Self {
21        Self { router }
22    }
23
24    pub fn server(self) -> TraceSubscribeServer<Self> {
25        TraceSubscribeServer::new(self)
26    }
27}
28
29#[tonic::async_trait]
30impl TraceSubscribe for TraceSubscribeService {
31    type SubscribeStream =
32        Pin<Box<dyn Stream<Item = Result<SubscribeResponse, Status>> + Send + 'static>>;
33
34    // TODO(jbott): handle disconnects
35    async fn subscribe(
36        &self,
37        request: Request<Streaming<SubscribeRequest>>,
38    ) -> Result<Response<Self::SubscribeStream>, Status> {
39        // Attach to our router, forwarding trace messages to the client using a stream
40        let (sink, stream) = self
41            .router
42            .subscribe_stream()
43            .await
44            .map_err(|e| Status::internal(format!("Failed to subscribe: {}", e)))?;
45
46        let stream = stream.chunks_timeout(CHUNK_SIZE, CHUNK_TIMEOUT).map(|m| {
47            Ok(SubscribeResponse::from_ipc(
48                m.into_iter().map(|msg| msg.into()).collect(),
49            ))
50        });
51
52        // Handle messages from the client
53        let mut req_stream = request.into_inner();
54        tokio::task::spawn(async move {
55            while let Some(req) = req_stream.message().await? {
56                if let Some(cmd) = req.cmd {
57                    match cmd {
58                        Cmd::Subscribe(subscribe) => {
59                            let filter = match &subscribe.filter {
60                                Some(f) => Filter::parse(f),
61                                None => Ok(Filter::any()),
62                            };
63
64                            match filter {
65                                Ok(f) => sink.subscribe(f).await,
66                                Err(e) => tracing::error!("Failed to parse filter: {}", e),
67                            }
68                        }
69                        Cmd::Unsubscribe(unsubscribe) => {
70                            let filter = match &unsubscribe.filter {
71                                Some(f) => Filter::parse(f),
72                                None => Ok(Filter::any()),
73                            };
74
75                            match filter {
76                                Ok(f) => sink.unsubscribe(f).await,
77                                Err(e) => tracing::error!("Failed to parse filter: {}", e),
78                            }
79                        }
80                    }
81                }
82            }
83
84            Ok::<_, anyhow::Error>(())
85        });
86
87        Ok(Response::new(Box::pin(stream)))
88    }
89}
zelos_trace_grpc/subscribe/service.rs

zelos_trace_grpc/subscribe/
service.rs
