feat: implement watch_state

Signed-off-by: Rustin170506 <29879298+Rustin170506@users.noreply.github.com>
tokio-rs · Sep 5, 2024 · d0810c6 · d0810c6
1 parent 1b485cc
commit d0810c6
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 11 deletions.
diff --git a/console-subscriber/src/aggregator/mod.rs b/console-subscriber/src/aggregator/mod.rs
@@ -89,7 +89,7 @@ pub struct Aggregator {
     poll_ops: Vec<proto::resources::PollOp>,
 
     /// The time "state" of the aggregator, such as paused or live.
-    temporality: Temporality,
+    temporality:  proto::instrument::Temporality,
 
     /// Used to anchor monotonic timestamps to a base `SystemTime`, to produce a
     /// timestamp that can be sent over the wire.
@@ -102,11 +102,6 @@ pub(crate) struct Flush {
     triggered: AtomicBool,
 }
 
-#[derive(Debug)]
-enum Temporality {
-    Live,
-    Paused,
-}
 // Represent static data for resources
 struct Resource {
     id: Id,
@@ -162,7 +157,7 @@ impl Aggregator {
             async_ops: IdData::default(),
             async_op_stats: IdData::default(),
             poll_ops: Default::default(),
-            temporality: Temporality::Live,
+            temporality: proto::instrument::Temporality::Live,
             base_time,
         }
     }
@@ -179,8 +174,8 @@ impl Aggregator {
                 // if the flush interval elapses, flush data to the client
                 _ = publish.tick() => {
                     match self.temporality {
-                        Temporality::Live => true,
-                        Temporality::Paused => false,
+                        proto::instrument::Temporality::Live => true,
+                        proto::instrument::Temporality::Paused => false,
                     }
                 }
 
@@ -199,11 +194,17 @@ impl Aggregator {
                         Some(Command::WatchTaskDetail(watch_request)) => {
                             self.add_task_detail_subscription(watch_request);
                         },
+                        Some(Command::WatchState(subscription)) => {
+                            let state = proto::instrument::State {
+                                temporality: self.temporality.into(),
+                            };
+                            subscription.update(&state);
+                        },
                         Some(Command::Pause) => {
-                            self.temporality = Temporality::Paused;
+                            self.temporality = proto::instrument::Temporality::Paused;
                         }
                         Some(Command::Resume) => {
-                            self.temporality = Temporality::Live;
+                            self.temporality = proto::instrument::Temporality::Live;
                         }
                         None => {
                             tracing::debug!("rpc channel closed, terminating");

diff --git a/console-subscriber/src/lib.rs b/console-subscriber/src/lib.rs
@@ -178,6 +178,7 @@ struct Watch<T>(mpsc::Sender<Result<T, tonic::Status>>);
 enum Command {
     Instrument(Watch<proto::instrument::Update>),
     WatchTaskDetail(WatchRequest<proto::tasks::TaskDetails>),
+    WatchState(Watch<proto::instrument::State>),
     Pause,
     Resume,
 }
@@ -1190,6 +1191,8 @@ impl proto::instrument::instrument_server::Instrument for Server {
         tokio_stream::wrappers::ReceiverStream<Result<proto::instrument::Update, tonic::Status>>;
     type WatchTaskDetailsStream =
         tokio_stream::wrappers::ReceiverStream<Result<proto::tasks::TaskDetails, tonic::Status>>;
+    type WatchStateStream =
+        tokio_stream::wrappers::ReceiverStream<Result<proto::instrument::State, tonic::Status>>;
     async fn watch_updates(
         &self,
         req: tonic::Request<proto::instrument::InstrumentRequest>,
@@ -1245,6 +1248,18 @@ impl proto::instrument::instrument_server::Instrument for Server {
         Ok(tonic::Response::new(stream))
     }
 
+    async fn watch_state(
+        &self,
+        _req: tonic::Request<proto::instrument::StateRequest>,
+    ) -> Result<tonic::Response<Self::WatchStateStream>, tonic::Status> {
+       let (stream_sender, stream_recv) = mpsc::channel(self.client_buffer);
+       self.subscribe.send(Command::WatchState(Watch(stream_sender))).await.map_err(|_| {
+            tonic::Status::internal("cannot get state, aggregation task is not running")
+        })?;
+        let stream = tokio_stream::wrappers::ReceiverStream::new(stream_recv);
+        Ok(tonic::Response::new(stream))
+    }
+
     async fn pause(
         &self,
         _req: tonic::Request<proto::instrument::PauseRequest>,