feat: add monitor for debugging

Author: SupernaviX

modules/clock/src/clock.rs

@@ -2,7 +2,7 @@
 //! Generates regular clock.tick events
 
 use anyhow::Result;
-use caryatid_sdk::{module, Context, MessageBounds, Module};
+use caryatid_sdk::{module, Context, MessageBounds};
 use chrono::{DateTime, Utc};
 use config::Config;
 use std::sync::Arc;
@@ -70,6 +70,7 @@ impl<M: From<ClockTickMessage> + MessageBounds> Clock<M> {
 mod tests {
     use super::*;
     use caryatid_sdk::mock_bus::MockBus;
+    use caryatid_sdk::Module;
     use config::{Config, FileFormat};
     use tokio::sync::{mpsc, watch::Sender, Notify};
     use tokio::time::{timeout, Duration};
@@ -99,6 +100,7 @@ mod tests {
     struct TestSetup {
         module: Arc<dyn Module<Message>>,
         context: Arc<Context<Message>>,
+        startup_watch: Sender<bool>,
     }
 
     impl TestSetup {
@@ -120,12 +122,9 @@ mod tests {
             // Create mock bus
             let bus = Arc::new(MockBus::<Message>::new(&config));
 
+            let startup_watch = Sender::new(false);
             // Create a context
-            let context = Arc::new(Context::new(
-                config.clone(),
-                bus.clone(),
-                Sender::<bool>::new(false),
-            ));
+            let context = Arc::new(Context::new(config.clone(), bus, startup_watch.subscribe()));
 
             // Create the clock
             let clock = Clock::<Message> {
@@ -136,11 +135,12 @@ mod tests {
             Self {
                 module: Arc::new(clock),
                 context,
+                startup_watch,
             }
         }
 
         fn start(&self) {
-            let _ = self.context.startup_watch.send(true);
+            let _ = self.startup_watch.send(true);
         }
     }
 

modules/rest_server/src/rest_server.rs

@@ -2,7 +2,7 @@
 //! Provides a REST endpoint which integrates with the message bus
 
 use anyhow::Result;
-use caryatid_sdk::{module, Context, MessageBounds, Module};
+use caryatid_sdk::{module, Context, MessageBounds};
 use config::Config;
 use std::{collections::HashMap, sync::Arc};
 use tracing::{error, info};
@@ -176,6 +176,7 @@ impl<M: From<RESTRequest> + GetRESTResponse + MessageBounds> RESTServer<M> {
 mod tests {
     use super::*;
     use caryatid_sdk::mock_bus::MockBus;
+    use caryatid_sdk::Module;
     use config::{Config, FileFormat};
     use futures::future;
     use hyper::Client;
@@ -227,6 +228,7 @@ mod tests {
     struct TestSetup {
         module: Arc<dyn Module<Message>>,
         context: Arc<Context<Message>>,
+        startup_watch: Sender<bool>,
     }
 
     impl TestSetup {
@@ -249,10 +251,11 @@ mod tests {
             let mock_bus = Arc::new(MockBus::<Message>::new(&config));
 
             // Create a context
+            let startup_watch = Sender::new(false);
             let context = Arc::new(Context::new(
                 config.clone(),
-                mock_bus.clone(),
-                Sender::<bool>::new(false),
+                mock_bus,
+                startup_watch.subscribe(),
             ));
 
             // Create the server
@@ -264,11 +267,12 @@ mod tests {
             Self {
                 module: Arc::new(rest_server),
                 context,
+                startup_watch,
             }
         }
 
         fn start(&self) {
-            let _ = self.context.startup_watch.send(true);
+            let _ = self.startup_watch.send(true);
         }
     }
 

process/Cargo.toml

@@ -14,6 +14,7 @@ anyhow = "1.0"
 tokio = { version = "1", features = ["full"] }
 config = "0.15.11"
 minicbor-serde = { version = "0.6", features = ["alloc"] }
+dashmap = "6"
 tracing = "0.1.40"
 serde = "1.0.210"
 serde_json = "1.0"

process/src/monitor.rs

@@ -0,0 +1,254 @@
+use std::{
+    collections::BTreeMap,
+    future::Future,
+    path::PathBuf,
+    sync::Arc,
+    task::Poll,
+    time::{Duration, Instant},
+};
+
+use anyhow::Result;
+use async_trait::async_trait;
+use caryatid_sdk::{MessageBounds, MessageBus, Subscription, SubscriptionBounds};
+use dashmap::DashMap;
+use futures::{future::BoxFuture, FutureExt};
+use serde::{Deserialize, Serialize};
+use tokio::{fs, time};
+
+#[derive(Default, Clone)]
+struct ReadStreamState {
+    read: u64,
+    pending_since: Option<Instant>,
+}
+
+#[derive(Default, Clone)]
+struct WriteStreamState {
+    written: u64,
+    pending_since: Option<Instant>,
+}
+
+#[derive(Default, Clone)]
+struct ModuleState {
+    reads: DashMap<String, ReadStreamState>,
+    writes: DashMap<String, WriteStreamState>,
+}
+
+#[derive(Serialize)]
+struct SerializedReadStreamState {
+    read: u64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    unread: Option<u64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pending_for: Option<String>,
+}
+
+#[derive(Serialize)]
+struct SerializedWriteStreamState {
+    written: u64,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pending_for: Option<String>,
+}
+
+#[derive(Serialize)]
+struct SerializedModuleState {
+    reads: BTreeMap<String, SerializedReadStreamState>,
+    writes: BTreeMap<String, SerializedWriteStreamState>,
+}
+
+const fn default_frequency() -> Duration {
+    Duration::from_secs(5)
+}
+
+#[derive(Deserialize)]
+pub struct MonitorConfig {
+    output: PathBuf,
+    #[serde(default = "default_frequency")]
+    frequency: Duration,
+}
+
+pub struct Monitor {
+    modules: BTreeMap<String, Arc<ModuleState>>,
+    stream_writes: Arc<DashMap<String, u64>>,
+    output_path: PathBuf,
+    write_frequency: Duration,
+}
+impl Monitor {
+    pub fn new(config: MonitorConfig) -> Self {
+        Self {
+            modules: BTreeMap::new(),
+            stream_writes: Arc::new(DashMap::new()),
+            output_path: config.output,
+            write_frequency: config.frequency,
+        }
+    }
+
+    pub fn spy_on_bus<M: MessageBounds>(
+        &mut self,
+        module_name: &str,
+        message_bus: Arc<dyn MessageBus<M>>,
+    ) -> Arc<dyn MessageBus<M>> {
+        let state = Arc::new(ModuleState::default());
+        self.modules.insert(module_name.to_string(), state.clone());
+
+        Arc::new(MonitorBus {
+            inner: message_bus,
+            stream_writes: self.stream_writes.clone(),
+            state,
+        })
+    }
+
+    pub async fn monitor(self) {
+        loop {
+            time::sleep(self.write_frequency).await;
+            let now = Instant::now();
+            let state = self
+                .modules
+                .iter()
+                .map(|(name, state)| {
+                    let reads = state
+                        .reads
+                        .iter()
+                        .map(|kvp| {
+                            let (topic, data) = kvp.pair();
+                            let read = data.read;
+                            let unread = self
+                                .stream_writes
+                                .get(topic)
+                                .and_then(|w| w.checked_sub(read))
+                                .filter(|u| *u > 0);
+                            let pending_for = data
+                                .pending_since
+                                .map(|d| format!("{:?}", now.duration_since(d)));
+                            let state = SerializedReadStreamState {
+                                read,
+                                unread,
+                                pending_for,
+                            };
+                            (topic.clone(), state)
+                        })
+                        .collect();
+
+                    let writes = state
+                        .writes
+                        .iter()
+                        .map(|kvp| {
+                            let (topic, data) = kvp.pair();
+                            let written = data.written;
+                            let pending_for = data
+                                .pending_since
+                                .map(|d| format!("{:?}", now.duration_since(d)));
+                            let state = SerializedWriteStreamState {
+                                written,
+                                pending_for,
+                            };
+                            (topic.clone(), state)
+                        })
+                        .collect();
+
+                    (name.clone(), SerializedModuleState { reads, writes })
+                })
+                .collect::<BTreeMap<_, _>>();
+            let serialized = serde_json::to_vec_pretty(&state).expect("could not serialize state");
+            fs::write(&self.output_path, serialized)
+                .await
+                .expect("could not write file");
+        }
+    }
+}
+
+pub struct MonitorBus<M: MessageBounds> {
+    inner: Arc<dyn MessageBus<M>>,
+    stream_writes: Arc<DashMap<String, u64>>,
+    state: Arc<ModuleState>,
+}
+
+#[async_trait]
+impl<M: MessageBounds> MessageBus<M> for MonitorBus<M> {
+    async fn publish(&self, topic: &str, message: Arc<M>) -> Result<()> {
+        self.state
+            .writes
+            .entry(topic.to_string())
+            .or_default()
+            .pending_since = Some(Instant::now());
+        let res = self.inner.publish(topic, message).await;
+        let mut writes = self.state.writes.entry(topic.to_string()).or_default();
+        writes.written += 1;
+        writes.pending_since = None;
+        *self.stream_writes.entry(topic.to_string()).or_default() += 1;
+        res
+    }
+
+    fn request_timeout(&self) -> std::time::Duration {
+        self.inner.request_timeout()
+    }
+
+    async fn subscribe(&self, topic: &str) -> Result<Box<dyn Subscription<M>>> {
+        self.state.reads.entry(topic.to_string()).or_default();
+        Ok(Box::new(MonitorSubscription {
+            inner: self.inner.subscribe(topic).await?,
+            state: self.state.clone(),
+            topic: topic.to_string(),
+        }))
+    }
+
+    async fn shutdown(&self) -> Result<()> {
+        self.inner.shutdown().await
+    }
+}
+
+struct MonitorSubscription<M: MessageBounds> {
+    inner: Box<dyn Subscription<M>>,
+    state: Arc<ModuleState>,
+    topic: String,
+}
+impl<M: MessageBounds> SubscriptionBounds for MonitorSubscription<M> {}
+
+impl<M: MessageBounds> Subscription<M> for MonitorSubscription<M> {
+    fn read(&mut self) -> BoxFuture<'_, Result<(String, Arc<M>)>> {
+        Box::pin(
+            MonitorReadFuture {
+                inner: self.inner.read(),
+                state: &self.state,
+                topic: &self.topic,
+            }
+            .fuse(),
+        )
+    }
+}
+
+struct MonitorReadFuture<'a, M: MessageBounds> {
+    inner: BoxFuture<'a, Result<(String, Arc<M>)>>,
+    state: &'a ModuleState,
+    topic: &'a str,
+}
+impl<'a, M: MessageBounds> Future for MonitorReadFuture<'a, M> {
+    type Output = Result<(String, Arc<M>)>;
+
+    fn poll(
+        mut self: std::pin::Pin<&mut Self>,
+        cx: &mut std::task::Context<'_>,
+    ) -> std::task::Poll<Self::Output> {
+        let res = self.inner.poll_unpin(cx);
+        let mut entry = self.state.reads.entry(self.topic.to_string()).or_default();
+        match &res {
+            Poll::Pending => {
+                if entry.pending_since.is_none() {
+                    entry.pending_since = Some(Instant::now());
+                }
+            }
+            Poll::Ready(r) => {
+                entry.pending_since = None;
+                if r.is_ok() {
+                    entry.read += 1;
+                }
+            }
+        }
+        res
+    }
+}
+impl<'a, M: MessageBounds> Drop for MonitorReadFuture<'a, M> {
+    fn drop(&mut self) {
+        let mut entry = self.state.reads.entry(self.topic.to_string()).or_default();
+        entry.pending_since = None;
+    }
+}