From 85c4cd20f0386edf515bf465f6c78dd48e2b810f Mon Sep 17 00:00:00 2001
From: Philip Jenvey <pjenvey@underboss.org>
Date: Tue, 20 Feb 2024 11:39:52 -0800
Subject: [PATCH] feat: have Bigtable match DynamoDB's metrics

(not including conversion errors which never occur)

Closes: SYNC-4150
---
 .../src/db/bigtable/bigtable_client/mod.rs    | 28 ++++++++++++++-----
 autopush-common/src/db/dual/mod.rs            | 10 ++++---
 autopush-common/src/db/dynamodb/mod.rs        |  3 ++
 3 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/autopush-common/src/db/bigtable/bigtable_client/mod.rs b/autopush-common/src/db/bigtable/bigtable_client/mod.rs
index 9d1256b7c..f7e854732 100644
--- a/autopush-common/src/db/bigtable/bigtable_client/mod.rs
+++ b/autopush-common/src/db/bigtable/bigtable_client/mod.rs
@@ -6,7 +6,7 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime};
 
 use async_trait::async_trait;
-use cadence::StatsdClient;
+use cadence::{CountedExt, StatsdClient};
 use futures_util::StreamExt;
 use google_cloud_rust_raw::bigtable::admin::v2::bigtable_table_admin::DropRowRangeRequest;
 use google_cloud_rust_raw::bigtable::admin::v2::bigtable_table_admin_grpc::BigtableTableAdminClient;
@@ -72,7 +72,7 @@ impl From<Uaid> for String {
 pub struct BigTableClientImpl {
     pub(crate) settings: BigTableDbSettings,
     /// Metrics client
-    _metrics: Arc<StatsdClient>,
+    metrics: Arc<StatsdClient>,
     /// Connection Channel (used for alternate calls)
     pool: BigTablePool,
     metadata: Metadata,
@@ -234,7 +234,7 @@ impl BigTableClientImpl {
         let admin_metadata = db_settings.admin_metadata()?;
         Ok(Self {
             settings: db_settings,
-            _metrics: metrics,
+            metrics: metrics,
             metadata,
             admin_metadata,
             pool,
@@ -928,7 +928,8 @@ impl DbClient for BigTableClientImpl {
 
         let mut cells: Vec<cell::Cell> = Vec::new();
 
-        let family = if message.topic.is_some() {
+        let is_topic = message.topic.is_some();
+        let family = if is_topic {
             MESSAGE_TOPIC_FAMILY
         } else {
             MESSAGE_FAMILY
@@ -973,7 +974,14 @@ impl DbClient for BigTableClientImpl {
         }
         row.add_cells(family, cells);
         trace!("🉑 Adding row");
-        self.write_row(row).await.map_err(|e| e.into())
+        self.write_row(row).await?;
+
+        self.metrics
+            .incr_with_tags("notification.message.stored")
+            .with_tag("topic", &is_topic.to_string())
+            .with_tag("database", &self.name())
+            .send();
+        Ok(())
     }
 
     /// Save a batch of messages to the database.
@@ -1022,7 +1030,8 @@ impl DbClient for BigTableClientImpl {
                 new_version_cell(expiry),
             ],
         );
-        self.write_row(row).await.map_err(|e| e.into())
+        self.write_row(row).await?;
+        Ok(())
     }
 
     /// Delete the notification from storage.
@@ -1034,7 +1043,12 @@ impl DbClient for BigTableClientImpl {
         );
         let row_key = format!("{}#{}", uaid.simple(), chidmessageid);
         debug!("🉑🔥 Deleting message {}", &row_key);
-        self.delete_row(&row_key).await.map_err(|e| e.into())
+        self.delete_row(&row_key).await?;
+        self.metrics
+            .incr_with_tags("notification.message.deleted")
+            .with_tag("database", &self.name())
+            .send();
+        Ok(())
     }
 
     /// Return `limit` pending messages from storage. `limit=0` for all messages.
diff --git a/autopush-common/src/db/dual/mod.rs b/autopush-common/src/db/dual/mod.rs
index 0f93853bf..5c998fd3c 100644
--- a/autopush-common/src/db/dual/mod.rs
+++ b/autopush-common/src/db/dual/mod.rs
@@ -99,6 +99,12 @@ impl DualClientImpl {
         let primary = BigTableClientImpl::new(metrics.clone(), &db_settings.primary)?;
         let secondary = DdbClientImpl::new(metrics.clone(), &db_settings.secondary)?;
         debug!("⚖ Got primary and secondary");
+        if let Some(median) = median {
+            metrics
+                .incr_with_tags("database.dual.allot")
+                .with_tag("median", &median.to_string())
+                .send();
+        }
         Ok(Self {
             primary,
             secondary: secondary.clone(),
@@ -128,10 +134,6 @@ impl DualClientImpl {
         } else {
             (Box::new(&self.primary), true)
         };
-        self.metrics
-            .incr_with_tags("database.dual.error")
-            .with_tag("target", &target.0.name())
-            .send();
         debug!("⚖ alloting to {}", target.0.name());
         Ok(target)
     }
diff --git a/autopush-common/src/db/dynamodb/mod.rs b/autopush-common/src/db/dynamodb/mod.rs
index fa8c18283..261c24158 100644
--- a/autopush-common/src/db/dynamodb/mod.rs
+++ b/autopush-common/src/db/dynamodb/mod.rs
@@ -547,6 +547,7 @@ impl DbClient for DdbClientImpl {
         self.metrics
             .incr_with_tags("notification.message.stored")
             .with_tag("topic", &topic)
+            .with_tag("database", &self.name())
             .send();
         Ok(())
     }
@@ -559,6 +560,7 @@ impl DbClient for DdbClientImpl {
                 self.metrics
                     .incr_with_tags("notification.message.stored")
                     .with_tag("topic", &n.topic.is_some().to_string())
+                    .with_tag("database", &self.name())
                     .send();
                 serde_dynamodb::to_hashmap(&NotificationRecord::from_notif(uaid, n))
                     .ok()
@@ -600,6 +602,7 @@ impl DbClient for DdbClientImpl {
             .await?;
         self.metrics
             .incr_with_tags("notification.message.deleted")
+            .with_tag("database", &self.name())
             .send();
         Ok(())
     }