src: hub: Add driver stats to the hub

joaoantoniocardoso · Sep 8, 2024 · 67a2027 · 67a2027
1 parent c63da3a
commit 67a2027
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 2 deletions.
diff --git a/src/hub/actor.rs b/src/hub/actor.rs
@@ -8,6 +8,7 @@ use tracing::*;
 use crate::{
     drivers::{Driver, DriverInfo},
     protocol::Protocol,
+    stats::driver::DriverStatsInfo,
 };
 
 use super::protocol::HubCommand;
@@ -40,6 +41,13 @@ impl HubActor {
                 HubCommand::GetSender { response } => {
                     let _ = response.send(self.bcst_sender.clone());
                 }
+                HubCommand::GetStats { response } => {
+                    let stats = self.get_stats().await;
+                    let _ = response.send(stats);
+                }
+                HubCommand::ResetAllStats { response } => {
+                    let _ = response.send(self.reset_all_stats().await);
+                }
             }
         }
     }
@@ -159,4 +167,31 @@ impl HubActor {
     pub fn get_sender(&self) -> broadcast::Sender<Arc<Protocol>> {
         self.bcst_sender.clone()
     }
+
+    #[instrument(level = "debug", skip(self))]
+    pub async fn get_stats(&self) -> Vec<(String, DriverStatsInfo)> {
+        let drivers = self.drivers.read().await;
+
+        let mut drivers_stats = Vec::with_capacity(drivers.len());
+        for (_id, driver) in drivers.iter() {
+            let stats = driver.stats().await;
+            let info = driver.info();
+            let name = info.name().to_owned();
+
+            drivers_stats.push((name, stats));
+        }
+
+        drivers_stats
+    }
+
+    #[instrument(level = "debug", skip(self))]
+    pub async fn reset_all_stats(&self) -> Result<()> {
+        let drivers = self.drivers.write().await;
+
+        for (_id, driver) in drivers.iter() {
+            driver.reset_stats().await;
+        }
+
+        Ok(())
+    }
 }
diff --git a/src/hub/mod.rs b/src/hub/mod.rs
@@ -9,8 +9,11 @@ use std::{
 use anyhow::Result;
 use tokio::sync::{broadcast, mpsc, oneshot, RwLock};
 
-use crate::drivers::{Driver, DriverInfo};
-use crate::protocol::Protocol;
+use crate::{
+    drivers::{Driver, DriverInfo},
+    protocol::Protocol,
+    stats::driver::DriverStatsInfo,
+};
 
 use actor::HubActor;
 use protocol::HubCommand;
@@ -77,4 +80,25 @@ impl Hub {
         let res = response_rx.await?;
         Ok(res)
     }
+
+    pub async fn stats(&self) -> Result<Vec<(String, DriverStatsInfo)>> {
+        let (response_tx, response_rx) = oneshot::channel();
+        self.sender
+            .send(HubCommand::GetStats {
+                response: response_tx,
+            })
+            .await?;
+        let res = response_rx.await?;
+        Ok(res)
+    }
+
+    pub async fn reset_all_stats(&self) -> Result<()> {
+        let (response_tx, response_rx) = oneshot::channel();
+        self.sender
+            .send(HubCommand::ResetAllStats {
+                response: response_tx,
+            })
+            .await?;
+        response_rx.await?
+    }
 }
diff --git a/src/hub/protocol.rs b/src/hub/protocol.rs
@@ -6,6 +6,7 @@ use tokio::sync::{broadcast, oneshot};
 use crate::{
     drivers::{Driver, DriverInfo},
     protocol::Protocol,
+    stats::driver::DriverStatsInfo,
 };
 
 pub enum HubCommand {
@@ -23,4 +24,10 @@ pub enum HubCommand {
     GetSender {
         response: oneshot::Sender<broadcast::Sender<Arc<Protocol>>>,
     },
+    GetStats {
+        response: oneshot::Sender<Vec<(String, DriverStatsInfo)>>,
+    },
+    ResetAllStats {
+        response: oneshot::Sender<Result<()>>,
+    },
 }