Repository: cloudflare/daphne
Pull request #296 (2023-04-14 20:23:59)

daphne: Add metric for inbound requests


Patch:
diff --git a/daphne/src/metrics.rs b/daphne/src/metrics.rs
index ec7ebcd..59daf58 100644
--- a/daphne/src/metrics.rs
+++ b/daphne/src/metrics.rs
@@ -10,6 +10,9 @@ use prometheus::{
 };
 
 pub struct DaphneMetrics {
+    /// Inbound request metrics: Successful requests served, broken down by type.
+    inbound_request_counter: IntCounterVec,
+
     /// Report metrics. How many reports have been rejected, aggregated, and collected. When
     /// a report is rejected, the failure type is recorded.
     report_counter: IntCounterVec,
@@ -28,6 +31,13 @@ impl DaphneMetrics {
             "".into()
         };
 
+        let inbound_request_counter = register_int_counter_vec_with_registry!(
+            format!("{front}inbound_request_counter"),
+            "Total number of successful inbound requests.",
+            &["host", "type"],
+            registry
+        )?;
+
         let report_counter = register_int_counter_vec_with_registry!(
             format!("{front}report_counter"),
             "Total number reports rejected, aggregated, and collected.",
@@ -43,6 +53,7 @@ impl DaphneMetrics {
         )?;
 
         Ok(Self {
+            inbound_request_counter,
             report_counter,
             aggregation_job_gauge,
         })
@@ -62,6 +73,20 @@ pub struct ContextualizedDaphneMetrics<'req> {
 }
 
 impl ContextualizedDaphneMetrics<'_> {
+    pub fn inbound_req_inc(&self, request_type: DaphneRequestType) {
+        let request_type_str = match request_type {
+            DaphneRequestType::HpkeConfig => "hpke_config",
+            DaphneRequestType::Upload => "upload",
+            DaphneRequestType::Aggregate => "aggregate",
+            DaphneRequestType::Collect => "collect",
+        };
+
+        self.metrics
+            .inbound_request_counter
+            .with_label_values(&[self.host, request_type_str])
+            .inc();
+    }
+
     pub fn report_inc_by(&self, status: &str, val: u64) {
         self.metrics
             .report_counter
@@ -83,3 +108,15 @@ impl ContextualizedDaphneMetrics<'_> {
             .dec();
     }
 }
+
+#[derive(Clone, Copy, Debug)]
+pub enum DaphneRequestType {
+    /// DAP request for fetching the Aggregator's HPKE config.
+    HpkeConfig,
+    /// DAP upload request.
+    Upload,
+    /// DAP aggregate request.
+    Aggregate,
+    /// DAP collect request.
+    Collect,
+}
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index ded20e6..8ef1a8a 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -13,7 +13,7 @@ use crate::{
         PartialBatchSelector, Query, Report, ReportId, ReportMetadata, TaskId, Time,
         TransitionFailure, TransitionVar,
     },
-    metrics::DaphneMetrics,
+    metrics::{DaphneMetrics, DaphneRequestType},
     DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
     DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition, DapOutputShare,
     DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig, DapVersion,
@@ -151,6 +151,8 @@ where
             return Err(DapAbort::version_unknown());
         }
 
+        let metrics = self.metrics().with_host(req.host());
+
         // Parse the task ID from the query string, ensuring that it is the only query parameter.
         let mut id = None;
         for (k, v) in req.url.query_pairs() {
@@ -195,6 +197,7 @@ where
             _ => unreachable!("unhandled version {:?}", req.version),
         };
 
+        metrics.inbound_req_inc(DaphneRequestType::HpkeConfig);
         Ok(DapResponse {
             version: req.version,
             media_type: DapMediaType::HpkeConfigList,
@@ -309,6 +312,7 @@ where
     /// Handle HTTP POST to `/upload`. The input is the encoded report sent in the body of the HTTP
     /// request.
     async fn http_post_upload(&'srv self, req: &'req DapRequest<S>) -> Result<(), DapAbort> {
+        let metrics = self.metrics().with_host(req.host());
         debug!("upload for task {}", req.task_id()?);
 
         // Check whether the DAP version indicated by the sender is supported.
@@ -362,16 +366,20 @@ where
         // Store the report for future processing. At this point, the report may be rejected if
         // the Leader detects that the report was replayed or pertains to a batch that has already
         // been collected.
-        Ok(self.put_report(&report, req.task_id()?).await?)
+        self.put_report(&report, req.task_id()?).await?;
+
+        metrics.inbound_req_inc(DaphneRequestType::Upload);
+        Ok(())
     }
 
     /// Handle HTTP POST to `/collect`. The input is a [`CollectReq`](crate::messages::CollectReq).
     /// The return value is a URI that the Collector can poll later on to get the corresponding
     /// [`CollectResp`](crate::messages::CollectResp).
     async fn http_post_collect(&'srv self, req: &'req DapRequest<S>) -> Result<Url, DapAbort> {
+        let now = self.get_current_time();
+        let metrics = self.metrics().with_host(req.host());
         let task_id = req.task_id()?;
         debug!("collect for task {task_id}");
-        let now = self.get_current_time();
 
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
@@ -443,9 +451,12 @@ where
             _ => unreachable!("unhandled resource {:?}", req.resource),
         };
 
-        Ok(self
+        let collect_job_uri = self
             .init_collect_job(task_id, &collect_job_id, &collect_req)
-            .await?)
+            .await?;
+
+        metrics.inbound_req_inc(DaphneRequestType::Collect);
+        Ok(collect_job_uri)
     }
 
     /// Run the aggregation sub-protocol for the given set of reports. Return the number of reports
@@ -779,12 +790,13 @@ where
         &'srv self,
         req: &'req DapRequest<S>,
     ) -> Result<DapResponse, DapAbort> {
+        let metrics = self.metrics().with_host(req.host());
+
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
             return Err(DapAbort::version_unknown());
         }
 
-        let metrics = self.metrics().with_host(req.host());
         let task_id = req.task_id()?;
 
         if let Some(reason) = self.unauthorized_reason(req).await? {
@@ -952,6 +964,7 @@ where
                 };
 
                 metrics.agg_job_inc();
+                metrics.inbound_req_inc(DaphneRequestType::Aggregate);
                 Ok(DapResponse {
                     version: req.version,
                     media_type: DapMediaType::AggregationJobResp,
@@ -1021,6 +1034,7 @@ where
 
                 metrics.report_inc_by("aggregated", out_shares_count);
                 metrics.agg_job_dec();
+                metrics.inbound_req_inc(DaphneRequestType::Aggregate);
                 Ok(DapResponse {
                     version: req.version,
                     media_type: DapMediaType::agg_job_cont_resp_for_version(task_config.version),
@@ -1041,6 +1055,7 @@ where
         req: &'req DapRequest<S>,
     ) -> Result<DapResponse, DapAbort> {
         let now = self.get_current_time();
+        let metrics = self.metrics().with_host(req.host());
 
         // Check whether the DAP version indicated by the sender is supported.
         if req.version == DapVersion::Unknown {
@@ -1049,7 +1064,6 @@ where
 
         check_request_content_type(req, DapMediaType::AggregateShareReq)?;
 
-        let metrics = self.metrics().with_host(req.host());
         let task_id = req.task_id()?;
 
         if let Some(reason) = self.unauthorized_reason(req).await? {
@@ -1133,6 +1147,7 @@ where
         };
 
         metrics.report_inc_by("collected", agg_share_req.report_count);
+        metrics.inbound_req_inc(DaphneRequestType::Collect);
         Ok(DapResponse {
             version: req.version,
             media_type: DapMediaType::AggregateShare,


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 54564f7..34503cf 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -982,6 +982,7 @@ async fn http_post_aggregate_failure_report_replayed(version: DapVersion) {
 
     assert_metrics_include!(t.prometheus_registry, {
         r#"test_helper_report_counter{host="helper.org",status="rejected_report_replayed"}"#: 1,
+        r#"test_helper_inbound_request_counter{host="helper.org",type="aggregate"}"#: 1,
         r#"test_helper_aggregation_job_gauge{host="helper.org"}"#: 1,
     });
 }
@@ -1039,6 +1040,7 @@ async fn http_post_aggregate_failure_batch_collected(version: DapVersion) {
 
     assert_metrics_include!(t.prometheus_registry, {
         r#"test_helper_report_counter{host="helper.org",status="rejected_batch_collected"}"#: 1,
+        r#"test_helper_inbound_request_counter{host="helper.org",type="aggregate"}"#: 1,
         r#"test_helper_aggregation_job_gauge{host="helper.org"}"#: 1,
     });
 }
@@ -1571,6 +1573,8 @@ async fn e2e_time_interval(version: DapVersion) {
     t.run_col_job(task_id, &query).await.unwrap();
 
     assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_inbound_request_counter{host="helper.org",type="aggregate"}"#: 2,
+        r#"test_helper_inbound_request_counter{host="helper.org",type="collect"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="aggregated"}"#: 1,
         r#"test_helper_report_counter{host="helper.org",status="aggregated"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="collected"}"#: 1,
@@ -1602,6 +1606,8 @@ async fn e2e_fixed_size(version: DapVersion) {
     t.run_col_job(task_id, &query).await.unwrap();
 
     assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_inbound_request_counter{host="helper.org",type="aggregate"}"#: 2,
+        r#"test_helper_inbound_request_counter{host="helper.org",type="collect"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="aggregated"}"#: 1,
         r#"test_helper_report_counter{host="helper.org",status="aggregated"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="collected"}"#: 1,
@@ -1701,6 +1707,8 @@ async fn e2e_taskprov(version: DapVersion) {
     t.run_col_job(&taskprov_id, &query).await.unwrap();
 
     assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_inbound_request_counter{host="helper.org",type="aggregate"}"#: 2,
+        r#"test_helper_inbound_request_counter{host="helper.org",type="collect"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="aggregated"}"#: 1,
         r#"test_helper_report_counter{host="helper.org",status="aggregated"}"#: 1,
         r#"test_leader_report_counter{host="leader.com",status="collected"}"#: 1,