zelos_trace_grpc/subscribe/
client.rs

1use anyhow::Result;
2use tokio::sync::mpsc::Sender;
3use tokio_stream::wrappers::ReceiverStream;
4use tokio_util::sync::CancellationToken;
5use tonic::Streaming;
6use zelos_proto::trace::{
7    subscribe_request, trace_subscribe_client, SubscribeCommand, SubscribeRequest,
8    SubscribeResponse, UnsubscribeCommand,
9};
10
11pub struct TraceSubscribeClient {
12    /// The sender for subscribe requests.
13    req_sender: Sender<SubscribeRequest>,
14}
15
16impl TraceSubscribeClient {
17    /// Create a new TraceSubscribeClient and connect to the given address.
18    pub async fn new(
19        sender: zelos_trace_types::ipc::Sender,
20        cancellation_token: CancellationToken,
21        address: String,
22    ) -> Result<(Self, impl Future<Output = Result<()>>)> {
23        // Connect to the gRPC server
24        let channel = zelos_proto::channel::create_channel(address)?;
25        let mut client = trace_subscribe_client::TraceSubscribeClient::new(channel)
26            .max_decoding_message_size(zelos_proto::MAX_GRPC_MESSAGE_SIZE)
27            .max_encoding_message_size(zelos_proto::MAX_GRPC_MESSAGE_SIZE);
28
29        // Initialize a channel for sending subscribe requests
30        let (req_sender, req_receiver) = tokio::sync::mpsc::channel(1);
31
32        // Attempt to call the subscribe streaming method, exiting early if we fail
33        let request_stream = tonic::Request::new(ReceiverStream::new(req_receiver));
34        let resp = client.subscribe(request_stream).await?;
35
36        // Run our task to forward from the response stream to the sender
37        let future = Self::run(resp.into_inner(), sender.clone(), cancellation_token);
38
39        Ok((Self { req_sender }, future))
40    }
41
42    /// Run a task to forward from the response stream to the sender
43    async fn run(
44        mut stream: Streaming<SubscribeResponse>,
45        sender: zelos_trace_types::ipc::Sender,
46        cancellation_token: CancellationToken,
47    ) -> Result<()> {
48        loop {
49            tokio::select! {
50                msg = stream.message() => {
51                    match msg {
52                        Ok(Some(response)) => {
53                            // Forward the message to the router
54                            let ipc = response.as_ipc()?;
55                            for m in ipc {
56                                sender.send_async(m).await?;
57                            }
58                        }
59                        Ok(None) => {
60                            // Stream ended
61                            return Ok(());
62                        }
63                        Err(e) => {
64                            // Error from the client
65                            return Err(e.into());
66                        }
67                    }
68                }
69                _ = cancellation_token.cancelled() => return Ok(())
70            }
71        }
72    }
73
74    /// Send a subscribe command with the given filter and start time
75    pub async fn subscribe(&self, filter: Option<String>, start_time: Option<i64>) -> Result<()> {
76        self.req_sender
77            .send(SubscribeRequest {
78                cmd: Some(subscribe_request::Cmd::Subscribe(SubscribeCommand {
79                    filter,
80                    start_time,
81                })),
82            })
83            .await?;
84
85        Ok(())
86    }
87
88    /// Send an unsubscribe command with the given filter
89    pub async fn unsubscribe(&self, filter: Option<String>) -> Result<()> {
90        self.req_sender
91            .send(SubscribeRequest {
92                cmd: Some(subscribe_request::Cmd::Unsubscribe(UnsubscribeCommand {
93                    filter,
94                })),
95            })
96            .await?;
97
98        Ok(())
99    }
100
101    /// Send a subscribe command with a blank filter
102    pub async fn subscribe_all(&self) -> Result<()> {
103        self.subscribe(None, None).await
104    }
105
106    /// Send an unsubscribe command with a blank filter
107    pub async fn unsubscribe_all(&self) -> Result<()> {
108        self.unsubscribe(None).await
109    }
110}
zelos_trace_grpc/subscribe/client.rs

zelos_trace_grpc/subscribe/
client.rs
