Repository: cloudflare/daphne
Pull request #206 (2023-01-23 16:26:19)

Implement fmt::Display for Id and ReportId, instead of ToString.
This is a tiny PR that implements fmt::Display instead of ToString for Id, and adds fmt::Display for ReportId as well.  This is the recommended practice, and implementing fmt::Display gives you a ToString implementation automatically.  This lets us have nicer tracing code, e.g. instead of 

    trace!("report {} is ok", report.metadata.id.to_string());

We can say:

    trace!("report {} is ok", report.metadata.id);


Patch:
diff --git a/daphne/src/messages/mod.rs b/daphne/src/messages/mod.rs
index 9724fdd..7dec038 100644
--- a/daphne/src/messages/mod.rs
+++ b/daphne/src/messages/mod.rs
@@ -14,7 +14,7 @@ use serde::{Deserialize, Serialize};
 use std::{
     collections::HashSet,
     convert::{TryFrom, TryInto},
-    fmt::Debug,
+    fmt,
     io::{Cursor, Read},
 };
 
@@ -71,9 +71,9 @@ impl AsRef<[u8]> for Id {
     }
 }
 
-impl ToString for Id {
-    fn to_string(&self) -> String {
-        self.to_hex()
+impl fmt::Display for Id {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.to_hex())
     }
 }
 
@@ -88,6 +88,13 @@ pub type Time = u64;
 #[allow(missing_docs)]
 pub struct ReportId(pub [u8; 16]);
 
+impl ReportId {
+    /// Return the ID encoded as a hex string.
+    pub fn to_hex(&self) -> String {
+        hex::encode(self.0)
+    }
+}
+
 impl Encode for ReportId {
     fn encode(&self, bytes: &mut Vec<u8>) {
         bytes.extend_from_slice(&self.0);
@@ -108,6 +115,12 @@ impl AsRef<[u8]> for ReportId {
     }
 }
 
+impl fmt::Display for ReportId {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(f, "{}", self.to_hex())
+    }
+}
+
 /// Report extensions.
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
 #[serde(rename_all = "snake_case")]
@@ -552,16 +565,16 @@ impl Decode for TransitionFailure {
 impl std::fmt::Display for TransitionFailure {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
         match self {
-            Self::BatchCollected => write!(f, "batch-collected({})", *self as u8),
-            Self::ReportReplayed => write!(f, "report-replayed({})", *self as u8),
-            Self::ReportDropped => write!(f, "report-dropped({})", *self as u8),
-            Self::HpkeUnknownConfigId => write!(f, "hpke-unknown-config-id({})", *self as u8),
-            Self::HpkeDecryptError => write!(f, "hpke-decrypt-error({})", *self as u8),
-            Self::VdafPrepError => write!(f, "vdaf-prep-error({})", *self as u8),
-            Self::BatchSaturated => write!(f, "batch-saturated({})", *self as u8),
-            Self::TaskExpired => write!(f, "task-expired({})", *self as u8),
-            Self::UnrecognizedMessage => write!(f, "unrecognized-message({})", *self as u8),
-            Self::ReportTooEarly => write!(f, "report-too-early({})", *self as u8),
+            Self::BatchCollected => write!(f, "batch_collected"),
+            Self::ReportReplayed => write!(f, "report_replayed"),
+            Self::ReportDropped => write!(f, "report_dropped"),
+            Self::HpkeUnknownConfigId => write!(f, "hpke_unknown_config_id"),
+            Self::HpkeDecryptError => write!(f, "hpke_decrypt_error"),
+            Self::VdafPrepError => write!(f, "vdaf_prep_error"),
+            Self::BatchSaturated => write!(f, "batch_saturated"),
+            Self::TaskExpired => write!(f, "task_expired"),
+            Self::UnrecognizedMessage => write!(f, "unrecognized_message"),
+            Self::ReportTooEarly => write!(f, "report_too_early"),
         }
     }
 }
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 1855542..7b430fa 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -393,8 +393,6 @@ where
 
         // Filter out early rejected reports.
         //
-        // TODO Add metrics for early rejects.
-        //
         // TODO Add a test similar to http_post_aggregate_init_expired_task() in roles_test.rs that
         // verifies that the Leader properly checks for expiration. This will require extending the
         // test framework to run run_agg_job() directly.
@@ -408,10 +406,14 @@ where
         let reports = reports
             .into_iter()
             .filter(|report| {
-                // Filter the reports, removing any with a timestamps beyond the expiry or that was
-                // early rejected.
-                report.metadata.time < task_config.expiration
-                    && early_rejects.get(&report.metadata.id).is_none()
+                if let Some(failure) = early_rejects.get(&report.metadata.id) {
+                    self.metrics()
+                        .report_counter
+                        .with_label_values(&[&format!("rejected_{}", failure)])
+                        .inc();
+                    return false;
+                }
+                true
             })
             .collect();
 
@@ -728,32 +730,46 @@ where
                     ));
                 }
 
-                // Remove reports that are rejected early.
-                let early_rejects = early_rejects_future.await?;
                 let agg_resp = match transition {
                     DapHelperTransition::Continue(mut state, mut agg_resp) => {
-                        let mut i = 0;
-                        while i < state.seq.len() {
-                            let (_vdaf_state, _time, report_id) = &state.seq[i];
-
-                            // TODO Emit metrics for failure reasons.
-                            if let Some(failure) =
-                                early_rejects.get(&agg_resp.transitions[i].report_id)
+                        // Filter out early rejected reports.
+                        let early_rejects = early_rejects_future.await?;
+                        let mut state_index = 0;
+                        for transition in agg_resp.transitions.iter_mut() {
+                            let early_failure = early_rejects.get(&transition.report_id);
+                            if !matches!(transition.var, TransitionVar::Failed(..))
+                                && early_failure.is_some()
                             {
-                                // Mark reports that were rejected early as TransitionVar::Failed.
-                                agg_resp.transitions[i].var = TransitionVar::Failed(*failure);
+                                // NOTE(cjpatton) Clippy wants us to use and `if let` statement to
+                                // unwrap `early_failure`. I don't think this works becauase we
+                                // only want to enter this loop if `early_failure.is_some()` and
+                                // the current `transition` is not a failure. As far as I know, `if
+                                // let` statements can't yet be combined with other conditions.
+                                #[allow(clippy::unnecessary_unwrap)]
+                                let failure = early_failure.unwrap();
+                                transition.var = TransitionVar::Failed(*failure);
 
                                 // Remove VDAF preparation state of reports that were rejected early.
-                                if report_id == &agg_resp.transitions[i].report_id {
-                                    let _val = state.seq.remove(i);
+                                if transition.report_id == state.seq[state_index].2 {
+                                    let _val = state.seq.remove(state_index);
                                 } else {
                                     // The report ID in the Helper state and Aggregate response
-                                    // must be aligned. Abort with an internal error if this is not
-                                    // the case.
+                                    // must be aligned. If not, handle as an internal error.
                                     return Err(DapError::fatal("report IDs not aligned").into());
                                 }
+
+                                // NOTE(cjpatton) Unlike the Leader, the Helper filters out early
+                                // rejects after processing all of the reports. (This is an
+                                // optimization intended to reduce latency.) To avoid overcounting
+                                // rejection metrics, the latter rejections take precedence. The
+                                // Leader has the opposite behavior: Early rejections are resolved
+                                // first, so take precedence.
+                                self.metrics()
+                                    .report_counter
+                                    .with_label_values(&[&format!("rejected_{}", failure)])
+                                    .inc();
                             } else {
-                                i += 1;
+                                state_index += 1;
                             }
                         }
 
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 4687ac6..d25c484 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -501,10 +501,6 @@ impl VdafConfig {
     ///
     /// * `agg_init_req` is the request sent by the Leader.
     ///
-    /// * `early_rejects` is a tableindicating the set of reports in `agg_init_req` that the Helper
-    /// knows in advance it must reject. Each key is the report ID and the corresponding value is
-    /// the transition failure the Helper is to transmit.
-    ///
     /// * `version` is the DapVersion to use.
     pub(crate) async fn handle_agg_init_req(
         &self,


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 4e61c91..beb8bd3 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -865,6 +865,10 @@ async fn http_post_aggregate_failure_report_replayed(version: DapVersion) {
         transition.var,
         TransitionVar::Failed(TransitionFailure::ReportReplayed)
     );
+
+    assert_metrics_include!(t.helper_prometheus_registry, {
+        r#"daphne_report_counter{status="rejected_report_replayed"}"#: 1,
+    });
 }
 
 async_test_versions! { http_post_aggregate_failure_report_replayed }
@@ -915,6 +919,10 @@ async fn http_post_aggregate_failure_batch_collected(version: DapVersion) {
         transition.var,
         TransitionVar::Failed(TransitionFailure::BatchCollected)
     );
+
+    assert_metrics_include!(t.helper_prometheus_registry, {
+        r#"daphne_report_counter{status="rejected_batch_collected"}"#: 1,
+    });
 }
 
 async_test_versions! { http_post_aggregate_failure_batch_collected }
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 2b2317d..627a711 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -22,7 +22,7 @@ use prio::vdaf::{
     PrepareTransition,
 };
 use rand::prelude::*;
-use std::{collections::HashMap, fmt::Debug, time::SystemTime};
+use std::{fmt::Debug, time::SystemTime};
 use url::Url;
 
 impl<M: Debug> DapLeaderTransition<M> {
@@ -584,7 +584,6 @@ pub(crate) struct Test {
     task_config: DapTaskConfig,
     leader_hpke_receiver_config: HpkeReceiverConfig,
     helper_hpke_receiver_config: HpkeReceiverConfig,
-    early_rejects: HashMap<ReportId, TransitionFailure>,
     client_hpke_config_list: Vec<HpkeConfig>,
     collector_hpke_receiver_config: HpkeReceiverConfig,
 }
@@ -615,7 +614,6 @@ impl Test {
             agg_job_id,
             leader_hpke_receiver_config,
             helper_hpke_receiver_config,
-            early_rejects: HashMap::default(),
             client_hpke_config_list: vec![leader_hpke_config, helper_hpke_config],
             collector_hpke_receiver_config,
             task_config: DapTaskConfig {
@@ -686,14 +684,6 @@ impl Test {
             .await
             .unwrap();
 
-        for report_share in agg_init_req.report_shares {
-            // Make sure the Leader doesn't try to aggregate these reports again.
-            self.early_rejects.insert(
-                report_share.metadata.id.clone(),
-                TransitionFailure::ReportReplayed,
-            );
-        }
-
         agg_resp
     }