From f892a50f2c5299a85f990d83dae1a9e707f1a5bf Mon Sep 17 00:00:00 2001
From: asonix <asonix@asonix.dog>
Date: Sat, 19 Nov 2022 17:45:01 -0600
Subject: [PATCH] Add metrics printer

---
 Cargo.lock          | 8 ++++----
 src/admin/client.rs | 5 +++++
 src/args.rs         | 9 ++++++++-
 src/collector.rs    | 8 ++++----
 src/config.rs       | 3 +++
 src/main.rs         | 5 +++++
 6 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 694caf4..e05d130 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -499,9 +499,9 @@ dependencies = [
 
 [[package]]
 name = "background-jobs-actix"
-version = "0.14.0"
+version = "0.14.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "ea433afb3fbbb6dc2614bad9f6671517f6cffcd523981e568cdb595dc7fa6399"
+checksum = "3aba1bc1cdff87161a6e2e00bf9dc1712412ee926a065d96cc0a03dc851b5fd3"
 dependencies = [
  "actix-rt",
  "anyhow",
@@ -520,9 +520,9 @@ dependencies = [
 
 [[package]]
 name = "background-jobs-core"
-version = "0.14.0"
+version = "0.14.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "b3071bf7a5e46638085076957dc189b2b0b147cd279eb09510b7af54e95085ef"
+checksum = "1274e49ae8eff1fc6b4943660e59ce2f2e13e65a23a707924a50a40c7b94fc4d"
 dependencies = [
  "actix-rt",
  "anyhow",
diff --git a/src/admin/client.rs b/src/admin/client.rs
index f63151f..3602487 100644
--- a/src/admin/client.rs
+++ b/src/admin/client.rs
@@ -1,5 +1,6 @@
 use crate::{
     admin::{AllowedDomains, BlockedDomains, ConnectedActors, Domains},
+    collector::Snapshot,
     config::{AdminUrlKind, Config},
     error::{Error, ErrorKind},
 };
@@ -50,6 +51,10 @@ pub(crate) async fn connected(client: &Client, config: &Config) -> Result<Connec
     get_results(client, config, AdminUrlKind::Connected).await
 }
 
+pub(crate) async fn stats(client: &Client, config: &Config) -> Result<Snapshot, Error> {
+    get_results(client, config, AdminUrlKind::Stats).await
+}
+
 async fn get_results<T: DeserializeOwned>(
     client: &Client,
     config: &Config,
diff --git a/src/args.rs b/src/args.rs
index 6b6c054..18a4059 100644
--- a/src/args.rs
+++ b/src/args.rs
@@ -14,11 +14,14 @@ pub(crate) struct Args {
 
     #[arg(short, long, help = "List allowed and blocked domains")]
     list: bool,
+
+    #[arg(short, long, help = "Get statistics from the server")]
+    stats: bool,
 }
 
 impl Args {
     pub(crate) fn any(&self) -> bool {
-        !self.blocks.is_empty() || !self.allowed.is_empty() || self.list
+        !self.blocks.is_empty() || !self.allowed.is_empty() || self.list || self.stats
     }
 
     pub(crate) fn new() -> Self {
@@ -40,4 +43,8 @@ impl Args {
     pub(crate) fn list(&self) -> bool {
         self.list
     }
+
+    pub(crate) fn stats(&self) -> bool {
+        self.stats
+    }
 }
diff --git a/src/collector.rs b/src/collector.rs
index 401307b..dc5d179 100644
--- a/src/collector.rs
+++ b/src/collector.rs
@@ -29,25 +29,25 @@ struct Inner {
     registry: Registry<Key, GenerationalStorage<AtomicStorage>>,
 }
 
-#[derive(serde::Serialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
 struct Counter {
     labels: Vec<(String, String)>,
     value: u64,
 }
 
-#[derive(serde::Serialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
 struct Gauge {
     labels: Vec<(String, String)>,
     value: f64,
 }
 
-#[derive(serde::Serialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
 struct Histogram {
     labels: Vec<(String, String)>,
     value: Vec<(f64, Option<f64>)>,
 }
 
-#[derive(serde::Serialize)]
+#[derive(Debug, serde::Deserialize, serde::Serialize)]
 pub(crate) struct Snapshot {
     counters: HashMap<String, Vec<Counter>>,
     gauges: HashMap<String, Vec<Gauge>>,
diff --git a/src/config.rs b/src/config.rs
index 284f807..1bb4f19 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -77,6 +77,7 @@ pub enum AdminUrlKind {
     Allowed,
     Blocked,
     Connected,
+    Stats,
 }
 
 impl std::fmt::Debug for Config {
@@ -338,6 +339,8 @@ impl Config {
                 .try_resolve(IriRelativeStr::new("api/v1/admin/blocked")?.as_ref())?,
             AdminUrlKind::Connected => FixedBaseResolver::new(self.base_uri.as_ref())
                 .try_resolve(IriRelativeStr::new("api/v1/admin/connected")?.as_ref())?,
+            AdminUrlKind::Stats => FixedBaseResolver::new(self.base_uri.as_ref())
+                .try_resolve(IriRelativeStr::new("api/v1/admin/stats")?.as_ref())?,
         };
 
         Ok(iri)
diff --git a/src/main.rs b/src/main.rs
index 9aebdb9..884e572 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -142,6 +142,11 @@ async fn main() -> Result<(), anyhow::Error> {
             println!("{report}");
         }
 
+        if args.stats() {
+            let stats = admin::client::stats(&client, &config).await?;
+            println!("{:#?}", stats);
+        }
+
         return Ok(());
     }