Add db entry metrics (#1412)

Author: spacebear21

payjoin-mailroom/src/db/mod.rs

@@ -8,6 +8,7 @@ use tower::util::BoxCloneSyncService;
 use tower::{Service, ServiceExt};
 
 pub mod files;
+use crate::metrics::{MetricsService, PayjoinVersion};
 
 pub trait SendableError:
     std::error::Error + std::marker::Send + std::marker::Sync + std::convert::Into<anyhow::Error>
@@ -235,3 +236,55 @@ impl Db for DbServiceAdapter {
 }
 
 pub use files::FilesDb;
+
+/// Db decorator that records write metrics
+#[derive(Clone)]
+pub struct MetricsDb<D: Db> {
+    inner: D,
+    metrics: MetricsService,
+}
+
+impl<D: Db> MetricsDb<D> {
+    pub fn new(inner: D, metrics: MetricsService) -> Self { Self { inner, metrics } }
+}
+
+impl<D: Db> Db for MetricsDb<D> {
+    type OperationalError = D::OperationalError;
+
+    async fn post_v2_payload(
+        &self,
+        mailbox_id: &ShortId,
+        data: Vec<u8>,
+    ) -> Result<Option<()>, Error<Self::OperationalError>> {
+        let result = self.inner.post_v2_payload(mailbox_id, data).await?;
+        if result.is_some() {
+            self.metrics.record_db_entry(PayjoinVersion::Two);
+        }
+        Ok(result)
+    }
+
+    async fn wait_for_v2_payload(
+        &self,
+        mailbox_id: &ShortId,
+    ) -> Result<Arc<Vec<u8>>, Error<Self::OperationalError>> {
+        self.inner.wait_for_v2_payload(mailbox_id).await
+    }
+
+    async fn post_v1_response(
+        &self,
+        mailbox_id: &ShortId,
+        data: Vec<u8>,
+    ) -> Result<(), Error<Self::OperationalError>> {
+        self.inner.post_v1_response(mailbox_id, data).await
+    }
+
+    async fn post_v1_request_and_wait_for_response(
+        &self,
+        mailbox_id: &ShortId,
+        data: Vec<u8>,
+    ) -> Result<Arc<Vec<u8>>, Error<Self::OperationalError>> {
+        let result = self.inner.post_v1_request_and_wait_for_response(mailbox_id, data).await?;
+        self.metrics.record_db_entry(PayjoinVersion::One);
+        Ok(result)
+    }
+}

payjoin-mailroom/src/directory.rs

@@ -806,4 +806,63 @@ mod tests {
         let html = landing_page_html();
         assert!(!html.contains("{{VERSION_STRING}}"));
     }
+
+    // MetricsDb decorator
+
+    #[tokio::test]
+    async fn post_mailbox_increments_v2_db_entry_metric() {
+        use opentelemetry_sdk::metrics::{
+            InMemoryMetricExporter, PeriodicReader, SdkMeterProvider,
+        };
+
+        use crate::db::MetricsDb;
+        use crate::metrics::{MetricsService, DB_ENTRIES};
+
+        let exporter = InMemoryMetricExporter::default();
+        let reader = PeriodicReader::builder(exporter.clone()).build();
+        let provider = SdkMeterProvider::builder().with_reader(reader).build();
+        let metrics = MetricsService::new(Some(provider.clone()));
+
+        let dir = tempfile::tempdir().expect("tempdir");
+        let db = FilesDb::init(Duration::from_millis(100), dir.keep()).await.expect("db init");
+        let db = MetricsDb::new(db, metrics);
+        let ohttp: ohttp::Server =
+            crate::key_config::gen_ohttp_server_config().expect("ohttp config").into();
+        let svc = Service::new(db, ohttp, SentinelTag::new([0u8; 32]), None);
+
+        let id = valid_short_id_path();
+        let res = svc
+            .post_mailbox(&id, Body::from(b"small payload under 65k limit".to_vec()))
+            .await
+            .expect("post_mailbox should succeed");
+        assert_eq!(res.status(), StatusCode::OK);
+
+        provider.force_flush().expect("flush failed");
+        let finished = exporter.get_finished_metrics().expect("metrics");
+        let db_metric = finished
+            .iter()
+            .flat_map(|rm| rm.scope_metrics())
+            .flat_map(|sm| sm.metrics())
+            .find(|m| m.name() == DB_ENTRIES)
+            .expect("missing db_entries_total metric");
+
+        use opentelemetry::KeyValue;
+        use opentelemetry_sdk::metrics::data::{AggregatedMetrics, MetricData};
+
+        // This checks that  counter value is 1 as post_mailbox was called once
+        // Also confirms the v2 label is recorded
+        match db_metric.data() {
+            AggregatedMetrics::U64(MetricData::Sum(sum)) => {
+                let points: Vec<_> = sum.data_points().collect();
+                assert_eq!(points.len(), 1, "expected exactly one data point");
+                assert_eq!(points[0].value(), 1, "expected counter value of 1");
+                let attrs: Vec<_> = points[0].attributes().collect();
+                assert!(
+                    attrs.contains(&&KeyValue::new("version", "2")),
+                    "expected version=2 attribute"
+                );
+            }
+            other => panic!("expected U64 Sum, got {other:?}"),
+        }
+    }
 }

payjoin-mailroom/src/lib.rs

@@ -29,9 +29,12 @@ pub mod ohttp_relay;
 use crate::metrics::MetricsService;
 use crate::middleware::{track_connections, track_metrics};
 
+type DirectoryService =
+    crate::directory::Service<crate::db::MetricsDb<crate::db::DbServiceAdapter>>;
+
 #[derive(Clone)]
 struct Services {
-    directory: crate::directory::Service<crate::db::DbServiceAdapter>,
+    directory: DirectoryService,
     relay: crate::ohttp_relay::Service,
     metrics: MetricsService,
     #[cfg(feature = "access-control")]
@@ -40,16 +43,17 @@ struct Services {
 
 pub async fn serve(config: Config, meter_provider: Option<SdkMeterProvider>) -> anyhow::Result<()> {
     let sentinel_tag = generate_sentinel_tag();
+    let metrics = MetricsService::new(meter_provider);
 
     #[cfg(feature = "access-control")]
     let geoip = init_geoip(&config).await?;
 
-    let directory = init_directory(&config, sentinel_tag).await?;
+    let directory = init_directory(&config, sentinel_tag, &metrics).await?;
 
     let services = Services {
         directory,
         relay: crate::ohttp_relay::Service::new(sentinel_tag).await,
-        metrics: MetricsService::new(meter_provider),
+        metrics,
         #[cfg(feature = "access-control")]
         geoip,
     };
@@ -84,11 +88,12 @@ pub async fn serve_manual_tls(
     use std::net::SocketAddr;
 
     let sentinel_tag = generate_sentinel_tag();
+    let metrics = MetricsService::new(None);
 
     #[cfg(feature = "access-control")]
     let geoip = init_geoip(&config).await?;
 
-    let directory = init_directory(&config, sentinel_tag).await?;
+    let directory = init_directory(&config, sentinel_tag, &metrics).await?;
 
     let services = Services {
         directory,
@@ -98,7 +103,7 @@ pub async fn serve_manual_tls(
             default_gateway,
         )
         .await,
-        metrics: MetricsService::new(None),
+        metrics,
         #[cfg(feature = "access-control")]
         geoip,
     };
@@ -153,16 +158,17 @@ pub async fn serve_acme(
         .ok_or_else(|| anyhow::anyhow!("ACME configuration is required for serve_acme"))?;
 
     let sentinel_tag = generate_sentinel_tag();
+    let metrics = MetricsService::new(meter_provider);
 
     #[cfg(feature = "access-control")]
     let geoip = init_geoip(&config).await?;
 
-    let directory = init_directory(&config, sentinel_tag).await?;
+    let directory = init_directory(&config, sentinel_tag, &metrics).await?;
 
     let services = Services {
         directory,
         relay: crate::ohttp_relay::Service::new(sentinel_tag).await,
-        metrics: MetricsService::new(meter_provider),
+        metrics,
         #[cfg(feature = "access-control")]
         geoip,
     };
@@ -220,10 +226,11 @@ impl Connected<IncomingStream<'_, Listener>> for middleware::MaybePeerIp {
 async fn init_directory(
     config: &Config,
     sentinel_tag: SentinelTag,
-) -> anyhow::Result<crate::directory::Service<crate::db::DbServiceAdapter>> {
+    metrics: &MetricsService,
+) -> anyhow::Result<DirectoryService> {
     let files_db = crate::db::FilesDb::init(config.timeout, config.storage_dir.clone()).await?;
     files_db.spawn_background_prune().await;
-    let db = crate::db::DbServiceAdapter::new(files_db);
+    let db = crate::db::MetricsDb::new(crate::db::DbServiceAdapter::new(files_db), metrics.clone());
 
     let ohttp_keys_dir = config.storage_dir.join("ohttp-keys");
     let ohttp_config = init_ohttp_config(&ohttp_keys_dir)?;
@@ -538,10 +545,11 @@ mod tests {
         );
 
         let sentinel_tag = generate_sentinel_tag();
+        let metrics = MetricsService::new(Some(provider.clone()));
         let services = Services {
-            directory: init_directory(&config, sentinel_tag).await.unwrap(),
+            directory: init_directory(&config, sentinel_tag, &metrics).await.unwrap(),
             relay: crate::ohttp_relay::Service::new(sentinel_tag).await,
-            metrics: MetricsService::new(Some(provider.clone())),
+            metrics,
             #[cfg(feature = "access-control")]
             geoip: None,
         };

payjoin-mailroom/src/metrics.rs

@@ -1,10 +1,13 @@
+use std::fmt;
+
 use opentelemetry::metrics::{Counter, MeterProvider, UpDownCounter};
 use opentelemetry::KeyValue;
 use opentelemetry_sdk::metrics::SdkMeterProvider;
 
 pub(crate) const TOTAL_CONNECTIONS: &str = "total_connections";
 pub(crate) const ACTIVE_CONNECTIONS: &str = "active_connections";
 pub(crate) const HTTP_REQUESTS: &str = "http_request_total";
+pub(crate) const DB_ENTRIES: &str = "db_entries_total";
 
 #[derive(Clone)]
 pub struct MetricsService {
@@ -14,6 +17,21 @@ pub struct MetricsService {
     total_connections: Counter<u64>,
     /// Number of active connections right now
     active_connections: UpDownCounter<i64>,
+    /// Total v1/v2 mailbox entries written, labelled by `version`
+    db_entries_total: Counter<u64>,
+}
+
+#[repr(u8)]
+#[derive(Clone, Copy, PartialEq, Eq)]
+pub enum PayjoinVersion {
+    /// BIP 78 Payjoin
+    One = 1,
+    /// BIP 77 Async Payjoin
+    Two = 2,
+}
+
+impl fmt::Display for PayjoinVersion {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { (*self as u8).fmt(f) }
 }
 
 impl MetricsService {
@@ -36,7 +54,12 @@ impl MetricsService {
             .with_description("Number of active connections")
             .build();
 
-        Self { http_requests_total, total_connections, active_connections }
+        let db_entries_total = meter
+            .u64_counter(DB_ENTRIES)
+            .with_description("Total mailbox entries stored by protocol version")
+            .build();
+
+        Self { http_requests_total, total_connections, active_connections, db_entries_total }
     }
 
     pub fn record_http_request(&self, endpoint: &str, method: &str, status_code: u16) {
@@ -56,4 +79,8 @@ impl MetricsService {
     }
 
     pub fn record_connection_close(&self) { self.active_connections.add(-1, &[]); }
+
+    pub fn record_db_entry(&self, version: PayjoinVersion) {
+        self.db_entries_total.add(1, &[KeyValue::new("version", version.to_string())]);
+    }
 }