Repository: cloudflare/daphne
Pull request #208 (2023-01-23 20:47:10)

Add spans around major async calls from the router.
Add spans around major async calls (i.e. API entry points) from the URL router.  The span encompasses all work done by the API.

Also adjust the level of a few noisy messages.




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..7106d9e 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -38,7 +38,7 @@ use std::{
     sync::{Arc, RwLock, RwLockReadGuard},
     time::Duration,
 };
-use tracing::{debug, info};
+use tracing::trace;
 use worker::{kv::KvStore, *};
 
 pub(crate) const KV_KEY_PREFIX_HPKE_RECEIVER_CONFIG: &str = "hpke_receiver_config";
@@ -173,7 +173,7 @@ impl DaphneWorkerConfig {
             DaphneWorkerDeployment::default()
         };
         if !matches!(deployment, DaphneWorkerDeployment::Prod) {
-            info!("DAP deployment override applied: {:?}", deployment);
+            trace!("DAP deployment override applied: {:?}", deployment);
         }
 
         let taskprov = if global.allow_taskprov {
@@ -208,7 +208,7 @@ impl DaphneWorkerConfig {
         let admin_token = match env.secret("DAP_ADMIN_BEARER_TOKEN") {
             Ok(raw) => Some(BearerToken::from(raw.to_string())),
             Err(err) => {
-                info!("DAP_ADMIN_BEARER_TOKEN not configured: {:?}", err);
+                trace!("DAP_ADMIN_BEARER_TOKEN not configured: {:?}", err);
                 None
             }
         };
@@ -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 {
@@ -539,7 +539,7 @@ impl<'srv> DaphneWorker<'srv> {
                 .delete(kv_key.name.as_str())
                 .await
                 .map_err(|e| DapError::Fatal(format!("kv_store: {}", e)))?;
-            debug!("deleted KV item {}", kv_key.name);
+            trace!("deleted KV item {}", kv_key.name);
         }
 
         future_delete_durable.await.map_err(dap_err)?;
diff --git a/daphne_worker/src/durable/garbage_collector.rs b/daphne_worker/src/durable/garbage_collector.rs
index 4a81d7f..2c7071f 100644
--- a/daphne_worker/src/durable/garbage_collector.rs
+++ b/daphne_worker/src/durable/garbage_collector.rs
@@ -6,7 +6,7 @@ use crate::{
     durable::{DurableConnector, DurableOrdered, DurableReference},
     initialize_tracing, int_err,
 };
-use tracing::debug;
+use tracing::trace;
 use worker::*;
 
 pub(crate) const DURABLE_GARBAGE_COLLECTOR_PUT: &str = "/internal/do/garbage_collector/put";
@@ -50,7 +50,7 @@ impl DurableObject for GarbageCollector {
 
                 let queued = DurableOrdered::new_roughly_ordered(durable_ref, "object");
                 queued.put(&self.state).await?;
-                debug!(
+                trace!(
                     "GarbageCollector: scheduled {} instance {} for deletion",
                     queued.as_ref().binding,
                     queued.as_ref().id_hex
@@ -81,9 +81,10 @@ impl DurableObject for GarbageCollector {
                             &(),
                         )
                         .await?;
-                    debug!(
+                    trace!(
                         "GarbageCollector: deleted {} instance {}",
-                        durable_ref.binding, durable_ref.id_hex
+                        durable_ref.binding,
+                        durable_ref.id_hex
                     );
                 }
 
diff --git a/daphne_worker/src/durable/helper_state_store.rs b/daphne_worker/src/durable/helper_state_store.rs
index dba91a7..d267747 100644
--- a/daphne_worker/src/durable/helper_state_store.rs
+++ b/daphne_worker/src/durable/helper_state_store.rs
@@ -3,7 +3,7 @@
 
 use crate::{config::DaphneWorkerConfig, durable::state_get, initialize_tracing, int_err};
 use daphne::{messages::Id, DapVersion};
-use tracing::debug;
+use tracing::trace;
 use worker::*;
 
 pub(crate) fn durable_helper_state_name(
@@ -102,7 +102,7 @@ impl DurableObject for HelperStateStore {
     async fn alarm(&mut self) -> Result<Response> {
         self.state.storage().delete_all().await?;
         self.alarmed = false;
-        debug!(
+        trace!(
             "HelperStateStore: deleted instance {}",
             self.state.id().to_string()
         );
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index b132125..1df610e 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -163,7 +163,7 @@ use daphne::{
 use prio::codec::Encode;
 use serde::{Deserialize, Serialize};
 use std::{io, str, sync::Once};
-use tracing::{debug, error, info};
+use tracing::{debug, error, info, trace_span, Instrument};
 use tracing_subscriber::{fmt, layer::*, prelude::*, registry, EnvFilter};
 use worker::*;
 
@@ -236,7 +236,11 @@ impl DaphneWorkerRouter {
             .get_async("/:version/hpke_config", |req, ctx| async move {
                 let daph = ctx.data.handler(&ctx.env);
                 let req = daph.worker_request_to_dap(req).await?;
-                match daph.http_get_hpke_config(&req).await {
+                match daph
+                    .http_get_hpke_config(&req)
+                    .instrument(trace_span!("hpke_config"))
+                    .await
+                {
                     Ok(req) => dap_response_to_worker(req),
                     Err(e) => abort(e),
                 }
@@ -258,6 +262,7 @@ impl DaphneWorkerRouter {
 
                 let cmd: InternalTestAddTask = req.json().await?;
                 daph.internal_add_task(daph.state.config.default_version, cmd)
+                    .instrument(trace_span!("task"))
                     .await?;
                 Response::empty()
             });
@@ -269,7 +274,11 @@ impl DaphneWorkerRouter {
                         let daph = ctx.data.handler(&ctx.env);
                         let req = daph.worker_request_to_dap(req).await?;
 
-                        match daph.http_post_upload(&req).await {
+                        match daph
+                            .http_post_upload(&req)
+                            .instrument(trace_span!("upload"))
+                            .await
+                        {
                             Ok(()) => Response::empty(),
                             Err(e) => abort(e),
                         }
@@ -278,7 +287,11 @@ impl DaphneWorkerRouter {
                         let daph = ctx.data.handler(&ctx.env);
                         let req = daph.worker_request_to_dap(req).await?;
 
-                        match daph.http_post_collect(&req).await {
+                        match daph
+                            .http_post_collect(&req)
+                            .instrument(trace_span!("collect"))
+                            .await
+                        {
                             Ok(collect_uri) => {
                                 let mut headers = Headers::new();
                                 headers.set("Location", collect_uri.as_str())?;
@@ -296,7 +309,11 @@ impl DaphneWorkerRouter {
                             let task_id = parse_id!(ctx.param("task_id"));
                             let collect_id = parse_id!(ctx.param("collect_id"));
                             let daph = ctx.data.handler(&ctx.env);
-                            match daph.poll_collect_job(&task_id, &collect_id).await {
+                            match daph
+                                .poll_collect_job(&task_id, &collect_id)
+                                .instrument(trace_span!("poll_collect_job"))
+                                .await
+                            {
                                 Ok(DapCollectJob::Done(collect_resp)) => {
                                     dap_response_to_worker(DapResponse {
                                         media_type: Some(constants::MEDIA_TYPE_COLLECT_RESP),
@@ -318,7 +335,11 @@ impl DaphneWorkerRouter {
                         // TODO(cjpatton) Only enable this if `self.enable_internal_test` is set.
                         let daph = ctx.data.handler(&ctx.env);
                         let report_sel: DaphneWorkerReportSelector = req.json().await?;
-                        match daph.process(&report_sel).await {
+                        match daph
+                            .process(&report_sel)
+                            .instrument(trace_span!("process"))
+                            .await
+                        {
                             Ok(telem) => {
                                 debug!("{:?}", telem);
                                 Response::from_json(&telem)
@@ -335,7 +356,11 @@ impl DaphneWorkerRouter {
                             // TODO(cjpatton) Only enable this if `self.enable_internal_test` is set.
                             let task_id = parse_id!(ctx.param("task_id"));
                             let daph = ctx.data.handler(&ctx.env);
-                            match daph.internal_current_batch(&task_id).await {
+                            match daph
+                                .internal_current_batch(&task_id)
+                                .instrument(trace_span!("current_batch"))
+                                .await
+                            {
                                 Ok(batch_id) => Response::from_bytes(
                                     batch_id.to_base64url().as_bytes().to_owned(),
                                 ),
@@ -350,7 +375,11 @@ impl DaphneWorkerRouter {
                     let daph = ctx.data.handler(&ctx.env);
                     let req = daph.worker_request_to_dap(req).await?;
 
-                    match daph.http_post_aggregate(&req).await {
+                    match daph
+                        .http_post_aggregate(&req)
+                        .instrument(trace_span!("aggregate"))
+                        .await
+                    {
                         Ok(resp) => dap_response_to_worker(resp),
                         Err(e) => abort(e),
                     }
@@ -359,7 +388,11 @@ impl DaphneWorkerRouter {
                     let daph = ctx.data.handler(&ctx.env);
                     let req = daph.worker_request_to_dap(req).await?;
 
-                    match daph.http_post_aggregate_share(&req).await {
+                    match daph
+                        .http_post_aggregate_share(&req)
+                        .instrument(trace_span!("aggregate_share"))
+                        .await
+                    {
                         Ok(resp) => dap_response_to_worker(resp),
                         Err(e) => abort(e),
                     }
@@ -372,7 +405,11 @@ impl DaphneWorkerRouter {
             router
                 .post_async("/internal/delete_all", |_req, ctx| async move {
                     let daph = ctx.data.handler(&ctx.env);
-                    match daph.internal_delete_all().await {
+                    match daph
+                        .internal_delete_all()
+                        .instrument(trace_span!("delete_all"))
+                        .await
+                    {
                         Ok(()) => Response::empty(),
                         Err(e) => abort(e.into()),
                     }
@@ -387,6 +424,7 @@ impl DaphneWorkerRouter {
                         let daph = ctx.data.handler(&ctx.env);
                         let cmd: InternalTestEndpointForTask = req.json().await?;
                         daph.internal_endpoint_for_task(daph.state.config.default_version, cmd)
+                            .instrument(trace_span!("endpoint_for_task"))
                             .await
                     },
                 )
@@ -396,13 +434,16 @@ impl DaphneWorkerRouter {
                         let daph = ctx.data.handler(&ctx.env);
                         let cmd: InternalTestEndpointForTask = req.json().await?;
                         let version = daph.extract_version_parameter(&req)?;
-                        daph.internal_endpoint_for_task(version, cmd).await
+                        daph.internal_endpoint_for_task(version, cmd)
+                            .instrument(trace_span!("endpoint_for_task"))
+                            .await
                     },
                 )
                 .post_async("/internal/test/add_task", |mut req, ctx| async move {
                     let daph = ctx.data.handler(&ctx.env);
                     let cmd: InternalTestAddTask = req.json().await?;
                     daph.internal_add_task(daph.state.config.default_version, cmd)
+                        .instrument(trace_span!("add_task"))
                         .await?;
                     Response::from_json(&serde_json::json!({
                         "status": "success",
@@ -414,7 +455,9 @@ impl DaphneWorkerRouter {
                         let daph = ctx.data.handler(&ctx.env);
                         let cmd: InternalTestAddTask = req.json().await?;
                         let version = daph.extract_version_parameter(&req)?;
-                        daph.internal_add_task(version, cmd).await?;
+                        daph.internal_add_task(version, cmd)
+                            .instrument(trace_span!("add_task"))
+                            .await?;
                         Response::from_json(&serde_json::json!({
                             "status": "success",
                         }))


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()
     }