From 612c11d3c67137acd7d3ef8401deb04e9da4e6f1 Mon Sep 17 00:00:00 2001
From: yellowhatpro <yellowhatpro3119@gmail.com>
Date: Wed, 14 Aug 2024 12:15:36 +0530
Subject: [PATCH] feat: added archival status metrics

---
 grafana/dashboards/metrics-dashboard.json | 22 ++++++++++++++++----
 src/archival/listener.rs                  |  3 +++
 src/archival/utils.rs                     |  2 ++
 src/metrics/mod.rs                        | 25 ++++++++++++++++++++++-
 4 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/grafana/dashboards/metrics-dashboard.json b/grafana/dashboards/metrics-dashboard.json
index 94b2c4d..a3f2978 100644
--- a/grafana/dashboards/metrics-dashboard.json
+++ b/grafana/dashboards/metrics-dashboard.json
@@ -133,21 +133,35 @@
           "range": true,
           "refId": "B",
           "useBackend": false
+        },
+        {
+          "datasource": {
+            "type": "prometheus",
+            "uid": "PBFA97CFB590B2093"
+          },
+          "disableTextWrap": false,
+          "editorMode": "builder",
+          "expr": "archival_requests_total",
+          "fullMetaSearch": false,
+          "hide": false,
+          "includeNullMetadata": true,
+          "instant": false,
+          "legendFormat": "__auto",
+          "range": true,
+          "refId": "C",
+          "useBackend": false
         }
       ],
       "title": "Rust app metrics panel",
       "type": "timeseries"
     }
   ],
+  "refresh": "",
   "schemaVersion": 39,
   "tags": [],
   "templating": {
     "list": []
   },
-  "time": {
-    "from": "now-6h",
-    "to": "now"
-  },
   "timepicker": {},
   "timezone": "",
   "title": "mb-ia-dashboard",
diff --git a/src/archival/listener.rs b/src/archival/listener.rs
index 27c483b..656639b 100644
--- a/src/archival/listener.rs
+++ b/src/archival/listener.rs
@@ -5,6 +5,7 @@ use crate::archival::utils::{
 
 use crate::archival::error::ArchivalError;
 use crate::configuration::Settings;
+use crate::metrics::Metrics;
 use crate::structs::internet_archive_urls::{ArchivalStatus, InternetArchiveUrls};
 use sentry::Level::Error;
 use sqlx::postgres::PgListener;
@@ -33,6 +34,7 @@ pub async fn handle_payload(
     url_row: InternetArchiveUrls,
     pool: &PgPool,
 ) -> Result<(), ArchivalError> {
+    let metrics = Metrics::new().await;
     if let Some(url) = url_row.url {
         let id = url_row.id;
         if url_row.retry_count >= Some(3) {
@@ -46,6 +48,7 @@ pub async fn handle_payload(
             set_status_with_message(pool, id, ArchivalStatus::Failed as i32, status_ext.as_str())
                 .await?;
             sentry::capture_message(status_ext.as_str(), Error);
+            metrics.record_archival_status("failed").await;
         } else {
             let archival_result = archive(url, url_row.id, pool).await;
             if let Err(e) = archival_result {
diff --git a/src/archival/utils.rs b/src/archival/utils.rs
index 76b4070..fa9740c 100644
--- a/src/archival/utils.rs
+++ b/src/archival/utils.rs
@@ -155,6 +155,7 @@ pub async fn schedule_status_check(
                 archival_status_response.status.as_str(),
             )
             .await?;
+            metrics.record_archival_status("success").await;
             return Ok(());
         } else {
             if attempt == 3 {
@@ -169,6 +170,7 @@ pub async fn schedule_status_check(
                 )
                 .await?;
             }
+            metrics.record_archival_status("error").await;
             eprintln!("Could not archive: {} attempt", attempt)
         }
     }
diff --git a/src/metrics/mod.rs b/src/metrics/mod.rs
index bb3c731..9dd608f 100644
--- a/src/metrics/mod.rs
+++ b/src/metrics/mod.rs
@@ -1,4 +1,4 @@
-use prometheus::{labels, push_metrics, Counter, Opts, Registry};
+use prometheus::{labels, push_metrics, Counter, CounterVec, Opts, Registry};
 use std::sync::Arc;
 use tokio::sync::Mutex;
 use tokio::task::spawn_blocking;
@@ -6,12 +6,14 @@ use tokio::task::spawn_blocking;
 pub struct Metrics {
     pub db_poll_counter: Counter,
     pub network_request_counter: Counter,
+    pub archival_status_counter: CounterVec, // New counter for archival statuses
     pub registry: Arc<Mutex<Registry>>,
 }
 
 impl Metrics {
     pub async fn new() -> Self {
         let registry = Arc::new(Mutex::new(Registry::new()));
+
         let poll_opts = Opts::new(
             "db_polls_total",
             "Total number of edit data and edit notes polls",
@@ -33,12 +35,26 @@ impl Metrics {
             .await
             .register(Box::new(network_request_counter.clone()))
             .unwrap();
+
+        let status_opts = Opts::new(
+            "archival_requests_total",
+            "Total number of archival requests by status",
+        );
+        let archival_status_counter = CounterVec::new(status_opts, &["status"]).unwrap();
+        registry
+            .lock()
+            .await
+            .register(Box::new(archival_status_counter.clone()))
+            .unwrap();
+
         Metrics {
             db_poll_counter,
             network_request_counter,
+            archival_status_counter, // Add the new CounterVec to the struct
             registry,
         }
     }
+
     pub async fn push_metrics(&self) {
         let registry = self.registry.clone();
         spawn_blocking(move || {
@@ -54,4 +70,11 @@ impl Metrics {
         .await
         .unwrap();
     }
+
+    pub async fn record_archival_status(&self, status: &str) {
+        self.archival_status_counter
+            .with_label_values(&[status])
+            .inc();
+        self.push_metrics().await;
+    }
 }