Repository: cloudflare/daphne
Pull request #207 (2023-01-23 19:48:01)

Add report rejection metrics
Partially addresses #27.

Plumb `DapAggregator::metrics()` to `VdafConfig` and count report rejections. Wherever the aggregator transitions a report to "failued", record the failure reason.

Relatedly:

* Refactor `VdafConfig::produce_report_with_extensions()` to allow us to construct invalid input shares in tests.

* Add a prefix to the metrics registered by `DaphneMetrics::register()`, allowing multiple instances of `DaphneMetrics` to be registered with the same registry. This simplifies the testing framework a bit.

* In-line variable names in a few `format!`  instances.

* Move `assert_metrics_include!` macro to the `testing` module so that we can reuse it in `vdaf/mod_test.rs`. Also, only compile it in tests.

* Rename some tests in `vdaf/mod_test.rs` to more closely correspond to the condition being tested.


Patch:
diff --git a/daphne/src/metrics.rs b/daphne/src/metrics.rs
index c5acd18..96a522f 100644
--- a/daphne/src/metrics.rs
+++ b/daphne/src/metrics.rs
@@ -16,9 +16,9 @@ pub struct DaphneMetrics {
 
 impl DaphneMetrics {
     /// Regstier Daphne metrics with the specified registry.
-    pub fn register(registry: &Registry) -> Result<Self, DapError> {
+    pub fn register(registry: &Registry, prefix: &str) -> Result<Self, DapError> {
         let report_counter = register_int_counter_vec_with_registry!(
-            "daphne_report_counter",
+            format!("{prefix}_report_counter"),
             "Total number reports rejected, aggregated, and collected.",
             &["status"],
             registry
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 7b430fa..37224da 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -409,7 +409,7 @@ where
                 if let Some(failure) = early_rejects.get(&report.metadata.id) {
                     self.metrics()
                         .report_counter
-                        .with_label_values(&[&format!("rejected_{}", failure)])
+                        .with_label_values(&[&format!("rejected_{failure}")])
                         .inc();
                     return false;
                 }
@@ -428,6 +428,7 @@ where
                 &agg_job_id,
                 part_batch_sel,
                 reports,
+                self.metrics(),
             )
             .await?;
         let (state, agg_init_req) = match transition {
@@ -450,9 +451,13 @@ where
         let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
 
         // Prepare AggreagteContinueReq.
-        let transition = task_config
-            .vdaf
-            .handle_agg_resp(task_id, &agg_job_id, state, agg_resp)?;
+        let transition = task_config.vdaf.handle_agg_resp(
+            task_id,
+            &agg_job_id,
+            state,
+            agg_resp,
+            self.metrics(),
+        )?;
         let (uncommited, agg_cont_req) = match transition {
             DapLeaderTransition::Uncommitted(uncommited, agg_cont_req) => {
                 (uncommited, agg_cont_req)
@@ -475,9 +480,10 @@ where
         let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
 
         // Commit the output shares.
-        let out_shares = task_config
-            .vdaf
-            .handle_final_agg_resp(uncommited, agg_resp)?;
+        let out_shares =
+            task_config
+                .vdaf
+                .handle_final_agg_resp(uncommited, agg_resp, self.metrics())?;
         let out_shares_count = out_shares.len() as u64;
         self.put_out_shares(task_id, part_batch_sel, out_shares)
             .await?;
@@ -719,7 +725,7 @@ where
 
                 let transition = task_config
                     .vdaf
-                    .handle_agg_init_req(self, task_config, &agg_init_req)
+                    .handle_agg_init_req(self, task_config, &agg_init_req, self.metrics())
                     .await?;
 
                 // Check that helper state with task_id and agg_job_id does not exist.
@@ -766,7 +772,7 @@ where
                                 // first, so take precedence.
                                 self.metrics()
                                     .report_counter
-                                    .with_label_values(&[&format!("rejected_{}", failure)])
+                                    .with_label_values(&[&format!("rejected_{failure}")])
                                     .inc();
                             } else {
                                 state_index += 1;
@@ -813,7 +819,10 @@ where
                     .await?
                     .ok_or(DapAbort::UnrecognizedAggregationJob)?;
                 let part_batch_sel = state.part_batch_sel.clone();
-                let transition = task_config.vdaf.handle_agg_cont_req(state, &agg_cont_req)?;
+                let transition =
+                    task_config
+                        .vdaf
+                        .handle_agg_cont_req(state, &agg_cont_req, self.metrics())?;
 
                 let (agg_resp, out_shares_count) = match transition {
                     DapHelperTransition::Continue(..) => {
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index d25c484..2237c0c 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -12,6 +12,7 @@ use crate::{
         PlaintextInputShare, Report, ReportId, ReportMetadata, ReportShare, Time, Transition,
         TransitionFailure, TransitionVar,
     },
+    metrics::DaphneMetrics,
     vdaf::{
         prio2::{
             prio2_encode_prepare_message, prio2_helper_prepare_finish, prio2_leader_prepare_finish,
@@ -171,6 +172,30 @@ impl VdafConfig {
         measurement: DapMeasurement,
         extensions: Vec<Extension>,
         version: DapVersion,
+    ) -> Result<Report, DapError> {
+        let (public_share, input_shares) = self.produce_input_shares(measurement)?;
+        self.produce_report_with_extensions_for_shares(
+            public_share,
+            input_shares,
+            hpke_config_list,
+            time,
+            task_id,
+            extensions,
+            version,
+        )
+    }
+
+    /// Generate a report for the given public and input shares with the given extensions.
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) fn produce_report_with_extensions_for_shares(
+        &self,
+        public_share: Vec<u8>,
+        mut input_shares: Vec<Vec<u8>>,
+        hpke_config_list: &[HpkeConfig],
+        time: Time,
+        task_id: &Id,
+        extensions: Vec<Extension>,
+        version: DapVersion,
     ) -> Result<Report, DapError> {
         let mut rng = thread_rng();
         let report_extensions = match version {
@@ -183,24 +208,19 @@ impl VdafConfig {
             extensions: report_extensions,
         };
 
-        let public_share = Vec::new();
-        let mut encoded_input_shares = match self {
-            Self::Prio3(prio3_config) => prio3_shard(prio3_config, measurement)?,
-            Self::Prio2 { dimension } => prio2_shard(*dimension, measurement)?,
-        };
         if version != DapVersion::Draft02 {
             let mut encoded: Vec<Vec<u8>> = Vec::new();
-            for share in encoded_input_shares.into_iter() {
+            for share in input_shares.into_iter() {
                 let input_share = PlaintextInputShare {
                     extensions: extensions.clone(),
                     payload: share,
                 };
                 encoded.push(PlaintextInputShare::get_encoded(&input_share));
             }
-            encoded_input_shares = encoded;
+            input_shares = encoded;
         }
 
-        if hpke_config_list.len() != encoded_input_shares.len() {
+        if hpke_config_list.len() != input_shares.len() {
             return Err(DapError::Fatal("unexpected number of HPKE configs".into()));
         }
 
@@ -225,11 +245,9 @@ impl VdafConfig {
         // it here.
         encode_u32_bytes(&mut aad, &public_share);
 
-        let mut encrypted_input_shares = Vec::with_capacity(encoded_input_shares.len());
-        for (i, (hpke_config, input_share_data)) in hpke_config_list
-            .iter()
-            .zip(encoded_input_shares)
-            .enumerate()
+        let mut encrypted_input_shares = Vec::with_capacity(input_shares.len());
+        for (i, (hpke_config, input_share_data)) in
+            hpke_config_list.iter().zip(input_shares).enumerate()
         {
             info[n + 1] = if i == 0 {
                 CTX_ROLE_LEADER
@@ -253,6 +271,19 @@ impl VdafConfig {
         })
     }
 
+    /// Generate shares for a measurement.
+    pub(crate) fn produce_input_shares(
+        &self,
+        measurement: DapMeasurement,
+    ) -> Result<(Vec<u8>, Vec<Vec<u8>>), DapError> {
+        let public_share = Vec::new();
+        let input_shares = match self {
+            Self::Prio3(prio3_config) => prio3_shard(prio3_config, measurement)?,
+            Self::Prio2 { dimension } => prio2_shard(*dimension, measurement)?,
+        };
+        Ok((public_share, input_shares))
+    }
+
     /// Generate a report for a measurement. This method is run by the Client.
     ///
     /// # Inputs
@@ -319,6 +350,10 @@ impl VdafConfig {
             return Err(DapError::Transition(TransitionFailure::TaskExpired));
         }
 
+        if !public_share.is_empty() {
+            return Err(DapError::Transition(TransitionFailure::VdafPrepError));
+        }
+
         let input_share_text = match task_config.version {
             DapVersion::Draft02 => CTX_INPUT_SHARE_DRAFT02,
             DapVersion::Draft03 => CTX_INPUT_SHARE_DRAFT03,
@@ -406,6 +441,7 @@ impl VdafConfig {
         agg_job_id: &Id,
         part_batch_sel: &PartialBatchSelector,
         reports: Vec<Report>,
+        metrics: &DaphneMetrics,
     ) -> Result<DapLeaderTransition<AggregateInitializeReq>, DapAbort> {
         let mut processed = HashSet::with_capacity(reports.len());
         let mut states = Vec::with_capacity(reports.len());
@@ -457,9 +493,10 @@ impl VdafConfig {
                 }
 
                 // Skip report that can't be processed any further.
-                //
-                // TODO Emit metric for failure reason
-                Err(DapError::Transition(..)) => (),
+                Err(DapError::Transition(failure)) => metrics
+                    .report_counter
+                    .with_label_values(&[&format!("rejected_{failure}")])
+                    .inc(),
 
                 Err(e) => return Err(DapAbort::Internal(Box::new(e))),
             };
@@ -507,6 +544,7 @@ impl VdafConfig {
         decrypter: &impl HpkeDecrypter<'_>,
         task_config: &DapTaskConfig,
         agg_init_req: &AggregateInitializeReq,
+        metrics: &DaphneMetrics,
     ) -> Result<DapHelperTransition<AggregateResp>, DapAbort> {
         let num_reports = agg_init_req.report_shares.len();
         let mut processed = HashSet::with_capacity(num_reports);
@@ -543,7 +581,13 @@ impl VdafConfig {
                     TransitionVar::Continued(message_data)
                 }
 
-                Err(DapError::Transition(failure_reason)) => TransitionVar::Failed(failure_reason),
+                Err(DapError::Transition(failure)) => {
+                    metrics
+                        .report_counter
+                        .with_label_values(&[&format!("rejected_{failure}")])
+                        .inc();
+                    TransitionVar::Failed(failure)
+                }
 
                 Err(e) => return Err(DapAbort::Internal(Box::new(e))),
             };
@@ -581,6 +625,7 @@ impl VdafConfig {
         agg_job_id: &Id,
         state: DapLeaderState,
         agg_resp: AggregateResp,
+        metrics: &DaphneMetrics,
     ) -> Result<DapLeaderTransition<AggregateContinueReq>, DapAbort> {
         if agg_resp.transitions.len() != state.seq.len() {
             return Err(DapAbort::UnrecognizedMessage);
@@ -600,9 +645,13 @@ impl VdafConfig {
                 TransitionVar::Continued(message) => message,
 
                 // Skip report that can't be processed any further.
-                //
-                // TODO Log the reason the report was skipped.
-                TransitionVar::Failed(..) => continue,
+                TransitionVar::Failed(failure) => {
+                    metrics
+                        .report_counter
+                        .with_label_values(&[&format!("rejected_{failure}")])
+                        .inc();
+                    continue;
+                }
 
                 // TODO Log the fact that the helper sent an unexpected message.
                 TransitionVar::Finished => return Err(DapAbort::UnrecognizedMessage),
@@ -646,9 +695,13 @@ impl VdafConfig {
                 }
 
                 // Skip report that can't be processed any further.
-                //
-                // TODO Log the reason the report was skipped.
-                Err(VdafError::Codec(..)) | Err(VdafError::Vdaf(..)) => (),
+                Err(VdafError::Codec(..)) | Err(VdafError::Vdaf(..)) => {
+                    let failure = TransitionFailure::VdafPrepError;
+                    metrics
+                        .report_counter
+                        .with_label_values(&[&format!("rejected_{failure}")])
+                        .inc();
+                }
             };
         }
 
@@ -680,6 +733,7 @@ impl VdafConfig {
         &self,
         state: DapHelperState,
         agg_cont_req: &AggregateContinueReq,
+        metrics: &DaphneMetrics,
     ) -> Result<DapHelperTransition<AggregateResp>, DapAbort> {
         let mut processed = HashSet::with_capacity(state.seq.len());
         let mut recognized = HashSet::with_capacity(state.seq.len());
@@ -744,7 +798,12 @@ impl VdafConfig {
                     }
 
                     Err(VdafError::Codec(..)) | Err(VdafError::Vdaf(..)) => {
-                        TransitionVar::Failed(TransitionFailure::VdafPrepError)
+                        let failure = TransitionFailure::VdafPrepError;
+                        metrics
+                            .report_counter
+                            .with_label_values(&[&format!("rejected_{failure}")])
+                            .inc();
+                        TransitionVar::Failed(failure)
                     }
                 };
 
@@ -780,6 +839,7 @@ impl VdafConfig {
         &self,
         uncommitted: DapLeaderUncommitted,
         agg_resp: AggregateResp,
+        metrics: &DaphneMetrics,
     ) -> Result<Vec<DapOutputShare>, DapAbort> {
         if agg_resp.transitions.len() != uncommitted.seq.len() {
             return Err(DapAbort::UnrecognizedMessage);
@@ -801,9 +861,13 @@ impl VdafConfig {
                 TransitionVar::Continued(..) => return Err(DapAbort::UnrecognizedMessage),
 
                 // Skip report that can't be processed any further.
-                //
-                // TODO Log the reason the report was skipped.
-                TransitionVar::Failed(..) => continue,
+                TransitionVar::Failed(failure) => {
+                    metrics
+                        .report_counter
+                        .with_label_values(&[&format!("rejected_{failure}")])
+                        .inc();
+                    continue;
+                }
 
                 TransitionVar::Finished => out_shares.push(out_share),
             };
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index e452602..7e9ced6 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -316,7 +316,7 @@ impl DaphneWorkerState {
 
         // TODO(cjpatton) Push metrics to gateway after handling the request.
         let prometheus_registry = Registry::new();
-        let daphne_metrics = DaphneMetrics::register(&prometheus_registry)
+        let daphne_metrics = DaphneMetrics::register(&prometheus_registry, "daphne_worker")
             .map_err(|e| Error::RustError(format!("failed to register daphne metrics: {}", e)))?;
 
         Ok(Self {


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index beb8bd3..b1dcffb 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -2,7 +2,8 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use crate::{
-    async_test_version, async_test_versions,
+    assert_metrics_include, assert_metrics_include_auxiliary_function, async_test_version,
+    async_test_versions,
     auth::BearerToken,
     constants::{
         DRAFT02_MEDIA_TYPE_HPKE_CONFIG, MEDIA_TYPE_AGG_CONT_REQ, MEDIA_TYPE_AGG_INIT_REQ,
@@ -49,69 +50,16 @@ macro_rules! get_reports {
     }};
 }
 
-// Helper macro used by `assert_metrics_include`.
-macro_rules! build_metrics_list {
-    ($set:expr, $k:tt: $v:expr,) => {{
-        let line = format!("{} {}", $k, $v);
-        $set.insert(line);
-    }};
-
-    ($set:expr, $k:tt: $v:expr, $($ks:tt: $vs:expr),+,) => {{
-        let line = format!("{} {}", $k, $v);
-        $set.insert(line);
-        build_metrics_list!($set, $($ks: $vs),+)
-    }}
-}
-
-// Gather metrics from a registry and assert that a list of metrics are present and have the
-// correct value. For example:
-// ```
-// let registry = prometheus::Registry::new();
-//
-// // ... Register a metric called "report_counter" and use it.
-//
-// assert_metrics_include!(t.helper_prometheus_registry, {
-//      r#"report_counter{status="aggregated"}"#: 23,
-// });
-// ```
-macro_rules! assert_metrics_include {
-    ($registry:expr, {$($ks:tt: $vs:expr),+,}) => {{
-        use prometheus::{Encoder, TextEncoder};
-        use std::collections::HashSet;
-
-        let mut want: HashSet<String> = HashSet::new();
-        build_metrics_list!(&mut want, $($ks: $vs),+,);
-
-        // Encode the metrics and iterate over each line. For each line, if the line appears in the
-        // list of expected output lines, then remove it.
-        let mut got_buf = Vec::new();
-        let encoder = TextEncoder::new();
-        encoder.encode(&$registry.gather(), &mut got_buf).unwrap();
-        let got_str = String::from_utf8(got_buf).unwrap();
-        for line in got_str.split('\n') {
-            want.remove(line);
-        }
-
-        // The metrics contain the expected lines if the the set is now empty.
-        if !want.is_empty() {
-            panic!("unexpected metrics: got:\n{}\nmust contain:\n{}\n",
-                   got_str, want.into_iter().collect::<Vec<String>>().join("\n"));
-        }
-    }}
-}
-
 struct Test {
     now: Time,
     leader: Arc<MockAggregator>,
     helper: Arc<MockAggregator>,
-    #[allow(dead_code)]
-    leader_prometheus_registry: prometheus::Registry,
-    helper_prometheus_registry: prometheus::Registry,
     collector_token: BearerToken,
     time_interval_task_id: Id,
     fixed_size_task_id: Id,
     expired_task_id: Id,
     version: DapVersion,
+    prometheus_registry: prometheus::Registry,
 }
 
 impl Test {
@@ -202,11 +150,12 @@ impl Test {
         let mut taskprov_vdaf_verify_key_init = vec![0; 32];
         rng.fill(&mut taskprov_vdaf_verify_key_init[..]);
 
+        let prometheus_registry = prometheus::Registry::new();
+
         let helper_hpke_receiver_config_list = global_config
             .gen_hpke_receiver_config_list(rng.gen())
             .collect::<Result<Vec<HpkeReceiverConfig>, _>>()
             .expect("failed to generate HPKE receiver config");
-        let helper_prometheus_registry = prometheus::Registry::new();
         let helper = Arc::new(MockAggregator {
             global_config: global_config.clone(),
             tasks: Arc::new(Mutex::new(tasks.clone())),
@@ -219,7 +168,7 @@ impl Test {
             agg_store: Arc::new(Mutex::new(HashMap::new())),
             collector_hpke_config: collector_hpke_receiver_config.config.clone(),
             taskprov_vdaf_verify_key_init: taskprov_vdaf_verify_key_init.clone(),
-            metrics: DaphneMetrics::register(&helper_prometheus_registry).unwrap(),
+            metrics: DaphneMetrics::register(&prometheus_registry, "test_helper").unwrap(),
             peer: None,
         });
 
@@ -227,7 +176,6 @@ impl Test {
             .gen_hpke_receiver_config_list(rng.gen())
             .collect::<Result<Vec<HpkeReceiverConfig>, _>>()
             .expect("failed to generate HPKE receiver config");
-        let leader_prometheus_registry = prometheus::Registry::new();
         let leader = Arc::new(MockAggregator {
             global_config: global_config.clone(),
             tasks: Arc::new(Mutex::new(tasks.clone())),
@@ -240,7 +188,7 @@ impl Test {
             agg_store: Arc::new(Mutex::new(HashMap::new())),
             collector_hpke_config: collector_hpke_receiver_config.config.clone(),
             taskprov_vdaf_verify_key_init: taskprov_vdaf_verify_key_init.clone(),
-            metrics: DaphneMetrics::register(&leader_prometheus_registry).unwrap(),
+            metrics: DaphneMetrics::register(&prometheus_registry, "test_leader").unwrap(),
             peer: Some(Arc::clone(&helper)),
         });
 
@@ -248,13 +196,12 @@ impl Test {
             now,
             leader,
             helper,
-            leader_prometheus_registry,
-            helper_prometheus_registry,
             collector_token,
             time_interval_task_id,
             fixed_size_task_id,
             expired_task_id,
             version,
+            prometheus_registry,
         }
     }
 
@@ -866,8 +813,8 @@ async fn http_post_aggregate_failure_report_replayed(version: DapVersion) {
         TransitionVar::Failed(TransitionFailure::ReportReplayed)
     );
 
-    assert_metrics_include!(t.helper_prometheus_registry, {
-        r#"daphne_report_counter{status="rejected_report_replayed"}"#: 1,
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_report_counter{status="rejected_report_replayed"}"#: 1,
     });
 }
 
@@ -920,8 +867,8 @@ async fn http_post_aggregate_failure_batch_collected(version: DapVersion) {
         TransitionVar::Failed(TransitionFailure::BatchCollected)
     );
 
-    assert_metrics_include!(t.helper_prometheus_registry, {
-        r#"daphne_report_counter{status="rejected_batch_collected"}"#: 1,
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_report_counter{status="rejected_batch_collected"}"#: 1,
     });
 }
 
@@ -1441,11 +1388,9 @@ async fn e2e_time_interval(version: DapVersion) {
     let query = task_config.query_for_current_batch_window(t.now);
     t.run_col_job(task_id, &query).await.unwrap();
 
-    assert_metrics_include!(t.leader_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
-    });
-    assert_metrics_include!(t.helper_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="aggregated"}"#: 1,
+        r#"test_helper_report_counter{status="aggregated"}"#: 1,
     });
 }
 
@@ -1471,11 +1416,9 @@ async fn e2e_fixed_size(version: DapVersion) {
     };
     t.run_col_job(task_id, &query).await.unwrap();
 
-    assert_metrics_include!(t.leader_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
-    });
-    assert_metrics_include!(t.helper_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="aggregated"}"#: 1,
+        r#"test_helper_report_counter{status="aggregated"}"#: 1,
     });
 }
 
@@ -1566,11 +1509,9 @@ async fn e2e_taskprov(version: DapVersion) {
     };
     t.run_col_job(task_id, &query).await.unwrap();
 
-    assert_metrics_include!(t.leader_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
-    });
-    assert_metrics_include!(t.helper_prometheus_registry, {
-        r#"daphne_report_counter{status="aggregated"}"#: 1,
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="aggregated"}"#: 1,
+        r#"test_helper_report_counter{status="aggregated"}"#: 1,
     });
 }
 
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 4ffd8fd..edd21de 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -946,3 +946,61 @@ macro_rules! async_test_versions {
         )*
     };
 }
+
+/// Helper macro used by `assert_metrics_include`.
+//
+// TODO(cjpatton) Figure out how to bake this into `asssert_metrics_include` so that users don't
+// have to import both macros.
+#[cfg(test)]
+#[macro_export]
+macro_rules! assert_metrics_include_auxiliary_function {
+    ($set:expr, $k:tt: $v:expr,) => {{
+        let line = format!("{} {}", $k, $v);
+        $set.insert(line);
+    }};
+
+    ($set:expr, $k:tt: $v:expr, $($ks:tt: $vs:expr),+,) => {{
+        let line = format!("{} {}", $k, $v);
+        $set.insert(line);
+        assert_metrics_include_auxiliary_function!($set, $($ks: $vs),+,)
+    }}
+}
+
+/// Gather metrics from a registry and assert that a list of metrics are present and have the
+/// correct value. For example:
+/// ```
+/// let registry = prometheus::Registry::new();
+///
+/// // ... Register a metric called "report_counter" and use it.
+///
+/// assert_metrics_include!(t.helper_prometheus_registry, {
+///      r#"report_counter{status="aggregated"}"#: 23,
+/// });
+/// ```
+#[cfg(test)]
+#[macro_export]
+macro_rules! assert_metrics_include {
+    ($registry:expr, {$($ks:tt: $vs:expr),+,}) => {{
+        use prometheus::{Encoder, TextEncoder};
+        use std::collections::HashSet;
+
+        let mut want: HashSet<String> = HashSet::new();
+        assert_metrics_include_auxiliary_function!(&mut want, $($ks: $vs),+,);
+
+        // Encode the metrics and iterate over each line. For each line, if the line appears in the
+        // list of expected output lines, then remove it.
+        let mut got_buf = Vec::new();
+        let encoder = TextEncoder::new();
+        encoder.encode(&$registry.gather(), &mut got_buf).unwrap();
+        let got_str = String::from_utf8(got_buf).unwrap();
+        for line in got_str.split('\n') {
+            want.remove(line);
+        }
+
+        // The metrics contain the expected lines if the the set is now empty.
+        if !want.is_empty() {
+            panic!("unexpected metrics: got:\n{}\nmust contain:\n{}\n",
+                   got_str, want.into_iter().collect::<Vec<String>>().join("\n"));
+        }
+    }}
+}
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 627a711..e070e81 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -2,13 +2,15 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use crate::{
-    async_test_version, async_test_versions,
+    assert_metrics_include, assert_metrics_include_auxiliary_function, async_test_version,
+    async_test_versions,
     hpke::HpkeReceiverConfig,
     messages::{
         AggregateContinueReq, AggregateInitializeReq, AggregateResp, BatchSelector, HpkeAeadId,
         HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Interval, PartialBatchSelector,
-        Report, ReportId, Time, Transition, TransitionFailure, TransitionVar,
+        Report, ReportId, ReportShare, Time, Transition, TransitionFailure, TransitionVar,
     },
+    metrics::DaphneMetrics,
     test_version, test_versions, DapAbort, DapAggregateResult, DapAggregateShare, DapError,
     DapHelperState, DapHelperTransition, DapLeaderState, DapLeaderTransition, DapLeaderUncommitted,
     DapMeasurement, DapOutputShare, DapQueryConfig, DapTaskConfig, DapVersion, Prio3Config,
@@ -173,7 +175,7 @@ fn roundtrip_report_unsupported_hpke_suite(version: DapVersion) {
 
 test_versions! { roundtrip_report_unsupported_hpke_suite }
 
-async fn agg_init_req(version: DapVersion) {
+async fn produce_agg_init_req(version: DapVersion) {
     let mut t = Test::new(TEST_VDAF, version);
     let reports = t.produce_reports(vec![
         DapMeasurement::U64(1),
@@ -201,9 +203,9 @@ async fn agg_init_req(version: DapVersion) {
     }
 }
 
-async_test_versions! { agg_init_req }
+async_test_versions! { produce_agg_init_req }
 
-async fn agg_init_req_fail_hpke_decrypt_err(version: DapVersion) {
+async fn produce_agg_init_req_skip_hpke_decrypt_err(version: DapVersion) {
     let t = Test::new(TEST_VDAF, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
@@ -214,11 +216,15 @@ async fn agg_init_req_fail_hpke_decrypt_err(version: DapVersion) {
         t.produce_agg_init_req(reports).await,
         DapLeaderTransition::Skip
     );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="rejected_hpke_decrypt_error"}"#: 1,
+    });
 }
 
-async_test_versions! { agg_init_req_fail_hpke_decrypt_err }
+async_test_versions! { produce_agg_init_req_skip_hpke_decrypt_err }
 
-async fn agg_init_req_fail_hpke_decrypt_err_wrong_config_id(version: DapVersion) {
+async fn produce_agg_init_req_skip_hpke_unknown_config_id(version: DapVersion) {
     let t = Test::new(TEST_VDAF, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
@@ -229,11 +235,34 @@ async fn agg_init_req_fail_hpke_decrypt_err_wrong_config_id(version: DapVersion)
         t.produce_agg_init_req(reports).await,
         DapLeaderTransition::Skip
     );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="rejected_hpke_unknown_config_id"}"#: 1,
+    });
 }
 
-async_test_versions! { agg_init_req_fail_hpke_decrypt_err_wrong_config_id }
+async_test_versions! { produce_agg_init_req_skip_hpke_unknown_config_id }
+
+async fn produce_agg_init_req_skip_vdaf_prep_error(version: DapVersion) {
+    let t = Test::new(TEST_VDAF, version);
+    let reports = vec![
+        t.produce_invalid_report_public_share_decode_failure(DapMeasurement::U64(1), version),
+        t.produce_invalid_report_input_share_decode_failure(DapMeasurement::U64(1), version),
+    ];
+
+    assert_matches!(
+        t.produce_agg_init_req(reports).await,
+        DapLeaderTransition::Skip
+    );
 
-async fn agg_resp_fail_hpke_decrypt_err(version: DapVersion) {
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="rejected_vdaf_prep_error"}"#: 2,
+    });
+}
+
+async_test_versions! { produce_agg_init_req_skip_vdaf_prep_error }
+
+async fn handle_agg_init_req_hpke_decrypt_err(version: DapVersion) {
     let mut t = Test::new(TEST_VDAF, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
@@ -251,11 +280,15 @@ async fn agg_resp_fail_hpke_decrypt_err(version: DapVersion) {
         agg_resp.transitions[0].var,
         TransitionVar::Failed(TransitionFailure::HpkeDecryptError)
     );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_report_counter{status="rejected_hpke_decrypt_error"}"#: 1,
+    });
 }
 
-async_test_versions! { agg_resp_fail_hpke_decrypt_err }
+async_test_versions! { handle_agg_init_req_hpke_decrypt_err }
 
-async fn agg_resp_fail_hpke_decrypt_err_wrong_id(version: DapVersion) {
+async fn handle_agg_init_req_hpke_unknown_config_id(version: DapVersion) {
     let mut t = Test::new(TEST_VDAF, version);
     let mut reports = t.produce_reports(vec![DapMeasurement::U64(1)]);
 
@@ -273,9 +306,58 @@ async fn agg_resp_fail_hpke_decrypt_err_wrong_id(version: DapVersion) {
         agg_resp.transitions[0].var,
         TransitionVar::Failed(TransitionFailure::HpkeUnknownConfigId)
     );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_report_counter{status="rejected_hpke_unknown_config_id"}"#: 1,
+    });
 }
 
-async_test_versions! { agg_resp_fail_hpke_decrypt_err_wrong_id }
+async_test_versions! { handle_agg_init_req_hpke_unknown_config_id }
+
+async fn handle_agg_init_req_vdaf_prep_error(version: DapVersion) {
+    let mut t = Test::new(TEST_VDAF, version);
+    let report0 =
+        t.produce_invalid_report_public_share_decode_failure(DapMeasurement::U64(1), version);
+    let report1 =
+        t.produce_invalid_report_input_share_decode_failure(DapMeasurement::U64(1), version);
+
+    let agg_req = AggregateInitializeReq {
+        task_id: t.task_id.clone(),
+        agg_job_id: t.agg_job_id.clone(),
+        agg_param: Vec::new(),
+        part_batch_sel: PartialBatchSelector::TimeInterval,
+        report_shares: vec![
+            ReportShare {
+                metadata: report0.metadata,
+                public_share: report0.public_share,
+                encrypted_input_share: report0.encrypted_input_shares[1].clone(),
+            },
+            ReportShare {
+                metadata: report1.metadata,
+                public_share: report1.public_share,
+                encrypted_input_share: report1.encrypted_input_shares[1].clone(),
+            },
+        ],
+    };
+
+    let (_, agg_resp) = t.handle_agg_init_req(agg_req).await.unwrap_continue();
+
+    assert_eq!(agg_resp.transitions.len(), 2);
+    assert_matches!(
+        agg_resp.transitions[0].var,
+        TransitionVar::Failed(TransitionFailure::VdafPrepError)
+    );
+    assert_matches!(
+        agg_resp.transitions[1].var,
+        TransitionVar::Failed(TransitionFailure::VdafPrepError)
+    );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_helper_report_counter{status="rejected_vdaf_prep_error"}"#: 2,
+    });
+}
+
+async_test_versions! { handle_agg_init_req_vdaf_prep_error }
 
 async fn agg_resp_abort_transition_out_of_order(version: DapVersion) {
     let mut t = Test::new(TEST_VDAF, version);
@@ -414,23 +496,21 @@ async_test_versions! { agg_cont_req }
 
 async fn agg_cont_req_skip_vdaf_prep_error(version: DapVersion) {
     let mut t = Test::new(TEST_VDAF, version);
-    let reports = t.produce_reports(vec![
-        DapMeasurement::U64(1),
-        DapMeasurement::U64(1),
-        DapMeasurement::U64(1),
-    ]);
+    let mut reports = t.produce_reports(vec![DapMeasurement::U64(1), DapMeasurement::U64(1)]);
+    reports.insert(
+        1,
+        t.produce_invalid_report_vdaf_prep_failure(DapMeasurement::U64(1), version),
+    );
+
     let (leader_state, agg_init_req) = t.produce_agg_init_req(reports).await.unwrap_continue();
     let (helper_state, agg_resp) = t
         .handle_agg_init_req(agg_init_req.clone())
         .await
         .unwrap_continue();
 
-    let (_, mut agg_cont_req) = t
+    let (_, agg_cont_req) = t
         .handle_agg_resp(leader_state, agg_resp)
         .unwrap_uncommitted();
-    // Simulate VDAF preparation error (due to the report being invalid, say). The leader will
-    // skip the report without processing it any further.
-    agg_cont_req.transitions.remove(1);
 
     let (helper_output_shares, agg_resp) = t
         .handle_agg_cont_req(helper_state, &agg_cont_req)
@@ -446,6 +526,10 @@ async fn agg_cont_req_skip_vdaf_prep_error(version: DapVersion) {
         agg_resp.transitions[1].report_id,
         agg_init_req.report_shares[2].metadata.id
     );
+
+    assert_metrics_include!(t.prometheus_registry, {
+        r#"test_leader_report_counter{status="rejected_vdaf_prep_error"}"#: 1,
+    });
 }
 
 async_test_versions! { agg_cont_req_skip_vdaf_prep_error }
@@ -586,6 +670,9 @@ pub(crate) struct Test {
     helper_hpke_receiver_config: HpkeReceiverConfig,
     client_hpke_config_list: Vec<HpkeConfig>,
     collector_hpke_receiver_config: HpkeReceiverConfig,
+    prometheus_registry: prometheus::Registry,
+    leader_metrics: DaphneMetrics,
+    helper_metrics: DaphneMetrics,
 }
 
 impl Test {
@@ -607,6 +694,9 @@ impl Test {
         let leader_hpke_config = leader_hpke_receiver_config.clone().config;
         let helper_hpke_config = helper_hpke_receiver_config.clone().config;
         let collector_hpke_config = collector_hpke_receiver_config.clone().config;
+        let prometheus_registry = prometheus::Registry::new();
+        let leader_metrics = DaphneMetrics::register(&prometheus_registry, "test_leader").unwrap();
+        let helper_metrics = DaphneMetrics::register(&prometheus_registry, "test_helper").unwrap();
 
         Test {
             now,
@@ -628,6 +718,9 @@ impl Test {
                 vdaf_verify_key,
                 collector_hpke_config,
             },
+            prometheus_registry,
+            leader_metrics,
+            helper_metrics,
         }
     }
 
@@ -651,6 +744,85 @@ impl Test {
         reports
     }
 
+    // Tweak the Helper's share so that decoding succeeds but preparation fails.
+    fn produce_invalid_report_vdaf_prep_failure(
+        &self,
+        measurement: DapMeasurement,
+        version: DapVersion,
+    ) -> Report {
+        let (invalid_public_share, mut invalid_input_shares) = self
+            .task_config
+            .vdaf
+            .produce_input_shares(measurement)
+            .unwrap();
+        invalid_input_shares[1][0] ^= 1; // The first bit is incorrect!
+        self.task_config
+            .vdaf
+            .produce_report_with_extensions_for_shares(
+                invalid_public_share,
+                invalid_input_shares,
+                &self.client_hpke_config_list,
+                self.now,
+                &self.task_id,
+                Vec::new(), // extensions
+                version,
+            )
+            .unwrap()
+    }
+
+    // Tweak the public share so that it can't be decoded.
+    fn produce_invalid_report_public_share_decode_failure(
+        &self,
+        measurement: DapMeasurement,
+        version: DapVersion,
+    ) -> Report {
+        let (mut invalid_public_share, invalid_input_shares) = self
+            .task_config
+            .vdaf
+            .produce_input_shares(measurement)
+            .unwrap();
+        invalid_public_share.push(1); // Add spurious byte at the end
+        self.task_config
+            .vdaf
+            .produce_report_with_extensions_for_shares(
+                invalid_public_share,
+                invalid_input_shares,
+                &self.client_hpke_config_list,
+                self.now,
+                &self.task_id,
+                Vec::new(), // extensions
+                version,
+            )
+            .unwrap()
+    }
+
+    // Tweak the input shares so that they can't be decoded.
+    fn produce_invalid_report_input_share_decode_failure(
+        &self,
+        measurement: DapMeasurement,
+        version: DapVersion,
+    ) -> Report {
+        let (invalid_public_share, mut invalid_input_shares) = self
+            .task_config
+            .vdaf
+            .produce_input_shares(measurement)
+            .unwrap();
+        invalid_input_shares[0].push(1); // Add a spurious byte to the Leader's share
+        invalid_input_shares[1].push(1); // Add a spurious byte to the Helper's share
+        self.task_config
+            .vdaf
+            .produce_report_with_extensions_for_shares(
+                invalid_public_share,
+                invalid_input_shares,
+                &self.client_hpke_config_list,
+                self.now,
+                &self.task_id,
+                Vec::new(), // extensions
+                version,
+            )
+            .unwrap()
+    }
+
     async fn produce_agg_init_req(
         &self,
         reports: Vec<Report>,
@@ -664,6 +836,7 @@ impl Test {
                 &self.agg_job_id,
                 &PartialBatchSelector::TimeInterval,
                 reports,
+                &self.leader_metrics,
             )
             .await
             .unwrap()
@@ -680,6 +853,7 @@ impl Test {
                 &self.helper_hpke_receiver_config,
                 &self.task_config,
                 &agg_init_req,
+                &self.helper_metrics,
             )
             .await
             .unwrap();
@@ -694,7 +868,13 @@ impl Test {
     ) -> DapLeaderTransition<AggregateContinueReq> {
         self.task_config
             .vdaf
-            .handle_agg_resp(&self.task_id, &self.agg_job_id, leader_state, agg_resp)
+            .handle_agg_resp(
+                &self.task_id,
+                &self.agg_job_id,
+                leader_state,
+                agg_resp,
+                &self.leader_metrics,
+            )
             .unwrap()
     }
 
@@ -705,7 +885,13 @@ impl Test {
     ) -> DapAbort {
         self.task_config
             .vdaf
-            .handle_agg_resp(&self.task_id, &self.agg_job_id, leader_state, agg_resp)
+            .handle_agg_resp(
+                &self.task_id,
+                &self.agg_job_id,
+                leader_state,
+                agg_resp,
+                &self.leader_metrics,
+            )
             .err()
             .expect("handle_agg_resp() succeeded; expected failure")
     }
@@ -717,7 +903,7 @@ impl Test {
     ) -> DapHelperTransition<AggregateResp> {
         self.task_config
             .vdaf
-            .handle_agg_cont_req(helper_state, agg_cont_req)
+            .handle_agg_cont_req(helper_state, agg_cont_req, &self.helper_metrics)
             .unwrap()
     }
 
@@ -728,7 +914,7 @@ impl Test {
     ) -> DapAbort {
         self.task_config
             .vdaf
-            .handle_agg_cont_req(helper_state, agg_cont_req)
+            .handle_agg_cont_req(helper_state, agg_cont_req, &self.helper_metrics)
             .err()
             .expect("handle_agg_cont_req() succeeded; expected failure")
     }
@@ -740,7 +926,7 @@ impl Test {
     ) -> Vec<DapOutputShare> {
         self.task_config
             .vdaf
-            .handle_final_agg_resp(leader_uncommitted, agg_resp)
+            .handle_final_agg_resp(leader_uncommitted, agg_resp, &self.leader_metrics)
             .unwrap()
     }