Repository: cloudflare/daphne
Pull request #362 (2023-07-07 00:21:59)

Prepare for offloading VDAF preparation
We want to be able to offload VDAF preparation to durable objects. This PR refactors the code to make this easier and resolves some prerequisites. Details in the commit messages, but at a high level:
* Split `consume_report_share()` into two steps and wrap the state after each step with a state object.
* Implement codec traits on VdafPrepState and VdafPrepMessage, which will be required for serializing the input and output to DO.
* Fix a bug that triggers abort if a reports plaintext input share is malformed. (We should just reject instead.)


Patch:
diff --git a/Cargo.lock b/Cargo.lock
index 6d95f1b..4cdb63f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -148,6 +148,12 @@ version = "1.0.71"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9c7d0618f0e0b7e8ff11427422b64564d5fb0be1940354bfe2e0529b18a9d9b8"
 
+[[package]]
+name = "arc-swap"
+version = "1.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bddcadddf5e9015d310179a59bb28c4d4b9920ad0f11e8e14dbadf654890c9a6"
+
 [[package]]
 name = "assert_matches"
 version = "1.5.0"
@@ -677,6 +683,7 @@ dependencies = [
 name = "daphne_worker"
 version = "0.3.0"
 dependencies = [
+ "arc-swap",
  "async-trait",
  "chrono",
  "daphne",
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 60dfd1f..b6693f0 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -46,11 +46,7 @@ use crate::{
         ReportMetadata, TaskId, Time,
     },
     taskprov::TaskprovVersion,
-    vdaf::{
-        prio2::prio2_decode_prepare_state,
-        prio3::{prio3_append_prepare_state, prio3_decode_prepare_state},
-        VdafAggregateShare, VdafMessage, VdafState, VdafVerifyKey,
-    },
+    vdaf::{VdafAggregateShare, VdafPrepMessage, VdafPrepState, VdafVerifyKey},
 };
 use constants::DapMediaType;
 #[cfg(test)]
@@ -58,7 +54,7 @@ use criterion as _;
 pub use error::DapError;
 use hpke::{HpkeConfig, HpkeKemId};
 use prio::{
-    codec::{Decode, Encode},
+    codec::{Decode, Encode, ParameterizedDecode},
     vdaf::Aggregatable as AggregatableTrait,
 };
 use rand::prelude::*;
@@ -450,7 +446,7 @@ pub enum DapAggregateResult {
 /// The Leader's state after sending an AggregateInitReq.
 #[derive(Debug)]
 pub struct DapLeaderState {
-    pub(crate) seq: Vec<(VdafState, VdafMessage, Time, ReportId)>,
+    pub(crate) seq: Vec<(VdafPrepState, VdafPrepMessage, Time, ReportId)>,
 }
 
 /// The Leader's state after sending an AggregateContReq.
@@ -463,36 +459,21 @@ pub struct DapLeaderUncommitted {
 #[derive(Clone, Debug, PartialEq)]
 pub struct DapHelperState {
     pub(crate) part_batch_sel: PartialBatchSelector,
-    pub(crate) seq: Vec<(VdafState, Time, ReportId)>,
+    pub(crate) seq: Vec<(VdafPrepState, Time, ReportId)>,
 }
 
-impl DapHelperState {
-    /// Encode the Helper state as a byte string.
-    ///
-    /// This method is used by the Helper in order to offload its state to the Leader. For
-    /// example, it might encrypt the output and add the ciphertext to an outgoing aggregate
-    /// response.
-    ///
-    /// Note that the encoding format is not specified by the DAP standard.
-    pub fn get_encoded(&self, vdaf_config: &VdafConfig) -> Result<Vec<u8>, DapError> {
-        let mut bytes = vec![];
-        self.part_batch_sel.encode(&mut bytes);
+impl Encode for DapHelperState {
+    fn encode(&self, bytes: &mut Vec<u8>) {
+        self.part_batch_sel.encode(bytes);
         for (state, time, report_id) in self.seq.iter() {
-            match (vdaf_config, state) {
-                (VdafConfig::Prio3(prio3_config), _) => {
-                    prio3_append_prepare_state(&mut bytes, prio3_config, state)?;
-                }
-                (VdafConfig::Prio2 { .. }, VdafState::Prio2(state)) => {
-                    state.encode(&mut bytes);
-                }
-                _ => return Err(fatal_error!(err = "VDAF config and prep state mismatch")),
-            }
-            time.encode(&mut bytes);
-            report_id.encode(&mut bytes);
+            state.encode(bytes);
+            time.encode(bytes);
+            report_id.encode(bytes);
         }
-        Ok(bytes)
     }
+}
 
+impl DapHelperState {
     /// Decode the Helper state from a byte string.
     pub fn get_decoded(vdaf_config: &VdafConfig, data: &[u8]) -> Result<Self, DapError> {
         let mut r = std::io::Cursor::new(data);
@@ -500,14 +481,8 @@ impl DapHelperState {
             .map_err(|e| DapAbort::from_codec_error(e, None))?;
         let mut seq = vec![];
         while (r.position() as usize) < data.len() {
-            let state = match vdaf_config {
-                VdafConfig::Prio3(ref prio3_config) => {
-                    prio3_decode_prepare_state(prio3_config, 1, &mut r)?
-                }
-                VdafConfig::Prio2 { dimension } => {
-                    prio2_decode_prepare_state(*dimension, 1, &mut r)?
-                }
-            };
+            let state = VdafPrepState::decode_with_param(vdaf_config, &mut r)
+                .map_err(|e| DapAbort::from_codec_error(e, None))?;
             let time = Time::decode(&mut r).map_err(|e| DapAbort::from_codec_error(e, None))?;
             let report_id =
                 ReportId::decode(&mut r).map_err(|e| DapAbort::from_codec_error(e, None))?;
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index b214b44..0e4c5d4 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -17,12 +17,12 @@ use crate::{
     metrics::ContextualizedDaphneMetrics,
     vdaf::{
         prio2::{
-            prio2_encode_prepare_message, prio2_helper_prepare_finish, prio2_leader_prepare_finish,
-            prio2_prepare_init, prio2_shard, prio2_unshard,
+            prio2_helper_prepare_finish, prio2_leader_prepare_finish, prio2_prepare_init,
+            prio2_shard, prio2_unshard,
         },
         prio3::{
-            prio3_encode_prepare_message, prio3_helper_prepare_finish, prio3_leader_prepare_finish,
-            prio3_prepare_init, prio3_shard, prio3_unshard,
+            prio3_helper_prepare_finish, prio3_leader_prepare_finish, prio3_prepare_init,
+            prio3_shard, prio3_unshard,
         },
     },
     DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
@@ -30,7 +30,7 @@ use crate::{
     DapTaskConfig, DapVersion, MetaAggregationJobId, VdafConfig,
 };
 use prio::{
-    codec::{CodecError, Decode, Encode, ParameterizedEncode},
+    codec::{CodecError, Decode, Encode, ParameterizedDecode, ParameterizedEncode},
     field::{Field128, Field64, FieldPrio2},
     vdaf::{
         prio2::{Prio2PrepareShare, Prio2PrepareState},
@@ -39,7 +39,9 @@ use prio::{
 };
 use rand::prelude::*;
 use serde::{Deserialize, Serialize};
-use std::{collections::HashSet, convert::TryInto};
+use std::{borrow::Cow, collections::HashSet, convert::TryInto};
+
+use self::{prio2::prio2_decode_prepare_state, prio3::prio3_decode_prepare_state};
 
 const CTX_INPUT_SHARE_DRAFT02: &[u8] = b"dap-02 input share";
 const CTX_INPUT_SHARE_DRAFT04: &[u8] = b"dap-04 input share";
@@ -75,20 +77,196 @@ impl AsRef<[u8]> for VdafVerifyKey {
     }
 }
 
+/// Report state during aggregation initialization after consuming the report share. This involves
+/// decryption as well a few validation steps.
 #[derive(Clone, Debug, PartialEq)]
-pub(crate) enum VdafState {
+pub enum EarlyReportStateConsumed<'req> {
+    Ready {
+        metadata: Cow<'req, ReportMetadata>,
+        public_share: Cow<'req, [u8]>,
+        input_share: Vec<u8>,
+    },
+    Rejected(TransitionFailure),
+}
+
+impl<'req> EarlyReportStateConsumed<'req> {
+    #[allow(clippy::too_many_arguments)]
+    pub(crate) async fn consume(
+        decrypter: &impl HpkeDecrypter,
+        is_leader: bool,
+        task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        metadata: &'req ReportMetadata,
+        public_share: &'req [u8],
+        encrypted_input_share: &HpkeCiphertext,
+    ) -> Result<EarlyReportStateConsumed<'req>, DapError> {
+        if metadata.time >= task_config.expiration {
+            return Ok(Self::Rejected(TransitionFailure::TaskExpired));
+        }
+
+        let input_share_text = match task_config.version {
+            DapVersion::Draft02 => CTX_INPUT_SHARE_DRAFT02,
+            DapVersion::Draft04 => CTX_INPUT_SHARE_DRAFT04,
+            _ => return Err(unimplemented_version()),
+        };
+        let n: usize = input_share_text.len();
+        let mut info = Vec::new();
+        info.reserve(n + 2);
+        info.extend_from_slice(input_share_text);
+        info.push(CTX_ROLE_CLIENT); // Sender role (receiver role set below)
+        info.push(if is_leader {
+            CTX_ROLE_LEADER
+        } else {
+            CTX_ROLE_HELPER
+        }); // Receiver role
+
+        let mut aad = Vec::with_capacity(58);
+        task_id.encode(&mut aad);
+        metadata.encode_with_param(&task_config.version, &mut aad);
+        // TODO spec: Consider folding the public share into a field called "header".
+        encode_u32_bytes(&mut aad, public_share);
+
+        let encoded_input_share = match decrypter
+            .hpke_decrypt(task_id, &info, &aad, encrypted_input_share)
+            .await
+        {
+            Ok(encoded_input_share) => encoded_input_share,
+            Err(DapError::Transition(failure)) => return Ok(Self::Rejected(failure)),
+            Err(e) => return Err(e),
+        };
+
+        // For Draft02, the encoded input share is the VDAF-specific payload, but for Draft03 and
+        // later it is a serialized PlaintextInputShare.  For simplicity in later code, we wrap the Draft02
+        // payload into a PlaintextInputShare.
+        let input_share = match task_config.version {
+            DapVersion::Draft02 => PlaintextInputShare {
+                extensions: vec![],
+                payload: encoded_input_share,
+            },
+            DapVersion::Draft04 => match PlaintextInputShare::get_decoded(&encoded_input_share) {
+                Ok(input_share) => input_share,
+                Err(..) => return Ok(Self::Rejected(TransitionFailure::UnrecognizedMessage)),
+            },
+            _ => return Err(unimplemented_version()),
+        };
+
+        Ok(Self::Ready {
+            metadata: Cow::Borrowed(metadata),
+            public_share: Cow::Borrowed(public_share),
+            input_share: input_share.payload,
+        })
+    }
+}
+
+/// Report state during aggregation initialization after the VDAF preparation step.
+pub enum EarlyReportStateInitialized {
+    Ready {
+        state: VdafPrepState,
+        message: VdafPrepMessage,
+    },
+    Rejected(TransitionFailure),
+}
+
+impl EarlyReportStateInitialized {
+    pub(crate) fn initialize(
+        is_leader: bool,
+        task_config: &DapTaskConfig,
+        early_report_state_consumed: EarlyReportStateConsumed,
+    ) -> Result<Self, DapError> {
+        let (metadata, public_share, input_share) = match early_report_state_consumed {
+            EarlyReportStateConsumed::Ready {
+                metadata,
+                public_share,
+                input_share,
+            } => (metadata, public_share, input_share),
+            EarlyReportStateConsumed::Rejected(failure) => return Ok(Self::Rejected(failure)),
+        };
+
+        let agg_id = usize::from(!is_leader);
+        let res = match (&task_config.vdaf, &task_config.vdaf_verify_key) {
+            (VdafConfig::Prio3(ref prio3_config), VdafVerifyKey::Prio3(ref verify_key)) => {
+                prio3_prepare_init(
+                    prio3_config,
+                    verify_key,
+                    agg_id,
+                    &metadata.as_ref().id.0,
+                    public_share.as_ref(),
+                    input_share.as_ref(),
+                )
+            }
+            (VdafConfig::Prio2 { dimension }, VdafVerifyKey::Prio2(ref verify_key)) => {
+                prio2_prepare_init(
+                    *dimension,
+                    verify_key,
+                    agg_id,
+                    &metadata.as_ref().id.0,
+                    public_share.as_ref(),
+                    input_share.as_ref(),
+                )
+            }
+            _ => return Err(fatal_error!(err = "VDAF verify key does not match config")),
+        };
+
+        let early_report_state_initialized = match res {
+            Ok((state, message)) => Self::Ready { state, message },
+            Err(..) => Self::Rejected(TransitionFailure::VdafPrepError),
+        };
+        Ok(early_report_state_initialized)
+    }
+}
+
+/// VDAF preparation state.
+#[derive(Clone, Debug, PartialEq)]
+pub enum VdafPrepState {
     Prio2(Prio2PrepareState),
     Prio3Field64(Prio3PrepareState<Field64, 16>),
     Prio3Field128(Prio3PrepareState<Field128, 16>),
 }
 
+impl Encode for VdafPrepState {
+    fn encode(&self, bytes: &mut Vec<u8>) {
+        match self {
+            Self::Prio3Field64(state) => state.encode(bytes),
+            Self::Prio3Field128(state) => state.encode(bytes),
+            Self::Prio2(state) => state.encode(bytes),
+        }
+    }
+}
+
+impl ParameterizedDecode<VdafConfig> for VdafPrepState {
+    fn decode_with_param(
+        vdaf_config: &VdafConfig,
+        bytes: &mut std::io::Cursor<&[u8]>,
+    ) -> Result<Self, CodecError> {
+        match vdaf_config {
+            VdafConfig::Prio3(ref prio3_config) => {
+                Ok(prio3_decode_prepare_state(prio3_config, 1, bytes)
+                    .map_err(|e| CodecError::Other(Box::new(e)))?)
+            }
+            VdafConfig::Prio2 { dimension } => Ok(prio2_decode_prepare_state(*dimension, 1, bytes)
+                .map_err(|e| CodecError::Other(Box::new(e)))?),
+        }
+    }
+}
+
+/// VDAF preparation message.
 #[derive(Clone, Debug)]
-pub(crate) enum VdafMessage {
+pub enum VdafPrepMessage {
     Prio2Share(Prio2PrepareShare),
     Prio3ShareField64(Prio3PrepareShare<Field64, 16>),
     Prio3ShareField128(Prio3PrepareShare<Field128, 16>),
 }
 
+impl Encode for VdafPrepMessage {
+    fn encode(&self, bytes: &mut Vec<u8>) {
+        match self {
+            Self::Prio3ShareField64(share) => share.encode(bytes),
+            Self::Prio3ShareField128(share) => share.encode(bytes),
+            Self::Prio2Share(share) => share.encode(bytes),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub(crate) enum VdafAggregateShare {
@@ -325,100 +503,6 @@ impl VdafConfig {
         )
     }
 
-    /// Consume a report share sent by the Client and return the initial Prepare step. This is run
-    /// by each Aggregator.
-    ///
-    /// # Inputs
-    ///
-    /// * `decryptor` is used to decrypt the input share.
-    ///
-    /// * `verify_key` is the secret VDAF verification key shared by the Aggregators.
-    ///
-    /// * `task_id` is the DAP task ID indicated by the report.
-    ///
-    /// * `report_id` is the report ID.
-    ///
-    /// * `encrypted_input_share` is the encrypted input share.
-    ///
-    /// * `version` is the DapVersion to use.
-    #[allow(clippy::too_many_arguments)]
-    pub(crate) async fn consume_report_share(
-        &self,
-        decrypter: &impl HpkeDecrypter,
-        is_leader: bool,
-        task_id: &TaskId,
-        task_config: &DapTaskConfig,
-        metadata: &ReportMetadata,
-        public_share: &[u8],
-        encrypted_input_share: &HpkeCiphertext,
-    ) -> Result<(VdafState, VdafMessage), DapError> {
-        if metadata.time >= task_config.expiration {
-            return Err(DapError::Transition(TransitionFailure::TaskExpired));
-        }
-
-        let input_share_text = match task_config.version {
-            DapVersion::Draft02 => CTX_INPUT_SHARE_DRAFT02,
-            DapVersion::Draft04 => CTX_INPUT_SHARE_DRAFT04,
-            _ => return Err(unimplemented_version()),
-        };
-        let n: usize = input_share_text.len();
-        let mut info = Vec::new();
-        info.reserve(n + 2);
-        info.extend_from_slice(input_share_text);
-        info.push(CTX_ROLE_CLIENT); // Sender role (receiver role set below)
-        info.push(if is_leader {
-            CTX_ROLE_LEADER
-        } else {
-            CTX_ROLE_HELPER
-        }); // Receiver role
-
-        let mut aad = Vec::with_capacity(58);
-        task_id.encode(&mut aad);
-        metadata.encode_with_param(&task_config.version, &mut aad);
-        // TODO spec: Consider folding the public share into a field called "header".
-        encode_u32_bytes(&mut aad, public_share);
-
-        let encoded_input_share = decrypter
-            .hpke_decrypt(task_id, &info, &aad, encrypted_input_share)
-            .await?;
-        // For Draft02, the encoded input share is the VDAF-specific payload, but for Draft03 and
-        // later it is a serialized PlaintextInputShare.  For simplicity in later code, we wrap the Draft02
-        // payload into a PlaintextInputShare.
-        let input_share = match task_config.version {
-            DapVersion::Draft02 => PlaintextInputShare {
-                extensions: vec![],
-                payload: encoded_input_share,
-            },
-            _ => PlaintextInputShare::get_decoded(&encoded_input_share)
-                .map_err(|e| DapAbort::from_codec_error(e, task_id.clone()))?,
-        };
-
-        let agg_id = usize::from(!is_leader);
-        match (self, &task_config.vdaf_verify_key) {
-            (Self::Prio3(ref prio3_config), VdafVerifyKey::Prio3(ref verify_key)) => {
-                Ok(prio3_prepare_init(
-                    prio3_config,
-                    verify_key,
-                    agg_id,
-                    &metadata.id.0,
-                    public_share,
-                    &input_share.payload,
-                )?)
-            }
-            (Self::Prio2 { dimension }, VdafVerifyKey::Prio2(ref verify_key)) => {
-                Ok(prio2_prepare_init(
-                    *dimension,
-                    verify_key,
-                    agg_id,
-                    &metadata.id.0,
-                    public_share,
-                    &input_share.payload,
-                )?)
-            }
-            _ => Err(fatal_error!(err = "VDAF verify key does not match config")),
-        }
-    }
-
     /// Initialize the aggregation flow for a sequence of reports. The outputs are the Leader's
     /// state for the aggregation flow and the initial aggregate request to be sent to the Helper.
     /// This method is called by the Leader.
@@ -451,21 +535,27 @@ impl VdafConfig {
                 (it.next().unwrap(), it.next().unwrap())
             };
 
-            match self
-                .consume_report_share(
-                    decrypter,
-                    true, // is_leader
-                    task_id,
-                    task_config,
-                    &report.report_metadata,
-                    &report.public_share,
-                    &leader_share,
-                )
-                .await
-            {
-                Ok((step, message)) => {
+            let early_report_state_consumed = EarlyReportStateConsumed::consume(
+                decrypter,
+                true,
+                task_id,
+                task_config,
+                &report.report_metadata,
+                &report.public_share,
+                &leader_share,
+            )
+            .await?;
+
+            let early_report_state_initialized = EarlyReportStateInitialized::initialize(
+                true,
+                task_config,
+                early_report_state_consumed,
+            )?;
+
+            match early_report_state_initialized {
+                EarlyReportStateInitialized::Ready { state, message } => {
                     states.push((
-                        step,
+                        state,
                         message,
                         report.report_metadata.time,
                         report.report_metadata.id.clone(),
@@ -477,13 +567,12 @@ impl VdafConfig {
                     });
                 }
 
-                // Skip report that can't be processed any further.
-                Err(DapError::Transition(failure)) => {
-                    metrics.report_inc_by(&format!("rejected_{failure}"), 1)
+                EarlyReportStateInitialized::Rejected(failure) => {
+                    // Skip report that can't be processed any further.
+                    metrics.report_inc_by(&format!("rejected_{failure}"), 1);
+                    continue;
                 }
-
-                Err(e) => return Err(DapAbort::Internal(Box::new(e))),
-            };
+            }
         }
 
         if seq.is_empty() {
@@ -547,37 +636,37 @@ impl VdafConfig {
             }
             processed.insert(report_share.report_metadata.id.clone());
 
-            let var = match self
-                .consume_report_share(
-                    decrypter,
-                    false, // is_leader
-                    task_id,
-                    task_config,
-                    &report_share.report_metadata,
-                    &report_share.public_share,
-                    &report_share.encrypted_input_share,
-                )
-                .await
-            {
-                Ok((step, message)) => {
-                    let message_data = match self {
-                        Self::Prio3(..) => prio3_encode_prepare_message(&message),
-                        Self::Prio2 { .. } => prio2_encode_prepare_message(&message),
-                    };
+            let early_report_state_consumed = EarlyReportStateConsumed::consume(
+                decrypter,
+                false,
+                task_id,
+                task_config,
+                &report_share.report_metadata,
+                &report_share.public_share,
+                &report_share.encrypted_input_share,
+            )
+            .await?;
+
+            let early_report_state_initialized = EarlyReportStateInitialized::initialize(
+                false,
+                task_config,
+                early_report_state_consumed,
+            )?;
+
+            let var = match early_report_state_initialized {
+                EarlyReportStateInitialized::Ready { state, message } => {
                     states.push((
-                        step,
+                        state,
                         report_share.report_metadata.time,
                         report_share.report_metadata.id.clone(),
                     ));
-                    TransitionVar::Continued(message_data)
+                    TransitionVar::Continued(message.get_encoded())
                 }
 
-                Err(DapError::Transition(failure)) => {
+                EarlyReportStateInitialized::Rejected(failure) => {
                     metrics.report_inc_by(&format!("rejected_{failure}"), 1);
                     TransitionVar::Failed(failure)
                 }
-
-                Err(e) => return Err(DapAbort::Internal(Box::new(e))),
             };
 
             transitions.push(Transition {
diff --git a/daphne/src/vdaf/prio2.rs b/daphne/src/vdaf/prio2.rs
index cc83e37..564d5fb 100644
--- a/daphne/src/vdaf/prio2.rs
+++ b/daphne/src/vdaf/prio2.rs
@@ -5,7 +5,8 @@
 //! [VDAF](https://datatracker.ietf.org/doc/draft-patton-cfrg-vdaf/).
 
 use crate::{
-    vdaf::VdafError, DapAggregateResult, DapMeasurement, VdafAggregateShare, VdafMessage, VdafState,
+    vdaf::VdafError, DapAggregateResult, DapMeasurement, VdafAggregateShare, VdafPrepMessage,
+    VdafPrepState,
 };
 use prio::{
     codec::{Decode, Encode, ParameterizedDecode},
@@ -46,25 +47,28 @@ pub(crate) fn prio2_prepare_init(
     nonce: &[u8; 16],
     public_share_data: &[u8],
     input_share_data: &[u8],
-) -> Result<(VdafState, VdafMessage), VdafError> {
+) -> Result<(VdafPrepState, VdafPrepMessage), VdafError> {
     let vdaf = Prio2::new(dimension)?;
     <()>::get_decoded_with_param(&vdaf, public_share_data)?;
     let input_share: Share<FieldPrio2, 32> =
         Share::get_decoded_with_param(&(&vdaf, agg_id), input_share_data)?;
     let (state, share) = vdaf.prepare_init(verify_key, agg_id, &(), nonce, &(), &input_share)?;
-    Ok((VdafState::Prio2(state), VdafMessage::Prio2Share(share)))
+    Ok((
+        VdafPrepState::Prio2(state),
+        VdafPrepMessage::Prio2Share(share),
+    ))
 }
 
 /// Consume the verifier shares and return the output share and serialized outbound message.
 pub(crate) fn prio2_leader_prepare_finish(
     dimension: usize,
-    leader_state: VdafState,
-    leader_share: VdafMessage,
+    leader_state: VdafPrepState,
+    leader_share: VdafPrepMessage,
     helper_share_data: &[u8],
 ) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
     let vdaf = Prio2::new(dimension)?;
     let (out_share, outbound) = match (leader_state, leader_share) {
-        (VdafState::Prio2(state), VdafMessage::Prio2Share(share)) => {
+        (VdafPrepState::Prio2(state), VdafPrepMessage::Prio2Share(share)) => {
             let helper_share =
                 Prio2PrepareShare::get_decoded_with_param(&state, helper_share_data)?;
             vdaf.prepare_preprocess([share, helper_share])?;
@@ -84,13 +88,13 @@ pub(crate) fn prio2_leader_prepare_finish(
 /// Consume the peer's prepare message and return an output share.
 pub(crate) fn prio2_helper_prepare_finish(
     dimension: usize,
-    helper_state: VdafState,
+    helper_state: VdafPrepState,
     leader_message_data: &[u8],
 ) -> Result<VdafAggregateShare, VdafError> {
     let vdaf = Prio2::new(dimension)?;
     <()>::get_decoded(leader_message_data)?;
     let out_share = match helper_state {
-        VdafState::Prio2(state) => match vdaf.prepare_step(state, ())? {
+        VdafPrepState::Prio2(state) => match vdaf.prepare_step(state, ())? {
             PrepareTransition::Continue(..) => {
                 panic!("prio2_helper_prepare_finish: unexpected transition (continued)")
             }
@@ -107,22 +111,14 @@ pub(crate) fn prio2_decode_prepare_state(
     dimension: usize,
     agg_id: usize,
     bytes: &mut Cursor<&[u8]>,
-) -> Result<VdafState, VdafError> {
+) -> Result<VdafPrepState, VdafError> {
     let vdaf = Prio2::new(dimension)?;
-    Ok(VdafState::Prio2(Prio2PrepareState::decode_with_param(
+    Ok(VdafPrepState::Prio2(Prio2PrepareState::decode_with_param(
         &(&vdaf, agg_id),
         bytes,
     )?))
 }
 
-/// Encode `message` as a byte string.
-pub(crate) fn prio2_encode_prepare_message(message: &VdafMessage) -> Vec<u8> {
-    match message {
-        VdafMessage::Prio2Share(message) => message.get_encoded(),
-        _ => panic!("prio2_encode_prepare_message: unexpected message type"),
-    }
-}
-
 /// Interpret `encoded_agg_shares` as a sequence of encoded aggregate shares and unshard them.
 pub(crate) fn prio2_unshard<M: IntoIterator<Item = Vec<u8>>>(
     dimension: usize,
diff --git a/daphne/src/vdaf/prio3.rs b/daphne/src/vdaf/prio3.rs
index 5f24a38..0702d56 100644
--- a/daphne/src/vdaf/prio3.rs
+++ b/daphne/src/vdaf/prio3.rs
@@ -5,7 +5,7 @@
 
 use crate::{
     vdaf::VdafError, DapAggregateResult, DapMeasurement, Prio3Config, VdafAggregateShare,
-    VdafMessage, VdafState,
+    VdafPrepMessage, VdafPrepState,
 };
 use prio::{
     codec::{Encode, ParameterizedDecode},
@@ -104,7 +104,7 @@ pub(crate) fn prio3_prepare_init(
     nonce: &[u8; 16],
     public_share_data: &[u8],
     input_share_data: &[u8],
-) -> Result<(VdafState, VdafMessage), VdafError> {
+) -> Result<(VdafPrepState, VdafPrepMessage), VdafError> {
     match &config {
         Prio3Config::Count => {
             let vdaf = Prio3::new_count(2)?;
@@ -117,8 +117,8 @@ pub(crate) fn prio3_prepare_init(
                 input_share_data
             );
             Ok((
-                VdafState::Prio3Field64(state),
-                VdafMessage::Prio3ShareField64(share),
+                VdafPrepState::Prio3Field64(state),
+                VdafPrepMessage::Prio3ShareField64(share),
             ))
         }
         Prio3Config::Histogram { buckets } => {
@@ -132,8 +132,8 @@ pub(crate) fn prio3_prepare_init(
                 input_share_data
             );
             Ok((
-                VdafState::Prio3Field128(state),
-                VdafMessage::Prio3ShareField128(share),
+                VdafPrepState::Prio3Field128(state),
+                VdafPrepMessage::Prio3ShareField128(share),
             ))
         }
         Prio3Config::Sum { bits } => {
@@ -147,8 +147,8 @@ pub(crate) fn prio3_prepare_init(
                 input_share_data
             );
             Ok((
-                VdafState::Prio3Field128(state),
-                VdafMessage::Prio3ShareField128(share),
+                VdafPrepState::Prio3Field128(state),
+                VdafPrepMessage::Prio3ShareField128(share),
             ))
         }
         Prio3Config::SumVec { bits, len } => {
@@ -162,8 +162,8 @@ pub(crate) fn prio3_prepare_init(
                 input_share_data
             );
             Ok((
-                VdafState::Prio3Field128(state),
-                VdafMessage::Prio3ShareField128(share),
+                VdafPrepState::Prio3Field128(state),
+                VdafPrepMessage::Prio3ShareField128(share),
             ))
         }
     }
@@ -197,15 +197,15 @@ macro_rules! leader_prep_fin {
 /// Consume the verifier shares and return the output share and serialized outbound message.
 pub(crate) fn prio3_leader_prepare_finish(
     config: &Prio3Config,
-    leader_state: VdafState,
-    leader_share: VdafMessage,
+    leader_state: VdafPrepState,
+    leader_share: VdafPrepMessage,
     helper_share_data: &[u8],
 ) -> Result<(VdafAggregateShare, Vec<u8>), VdafError> {
     let (agg_share, outbound) = match (&config, leader_state, leader_share) {
         (
             Prio3Config::Count,
-            VdafState::Prio3Field64(state),
-            VdafMessage::Prio3ShareField64(share),
+            VdafPrepState::Prio3Field64(state),
+            VdafPrepMessage::Prio3ShareField64(share),
         ) => {
             let vdaf = Prio3::new_count(2)?;
             let (out_share, outbound) = leader_prep_fin!(vdaf, state, share, helper_share_data);
@@ -214,8 +214,8 @@ pub(crate) fn prio3_leader_prepare_finish(
         }
         (
             Prio3Config::Histogram { buckets },
-            VdafState::Prio3Field128(state),
-            VdafMessage::Prio3ShareField128(share),
+            VdafPrepState::Prio3Field128(state),
+            VdafPrepMessage::Prio3ShareField128(share),
         ) => {
             let vdaf = Prio3::new_histogram(2, buckets)?;
             let (out_share, outbound) = leader_prep_fin!(vdaf, state, share, helper_share_data);
@@ -224,8 +224,8 @@ pub(crate) fn prio3_leader_prepare_finish(
         }
         (
             Prio3Config::Sum { bits },
-            VdafState::Prio3Field128(state),
-            VdafMessage::Prio3ShareField128(share),
+            VdafPrepState::Prio3Field128(state),
+            VdafPrepMessage::Prio3ShareField128(share),
         ) => {
             let vdaf = Prio3::new_sum(2, *bits)?;
             let (out_share, outbound) = leader_prep_fin!(vdaf, state, share, helper_share_data);
@@ -234,8 +234,8 @@ pub(crate) fn prio3_leader_prepare_finish(
         }
         (
             Prio3Config::SumVec { bits, len },
-            VdafState::Prio3Field128(state),
-            VdafMessage::Prio3ShareField128(share),
+            VdafPrepState::Prio3Field128(state),
+            VdafPrepMessage::Prio3ShareField128(share),
         ) => {
             let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
             let (out_share, outbound) = leader_prep_fin!(vdaf, state, share, helper_share_data);
@@ -272,26 +272,26 @@ macro_rules! helper_prep_fin {
 /// Consume the peer's prepare message and return an output share.
 pub(crate) fn prio3_helper_prepare_finish(
     config: &Prio3Config,
-    state: VdafState,
+    state: VdafPrepState,
     peer_message_data: &[u8],
 ) -> Result<VdafAggregateShare, VdafError> {
     let agg_share = match (&config, state) {
-        (Prio3Config::Count, VdafState::Prio3Field64(state)) => {
+        (Prio3Config::Count, VdafPrepState::Prio3Field64(state)) => {
             let vdaf = Prio3::new_count(2)?;
             let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
             VdafAggregateShare::Field64(vdaf.aggregate(&(), [out_share])?)
         }
-        (Prio3Config::Histogram { buckets }, VdafState::Prio3Field128(state)) => {
+        (Prio3Config::Histogram { buckets }, VdafPrepState::Prio3Field128(state)) => {
             let vdaf = Prio3::new_histogram(2, buckets)?;
             let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
             VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?)
         }
-        (Prio3Config::Sum { bits }, VdafState::Prio3Field128(state)) => {
+        (Prio3Config::Sum { bits }, VdafPrepState::Prio3Field128(state)) => {
             let vdaf = Prio3::new_sum(2, *bits)?;
             let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
             VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?)
         }
-        (Prio3Config::SumVec { bits, len }, VdafState::Prio3Field128(state)) => {
+        (Prio3Config::SumVec { bits, len }, VdafPrepState::Prio3Field128(state)) => {
             let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
             let out_share = helper_prep_fin!(vdaf, state, peer_message_data);
             VdafAggregateShare::Field128(vdaf.aggregate(&(), [out_share])?)
@@ -302,69 +302,40 @@ pub(crate) fn prio3_helper_prepare_finish(
     Ok(agg_share)
 }
 
-/// Interpret `step` as a prepare message for prio3 and append it to `bytes`. Returns an error if
-/// the `step` is not compatible with `param`.
-pub(crate) fn prio3_append_prepare_state(
-    bytes: &mut Vec<u8>,
-    config: &Prio3Config,
-    state: &VdafState,
-) -> Result<(), VdafError> {
-    match (&config, state) {
-        (Prio3Config::Count, VdafState::Prio3Field64(state)) => {
-            state.encode(bytes);
-        }
-        (Prio3Config::Histogram { buckets: _ }, VdafState::Prio3Field128(state))
-        | (Prio3Config::Sum { bits: _ }, VdafState::Prio3Field128(state)) => {
-            state.encode(bytes);
-        }
-        _ => panic!("prio3_append_prepare_state: {ERR_FIELD_TYPE}"),
-    }
-    Ok(())
-}
-
 /// Parse a prio3 prepare message from the front of `reader` whose type is compatible with `param`.
 pub(crate) fn prio3_decode_prepare_state(
     config: &Prio3Config,
     agg_id: usize,
     bytes: &mut Cursor<&[u8]>,
-) -> Result<VdafState, VdafError> {
+) -> Result<VdafPrepState, VdafError> {
     match config {
         Prio3Config::Count => {
             let vdaf = Prio3::new_count(2)?;
-            Ok(VdafState::Prio3Field64(
+            Ok(VdafPrepState::Prio3Field64(
                 Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
             ))
         }
         Prio3Config::Histogram { buckets } => {
             let vdaf = Prio3::new_histogram(2, buckets)?;
-            Ok(VdafState::Prio3Field128(
+            Ok(VdafPrepState::Prio3Field128(
                 Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
             ))
         }
         Prio3Config::Sum { bits } => {
             let vdaf = Prio3::new_sum(2, *bits)?;
-            Ok(VdafState::Prio3Field128(
+            Ok(VdafPrepState::Prio3Field128(
                 Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
             ))
         }
         Prio3Config::SumVec { bits, len } => {
             let vdaf = Prio3::new_sum_vec(2, *bits, *len)?;
-            Ok(VdafState::Prio3Field128(
+            Ok(VdafPrepState::Prio3Field128(
                 Prio3PrepareState::decode_with_param(&(&vdaf, agg_id), bytes)?,
             ))
         }
     }
 }
 
-/// Encode `message` as a byte string.
-pub(crate) fn prio3_encode_prepare_message(message: &VdafMessage) -> Vec<u8> {
-    match message {
-        VdafMessage::Prio3ShareField64(message) => message.get_encoded(),
-        VdafMessage::Prio3ShareField128(message) => message.get_encoded(),
-        _ => panic!("prio3_encode_prepare_message: unexpected message type"),
-    }
-}
-
 macro_rules! unshard {
     (
         $vdaf:ident,
diff --git a/daphne_worker/Cargo.toml b/daphne_worker/Cargo.toml
index e6dc7ab..3836439 100644
--- a/daphne_worker/Cargo.toml
+++ b/daphne_worker/Cargo.toml
@@ -18,6 +18,7 @@ readme = "../README.md"
 crate-type = ["cdylib", "rlib"]
 
 [dependencies]
+arc-swap = "1.6.0"
 async-trait.workspace = true
 chrono = { version = "0.4.26", default-features = false, features = ["clock", "wasmbind"] }
 daphne = { path = "../daphne" }
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index 87523ac..c30591d 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -451,6 +451,9 @@ impl<'srv> DaphneWorkerRequestState<'srv> {
         let prometheus_registry = Registry::new();
         let metrics = DaphneWorkerMetrics::register(&prometheus_registry, None)
             .map_err(|e| Error::RustError(format!("failed to register metrics: {e}")))?;
+
+        crate::tracing_utils::initialize_timing_histograms(&prometheus_registry, None)
+            .map_err(|e| Error::RustError(format!("failed to register metrics: {e}")))?;
         let host = req
             .url()?
             .host_str()
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 98cf8d3..1695c08 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -954,7 +954,7 @@ impl<'srv> DapHelper<DaphneWorkerAuth> for DaphneWorker<'srv> {
         helper_state: &DapHelperState,
     ) -> std::result::Result<bool, DapError> {
         let task_config = self.try_get_task_config(task_id).await?;
-        let helper_state_hex = hex::encode(helper_state.get_encoded(&task_config.as_ref().vdaf)?);
+        let helper_state_hex = hex::encode(helper_state.get_encoded());
         Ok(self
             .durable()
             .with_retry()
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 1422233..86db059 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -157,6 +157,7 @@ pub use crate::tracing_utils::initialize_tracing;
 use crate::{
     config::{DaphneWorkerIsolateState, DaphneWorkerRequestState},
     dap::dap_response_to_worker,
+    tracing_utils::MeasuredSpanName,
 };
 use daphne::{
     audit_log::{AuditLog, NoopAuditLog},
@@ -660,7 +661,15 @@ async fn handle_agg_job(
     let daph = ctx.data.handler(&ctx.env);
     let req = daph.worker_request_to_dap(req, &ctx).await?;
 
-    let span = info_span_from_dap_request!("aggregate", req);
+    let span = match req.media_type {
+        DapMediaType::AggregationJobInitReq => {
+            info_span_from_dap_request!(MeasuredSpanName::AggregateInit.as_str(), req)
+        }
+        DapMediaType::AggregationJobContinueReq => {
+            info_span_from_dap_request!(MeasuredSpanName::AggregateContinue.as_str(), req)
+        }
+        _ => info_span_from_dap_request!("aggregate", req),
+    };
 
     match daph.handle_agg_job_req(&req).instrument(span).await {
         Ok(resp) => dap_response_to_worker(resp),
@@ -675,7 +684,7 @@ async fn handle_agg_share_req(
     let daph = ctx.data.handler(&ctx.env);
     let req = daph.worker_request_to_dap(req, &ctx).await?;
 
-    let span = info_span_from_dap_request!("aggregate_share", req);
+    let span = info_span_from_dap_request!(MeasuredSpanName::AggregateShares.as_str(), req);
 
     match daph.handle_agg_share_req(&req).instrument(span).await {
         Ok(resp) => dap_response_to_worker(resp),
diff --git a/daphne_worker/src/metrics.rs b/daphne_worker/src/metrics.rs
index a52d986..a31bd7d 100644
--- a/daphne_worker/src/metrics.rs
+++ b/daphne_worker/src/metrics.rs
@@ -7,7 +7,7 @@ use crate::DapError;
 use daphne::{fatal_error, metrics::DaphneMetrics};
 use prometheus::{register_int_counter_vec_with_registry, IntCounterVec, Registry};
 
-pub(crate) struct DaphneWorkerMetrics {
+pub struct DaphneWorkerMetrics {
     /// Daphne metrics.
     pub(crate) daphne: DaphneMetrics,
 
diff --git a/daphne_worker/src/tracing_utils/mod.rs b/daphne_worker/src/tracing_utils/mod.rs
index 4238a0b..0a34867 100644
--- a/daphne_worker/src/tracing_utils/mod.rs
+++ b/daphne_worker/src/tracing_utils/mod.rs
@@ -3,17 +3,17 @@
 
 use chrono::{SecondsFormat, Utc};
 use std::{fmt::Result as FmtResult, io, str, sync::Once};
-use tracing::{event, Level, Subscriber};
-use tracing_core::span;
 use tracing_subscriber::{
     fmt,
     fmt::{format::Writer, time::FormatTime},
-    layer::Context as LayerContext,
     layer::*,
     prelude::*,
     registry, EnvFilter, Layer,
 };
-use worker::*;
+use worker::{console_log, Env};
+
+use wasm_timing::WasmTimingLayer;
+pub(crate) use wasm_timing::{initialize_timing_histograms, MeasuredSpanName};
 
 use self::workers_json_layer::WorkersJsonLayer;
 
@@ -72,79 +72,6 @@ impl io::Write for LogWriter {
     }
 }
 
-/// WasmTimingLayer provides a span's elapsed time.
-pub struct WasmTimingLayer {}
-
-impl WasmTimingLayer {
-    fn new() -> Self {
-        WasmTimingLayer {}
-    }
-}
-
-fn milli_now() -> i64 {
-    Date::now().as_millis() as i64
-}
-
-struct Timestamps {
-    busy: i64,
-    entered_at: i64,
-    started_at: i64,
-}
-
-impl Timestamps {
-    fn new() -> Self {
-        Timestamps {
-            busy: 0,
-            entered_at: 0,
-            started_at: milli_now(),
-        }
-    }
-}
-
-impl<S> Layer<S> for WasmTimingLayer
-where
-    S: Subscriber + for<'a> registry::LookupSpan<'a>,
-{
-    fn on_new_span(&self, _attrs: &span::Attributes<'_>, id: &span::Id, ctx: LayerContext<'_, S>) {
-        let span = ctx.span(id).expect("span should exist!");
-        let mut extensions = span.extensions_mut();
-        if extensions.get_mut::<Timestamps>().is_none() {
-            extensions.insert(Timestamps::new());
-        }
-    }
-
-    fn on_enter(&self, id: &span::Id, ctx: LayerContext<'_, S>) {
-        let span = ctx.span(id).expect("span should exist!");
-        let mut extensions = span.extensions_mut();
-        if let Some(timestamps) = extensions.get_mut::<Timestamps>() {
-            timestamps.entered_at = milli_now();
-        }
-    }
-
-    fn on_exit(&self, id: &span::Id, ctx: LayerContext<'_, S>) {
-        let span = ctx.span(id).expect("span should exist!");
-        let mut extensions = span.extensions_mut();
-        if let Some(timestamps) = extensions.get_mut::<Timestamps>() {
-            timestamps.busy += milli_now().saturating_sub(timestamps.entered_at);
-        }
-    }
-
-    fn on_close(&self, id: span::Id, ctx: LayerContext<'_, S>) {
-        let span = ctx.span(&id).expect("span should exist!");
-        let extensions = span.extensions();
-        if let Some(timestamps) = extensions.get::<Timestamps>() {
-            let elapsed = milli_now().saturating_sub(timestamps.started_at);
-            event!(
-                parent: id,
-                Level::INFO,
-                busy = timestamps.busy,
-                elapsed,
-                unit = "ms"
-            );
-        }
-    }
-}
-
 static INITIALIZE_TRACING: Once = Once::new();
 
 /// Setup logging.
@@ -168,7 +95,7 @@ pub fn initialize_tracing(env: &Env) {
         let (ansi, json) = match env.var("DAP_DEPLOYMENT") {
             Ok(format) if format.to_string() == "prod" => {
                 // JSON output
-                let json = WorkersJsonLayer.and_then(WasmTimingLayer::new());
+                let json = WorkersJsonLayer.and_then(WasmTimingLayer);
                 (None, Some(json))
             }
             Ok(_) | Err(_) => {
@@ -177,7 +104,7 @@ pub fn initialize_tracing(env: &Env) {
                     .with_ansi(true)
                     .with_writer(LogWriter::new)
                     .with_timer(WasmTime::new())
-                    .and_then(WasmTimingLayer::new());
+                    .and_then(WasmTimingLayer);
 
                 (Some(ansi), None)
             }
@@ -191,4 +118,5 @@ pub fn initialize_tracing(env: &Env) {
     });
 }
 
+mod wasm_timing;
 mod workers_json_layer;
diff --git a/daphne_worker/src/tracing_utils/wasm_timing.rs b/daphne_worker/src/tracing_utils/wasm_timing.rs
new file mode 100644
index 0000000..59af2a8
--- /dev/null
+++ b/daphne_worker/src/tracing_utils/wasm_timing.rs
@@ -0,0 +1,188 @@
+use arc_swap::ArcSwapOption;
+use daphne::{fatal_error, DapError};
+use prometheus::{register_histogram_with_registry, Histogram, Registry};
+use std::{str, sync::Arc};
+use tracing::{event, Level, Subscriber};
+use tracing_core::span;
+use tracing_subscriber::{layer::Context as LayerContext, registry, Layer};
+use worker::Date;
+
+#[derive(Debug)]
+#[allow(clippy::enum_variant_names)] // the common prefix makes sense, we just don't have more
+                                     // spans we want to measure yet.
+pub(crate) enum MeasuredSpanName {
+    AggregateInit,
+    AggregateContinue,
+    AggregateShares,
+}
+
+impl MeasuredSpanName {
+    pub const fn as_str(&self) -> &'static str {
+        match self {
+            Self::AggregateInit => "aggregate_init",
+            Self::AggregateContinue => "aggregate_continue",
+            Self::AggregateShares => "aggregate_shares",
+        }
+    }
+}
+
+impl TryFrom<&str> for MeasuredSpanName {
+    type Error = &'static str;
+
+    fn try_from(value: &str) -> std::result::Result<Self, Self::Error> {
+        match value {
+            "aggregate_init" => Ok(Self::AggregateInit),
+            "aggregate_continue" => Ok(Self::AggregateContinue),
+            "aggregate_shares" => Ok(Self::AggregateShares),
+            _ => Err("invalid request type"),
+        }
+    }
+}
+
+struct RequestTimeHistograms {
+    /// Helper: Aggregate request times
+    pub aggregation_init_request_times: Histogram,
+
+    /// Helper: Aggregate request times
+    pub aggregation_continue_request_times: Histogram,
+
+    /// Helper: Aggregate shares request times
+    pub aggregate_shares_request_times: Histogram,
+}
+
+static REQUEST_TIME_HISTOGRAMS: ArcSwapOption<RequestTimeHistograms> = ArcSwapOption::const_empty();
+
+pub(crate) fn initialize_timing_histograms(
+    registry: &Registry,
+    prefix: Option<&str>,
+) -> Result<(), DapError> {
+    let front = if let Some(prefix) = prefix {
+        format!("{prefix}_")
+    } else {
+        "".into()
+    };
+    let aggregation_init_request_times = register_histogram_with_registry!(
+        format!("{front}aggregation_init_request_times"),
+        "Time taken to process an aggregation",
+        prometheus::DEFAULT_BUCKETS
+            .iter()
+            .map(|i| i * 1000.0)
+            .collect(),
+        registry,
+    )
+    .map_err(|e| fatal_error!(err = e, "failed to register aggregation_request_times"))?;
+
+    let aggregation_continue_request_times = register_histogram_with_registry!(
+        format!("{front}aggregation_continue_request_times"),
+        "Time taken to process an aggregation",
+        prometheus::DEFAULT_BUCKETS
+            .iter()
+            .map(|i| i * 1000.0)
+            .collect(),
+        registry,
+    )
+    .map_err(|e| fatal_error!(err = e, "failed to register aggregate_shares_request_times"))?;
+
+    let aggregate_shares_request_times = register_histogram_with_registry!(
+        format!("{front}aggregate_shares_request_times"),
+        "Time taken to process an aggregation",
+        prometheus::DEFAULT_BUCKETS
+            .iter()
+            .map(|i| i * 1000.0)
+            .collect(),
+        registry,
+    )
+    .map_err(|e| fatal_error!(err = e, "failed to register upload_request_times"))?;
+
+    REQUEST_TIME_HISTOGRAMS.store(Some(Arc::new(RequestTimeHistograms {
+        aggregation_init_request_times,
+        aggregation_continue_request_times,
+        aggregate_shares_request_times,
+    })));
+    Ok(())
+}
+
+/// WasmTimingLayer provides a span's elapsed time.
+pub(super) struct WasmTimingLayer;
+
+fn milli_now() -> i64 {
+    Date::now().as_millis() as i64
+}
+
+struct Timestamps {
+    busy: i64,
+    entered_at: i64,
+    started_at: i64,
+}
+
+impl Timestamps {
+    fn new() -> Self {
+        Timestamps {
+            busy: 0,
+            entered_at: 0,
+            started_at: milli_now(),
+        }
+    }
+}
+
+impl<S> Layer<S> for WasmTimingLayer
+where
+    S: Subscriber + for<'a> registry::LookupSpan<'a>,
+{
+    fn on_new_span(&self, _attrs: &span::Attributes<'_>, id: &span::Id, ctx: LayerContext<'_, S>) {
+        let span = ctx.span(id).expect("span should exist!");
+        let mut extensions = span.extensions_mut();
+        if extensions.get_mut::<Timestamps>().is_none() {
+            extensions.insert(Timestamps::new());
+        }
+    }
+
+    fn on_enter(&self, id: &span::Id, ctx: LayerContext<'_, S>) {
+        let span = ctx.span(id).expect("span should exist!");
+        let mut extensions = span.extensions_mut();
+        if let Some(timestamps) = extensions.get_mut::<Timestamps>() {
+            timestamps.entered_at = milli_now();
+        }
+    }
+
+    fn on_exit(&self, id: &span::Id, ctx: LayerContext<'_, S>) {
+        let span = ctx.span(id).expect("span should exist!");
+        let mut extensions = span.extensions_mut();
+        if let Some(timestamps) = extensions.get_mut::<Timestamps>() {
+            timestamps.busy += milli_now().saturating_sub(timestamps.entered_at);
+        }
+    }
+
+    fn on_close(&self, id: span::Id, ctx: LayerContext<'_, S>) {
+        let span = ctx.span(&id).expect("span should exist!");
+        let extensions = span.extensions();
+        if let Some(timestamps) = extensions.get::<Timestamps>() {
+            let elapsed = milli_now().saturating_sub(timestamps.started_at);
+            if let Some(histograms) = &*REQUEST_TIME_HISTOGRAMS.load() {
+                match span.name().try_into() {
+                    Ok(ty) => match ty {
+                        MeasuredSpanName::AggregateInit => histograms
+                            .aggregation_init_request_times
+                            .observe(elapsed as _),
+                        MeasuredSpanName::AggregateContinue => histograms
+                            .aggregation_continue_request_times
+                            .observe(elapsed as _),
+                        MeasuredSpanName::AggregateShares => histograms
+                            .aggregate_shares_request_times
+                            .observe(elapsed as _),
+                    },
+                    Err(_ignored) => {}
+                }
+            }
+            event!(
+                parent: id,
+                Level::INFO,
+                busy = timestamps.busy,
+                elapsed,
+                unit = "ms",
+                "{} finished",
+                span.name(),
+            );
+        }
+    }
+}
diff --git a/daphne_worker/src/tracing_utils/workers_json_layer.rs b/daphne_worker/src/tracing_utils/workers_json_layer.rs
index 48bd7ca..a6a1479 100644
--- a/daphne_worker/src/tracing_utils/workers_json_layer.rs
+++ b/daphne_worker/src/tracing_utils/workers_json_layer.rs
@@ -124,8 +124,7 @@ where
 
         let metadata = event.metadata();
         if let (Some(file), Some(line)) = (metadata.file(), metadata.line()) {
-            fields.insert("source_file".to_owned(), file.into());
-            fields.insert("source_line".to_owned(), line.into());
+            fields.insert("at".to_owned(), format!("{}:{}", file, line).into());
         }
 
         // If there is no `message`, repurpose the error meessage is there is one or the
@@ -133,8 +132,8 @@ where
         if !fields.contains_key("message") {
             if let Some(error) = fields.get("error") {
                 fields.insert("message".into(), error.clone());
-            } else if let Some(span_name) = fields.remove("current_span") {
-                fields.insert("message".to_owned(), span_name);
+            } else {
+                fields.insert("message".to_owned(), "(no message)".into());
             }
         }
 


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 52c09f2..39ebb4a 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -24,6 +24,7 @@ use crate::{
 };
 use assert_matches::assert_matches;
 use async_trait::async_trait;
+use prio::codec::Encode;
 use rand::{thread_rng, Rng};
 use serde::{Deserialize, Serialize};
 use std::{
@@ -1396,13 +1397,8 @@ impl AggregationJobTest {
             .await else {
                 panic!("unexpected transition");
         };
-        let got = DapHelperState::get_decoded(
-            &self.task_config.vdaf,
-            &helper_state
-                .get_encoded(&self.task_config.vdaf)
-                .expect("failed to encode helper state"),
-        )
-        .expect("failed to decode helper state");
+        let got = DapHelperState::get_decoded(&self.task_config.vdaf, &helper_state.get_encoded())
+            .expect("failed to decode helper state");
         assert_eq!(got, helper_state);
 
         let DapLeaderTransition::Uncommitted(uncommitted, agg_cont) = self
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index a78beaa..eb80591 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -14,12 +14,13 @@ use crate::{
     testing::AggregationJobTest,
     DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
     DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement, DapOutputShare,
-    DapVersion, Prio3Config, VdafAggregateShare, VdafConfig, VdafMessage, VdafState,
+    DapVersion, Prio3Config, VdafAggregateShare, VdafConfig, VdafPrepMessage, VdafPrepState,
 };
 use assert_matches::assert_matches;
 use hpke_rs::HpkePublicKey;
 use paste::paste;
 use prio::{
+    codec::Encode,
     field::Field64,
     vdaf::{
         prio3::Prio3, Aggregatable, AggregateShare, Aggregator as VdafAggregator,
@@ -29,6 +30,8 @@ use prio::{
 use rand::prelude::*;
 use std::fmt::Debug;
 
+use super::{EarlyReportStateConsumed, EarlyReportStateInitialized};
+
 impl<M: Debug> DapLeaderTransition<M> {
     pub(crate) fn unwrap_continue(self) -> (DapLeaderState, M) {
         match self {
@@ -86,38 +89,44 @@ async fn roundtrip_report(version: DapVersion) {
         )
         .unwrap();
 
-    let (leader_step, leader_share) = TEST_VDAF
-        .consume_report_share(
-            &t.leader_hpke_receiver_config,
-            true, // is_leader
-            &t.task_id,
-            &t.task_config,
-            &report.report_metadata,
-            &report.public_share,
-            &report.encrypted_input_shares[0],
-        )
-        .await
-        .unwrap();
+    let early_report_state_consumed = EarlyReportStateConsumed::consume(
+        &t.leader_hpke_receiver_config,
+        true, // is_leader
+        &t.task_id,
+        &t.task_config,
+        &report.report_metadata,
+        &report.public_share,
+        &report.encrypted_input_shares[0],
+    )
+    .await
+    .unwrap();
+    let EarlyReportStateInitialized::Ready{ state: leader_step, message: leader_share } =
+        EarlyReportStateInitialized::initialize(true, &t.task_config, early_report_state_consumed).unwrap() else {
+        panic!("rejected unexpectedly");
+    };
 
-    let (helper_step, helper_share) = TEST_VDAF
-        .consume_report_share(
-            &t.helper_hpke_receiver_config,
-            false, // is_leader
-            &t.task_id,
-            &t.task_config,
-            &report.report_metadata,
-            &report.public_share,
-            &report.encrypted_input_shares[1],
-        )
-        .await
-        .unwrap();
+    let early_report_state_consumed = EarlyReportStateConsumed::consume(
+        &t.helper_hpke_receiver_config,
+        false, // is_helper
+        &t.task_id,
+        &t.task_config,
+        &report.report_metadata,
+        &report.public_share,
+        &report.encrypted_input_shares[1],
+    )
+    .await
+    .unwrap();
+    let EarlyReportStateInitialized::Ready{ state: helper_step, message: helper_share } =
+        EarlyReportStateInitialized::initialize(false, &t.task_config, early_report_state_consumed).unwrap() else {
+        panic!("rejected unexpectedly");
+    };
 
     match (leader_step, helper_step, leader_share, helper_share) {
         (
-            VdafState::Prio3Field64(leader_step),
-            VdafState::Prio3Field64(helper_step),
-            VdafMessage::Prio3ShareField64(leader_share),
-            VdafMessage::Prio3ShareField64(helper_share),
+            VdafPrepState::Prio3Field64(leader_step),
+            VdafPrepState::Prio3Field64(helper_step),
+            VdafPrepMessage::Prio3ShareField64(leader_share),
+            VdafPrepMessage::Prio3ShareField64(helper_share),
         ) => {
             let vdaf = Prio3::new_count(2).unwrap();
             let message = vdaf
@@ -716,8 +725,7 @@ async fn helper_state_serialization(version: DapVersion) {
         .await
         .unwrap_continue();
 
-    let got =
-        DapHelperState::get_decoded(TEST_VDAF, &want.get_encoded(TEST_VDAF).unwrap()).unwrap();
+    let got = DapHelperState::get_decoded(TEST_VDAF, &want.get_encoded()).unwrap();
     assert_eq!(got, want);
 
     assert!(DapHelperState::get_decoded(TEST_VDAF, b"invalid helper state").is_err())
diff --git a/daphne/src/vdaf/prio3_test.rs b/daphne/src/vdaf/prio3_test.rs
index 97998b6..62fd8b3 100644
--- a/daphne/src/vdaf/prio3_test.rs
+++ b/daphne/src/vdaf/prio3_test.rs
@@ -4,8 +4,8 @@
 use crate::{
     vdaf::{
         prio3::{
-            prio3_encode_prepare_message, prio3_helper_prepare_finish, prio3_leader_prepare_finish,
-            prio3_prepare_init, prio3_shard, prio3_unshard,
+            prio3_helper_prepare_finish, prio3_leader_prepare_finish, prio3_prepare_init,
+            prio3_shard, prio3_unshard,
         },
         VdafError,
     },
@@ -96,10 +96,12 @@ fn test_prepare(
         &encoded_input_shares[1],
     )?;
 
-    let helper_share_data = prio3_encode_prepare_message(&helper_share);
-
-    let (leader_out_share, leader_message_data) =
-        prio3_leader_prepare_finish(config, leader_state, leader_share, &helper_share_data)?;
+    let (leader_out_share, leader_message_data) = prio3_leader_prepare_finish(
+        config,
+        leader_state,
+        leader_share,
+        &helper_share.get_encoded(),
+    )?;
 
     let helper_out_share = prio3_helper_prepare_finish(config, helper_state, &leader_message_data)?;