Repository: cloudflare/daphne
Pull request #366 (2023-07-10 02:43:42)

Parallelize VDAF preparation initialization
Define a trait `DapReportInitializer` that takes a sequence of reports
in the "consumed" state and transitions each report to the "initialized"
state. Implementers are expected to perform the early validation steps
(i.e., the protocol logic previously implemented by
`DapAggregator::check_early_reject()`) and initialize VDAF preparation.

The primary motivation for this change is to improve the performance of
`DaphneWorker`. VDAF preparation initialization is expensive, so we'd
like to parallelize the computation. Because the Workers runtime is
single-threaded, the best option is to offload VDAF preparation to
durable objects. In particular, `DaphneWorker`'s implementation of
`DapReportInitializer` piggy-backs this computation on the requests to
`ReportsProcessed`.

cjpatton: Squashed.
cjpatton: Updated commit message.

Patch:
diff --git a/daphne/benches/aggregation.rs b/daphne/benches/aggregation.rs
index c99d9da..f944c12 100644
--- a/daphne/benches/aggregation.rs
+++ b/daphne/benches/aggregation.rs
@@ -1,7 +1,7 @@
 // Copyright (c) 2023 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use criterion::{criterion_group, criterion_main, Criterion};
+use criterion::{black_box, criterion_group, criterion_main, Criterion};
 use daphne::{
     hpke::HpkeKemId, testing::AggregationJobTest, DapLeaderTransition, DapMeasurement, DapVersion,
     Prio3Config, VdafConfig,
@@ -46,9 +46,7 @@ fn handle_agg_job_init_req(c: &mut Criterion) {
 
         c.bench_function(&format!("handle_agg_job_init_req {vdaf:?}"), |b| {
             b.to_async(&rt).iter(|| async {
-                agg_job_test
-                    .handle_agg_job_init_req(&agg_job_init_req)
-                    .await
+                black_box(agg_job_test.handle_agg_job_init_req(&agg_job_init_req)).await
             })
         });
     }
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index b6693f0..687d8ba 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -42,8 +42,7 @@ use crate::{
     hpke::HpkeReceiverConfig,
     messages::{
         AggregationJobId, BatchId, BatchSelector, Collection, CollectionJobId,
-        Draft02AggregationJobId, Duration, Interval, PartialBatchSelector, ReportId,
-        ReportMetadata, TaskId, Time,
+        Draft02AggregationJobId, Duration, Interval, PartialBatchSelector, ReportId, TaskId, Time,
     },
     taskprov::TaskprovVersion,
     vdaf::{VdafAggregateShare, VdafPrepMessage, VdafPrepState, VdafVerifyKey},
@@ -66,6 +65,7 @@ use std::{
     fmt::{Debug, Display},
 };
 use url::Url;
+use vdaf::{EarlyReportState, EarlyReportStateConsumed};
 
 /// DAP version used for a task.
 #[derive(Clone, Copy, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
@@ -364,12 +364,13 @@ impl DapTaskConfig {
         }
     }
 
-    /// Return the batch span of a set of reports with the given metadata.
+    /// Return the batch span of a set of reports.
     pub fn batch_span_for_meta<'sel, 'rep>(
         &self,
         part_batch_sel: &'sel PartialBatchSelector,
-        report_meta: impl Iterator<Item = &'rep ReportMetadata>,
-    ) -> Result<HashMap<DapBatchBucket<'sel>, Vec<&'rep ReportMetadata>>, DapError> {
+        consumed_reports: impl Iterator<Item = &'rep EarlyReportStateConsumed<'rep>>,
+    ) -> Result<HashMap<DapBatchBucket<'sel>, Vec<&'rep EarlyReportStateConsumed<'rep>>>, DapError>
+    {
         if !self.query.is_valid_part_batch_sel(part_batch_sel) {
             return Err(fatal_error!(
                 err = "partial batch selector not compatible with task",
@@ -377,18 +378,19 @@ impl DapTaskConfig {
         }
 
         let mut span: HashMap<_, Vec<_>> = HashMap::new();
-        for metadata in report_meta {
+        for consumed_report in consumed_reports.filter(|consumed_report| consumed_report.is_ready())
+        {
             let bucket = match part_batch_sel {
                 PartialBatchSelector::TimeInterval => DapBatchBucket::TimeInterval {
-                    batch_window: self.quantized_time_lower_bound(metadata.time),
+                    batch_window: self.quantized_time_lower_bound(consumed_report.metadata().time),
                 },
                 PartialBatchSelector::FixedSizeByBatchId { batch_id } => {
                     DapBatchBucket::FixedSize { batch_id }
                 }
             };
 
-            let report_ids = span.entry(bucket).or_default();
-            report_ids.push(metadata);
+            let consumed_reports_per_bucket = span.entry(bucket).or_default();
+            consumed_reports_per_bucket.push(consumed_report);
         }
 
         Ok(span)
@@ -481,7 +483,7 @@ impl DapHelperState {
             .map_err(|e| DapAbort::from_codec_error(e, None))?;
         let mut seq = vec![];
         while (r.position() as usize) < data.len() {
-            let state = VdafPrepState::decode_with_param(vdaf_config, &mut r)
+            let state = VdafPrepState::decode_with_param(&(vdaf_config, false), &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 =
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 28c7548..18131da 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -12,14 +12,15 @@ use crate::{
         constant_time_eq, decode_base64url, AggregateShare, AggregateShareReq,
         AggregationJobContinueReq, AggregationJobInitReq, AggregationJobResp, BatchId,
         BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeConfigList, Interval,
-        PartialBatchSelector, Query, Report, ReportId, ReportMetadata, TaskId, Time,
-        TransitionFailure, TransitionVar,
+        PartialBatchSelector, Query, Report, ReportMetadata, TaskId, Time, TransitionFailure,
     },
     metrics::{DaphneMetrics, DaphneRequestType},
-    taskprov, DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig,
-    DapHelperState, DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition,
-    DapOutputShare, DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig,
-    DapVersion, MetaAggregationJobId,
+    taskprov,
+    vdaf::{EarlyReportStateConsumed, EarlyReportStateInitialized},
+    DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
+    DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition, DapOutputShare,
+    DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig, DapVersion,
+    MetaAggregationJobId,
 };
 use async_trait::async_trait;
 use futures::TryFutureExt;
@@ -41,9 +42,26 @@ pub trait DapAuthorizedSender<S> {
     ) -> Result<S, DapError>;
 }
 
+/// Report initializer. Used by a DAP Aggregator [`DapAggregator`] when initializing an aggregation
+/// job.
+#[async_trait(?Send)]
+pub trait DapReportInitializer {
+    /// Initialize a sequence of reports that are in the "consumed" state by performing the early
+    /// validation steps (check if the report was replayed, belongs to a batch that has been
+    /// collected) and initializing VDAF preparation.
+    async fn initialize_reports<'req>(
+        &self,
+        is_leader: bool,
+        task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        part_batch_sel: &PartialBatchSelector,
+        consumed_reports: Vec<EarlyReportStateConsumed<'req>>,
+    ) -> Result<Vec<EarlyReportStateInitialized<'req>>, DapError>;
+}
+
 /// DAP Aggregator functionality.
 #[async_trait(?Send)]
-pub trait DapAggregator<S>: HpkeDecrypter + Sized {
+pub trait DapAggregator<S>: HpkeDecrypter + DapReportInitializer + Sized {
     /// A refernce to a task configuration stored by the Aggregator.
     type WrappedDapTaskConfig<'a>: AsRef<DapTaskConfig>;
 
@@ -120,16 +138,6 @@ pub trait DapAggregator<S>: HpkeDecrypter + Sized {
         batch_sel: &BatchSelector,
     ) -> Result<DapAggregateShare, DapError>;
 
-    /// Ensure a set of reorts can be aggregated. Return a transition failure for each report
-    /// that must be rejected early, due to the repot being replayed, the bucket that contains the
-    /// report being collected, etc.
-    async fn check_early_reject(
-        &self,
-        task_id: &TaskId,
-        part_batch_sel: &PartialBatchSelector,
-        report_meta: impl Iterator<Item = &ReportMetadata>,
-    ) -> Result<HashMap<ReportId, TransitionFailure>, DapError>;
-
     /// Mark a batch as collected.
     async fn mark_collected(
         &self,
@@ -477,34 +485,12 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     ) -> Result<u64, DapAbort> {
         let metrics = self.metrics().with_host(host);
 
-        // Filter out early rejected reports.
-        //
-        // TODO Add a test similar to http_post_aggregate_init_expired_task() in roles_test.rs that
-        // verifies that the Leader properly checks for expiration. This will require extending the
-        // test framework to run run_agg_job() directly.
-        let early_rejects = self
-            .check_early_reject(
-                task_id,
-                part_batch_sel,
-                reports.iter().map(|report| &report.report_metadata),
-            )
-            .await?;
-        let reports = reports
-            .into_iter()
-            .filter(|report| {
-                if let Some(failure) = early_rejects.get(&report.report_metadata.id) {
-                    metrics.report_inc_by(&format!("rejected_{failure}"), 1);
-                    return false;
-                }
-                true
-            })
-            .collect();
-
         // Prepare AggregationJobInitReq.
         let agg_job_id = MetaAggregationJobId::gen_for_version(&task_config.version);
         let transition = task_config
             .vdaf
             .produce_agg_job_init_req(
+                self,
                 self,
                 task_id,
                 task_config,
@@ -883,68 +869,21 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     &agg_job_init_req.agg_param,
                 )?;
 
-                let early_rejects_future = self.check_early_reject(
-                    task_id,
-                    &agg_job_init_req.part_batch_sel,
-                    agg_job_init_req
-                        .report_shares
-                        .iter()
-                        .map(|report_share| &report_share.report_metadata),
-                );
-
-                let transition_future = task_config
+                let transition = task_config
                     .vdaf
                     .handle_agg_job_init_req(
+                        self,
                         self,
                         task_id,
                         task_config,
                         &agg_job_init_req,
                         &metrics,
                     )
-                    .map_err(DapError::Abort);
-
-                let (early_rejects, transition) =
-                    futures::try_join!(early_rejects_future, transition_future)?;
+                    .map_err(DapError::Abort)
+                    .await?;
 
                 let agg_job_resp = match transition {
-                    DapHelperTransition::Continue(mut state, mut agg_job_resp) => {
-                        // Filter out early rejected reports.
-                        let mut state_index = 0;
-                        for transition in agg_job_resp.transitions.iter_mut() {
-                            let early_failure = early_rejects.get(&transition.report_id);
-                            if !matches!(transition.var, TransitionVar::Failed(..))
-                                && early_failure.is_some()
-                            {
-                                // NOTE(cjpatton) Clippy wants us to use and `if let` statement to
-                                // unwrap `early_failure`. I don't think this works becauase we
-                                // only want to enter this loop if `early_failure.is_some()` and
-                                // the current `transition` is not a failure. As far as I know, `if
-                                // let` statements can't yet be combined with other conditions.
-                                #[allow(clippy::unnecessary_unwrap)]
-                                let failure = early_failure.unwrap();
-                                transition.var = TransitionVar::Failed(*failure);
-
-                                // Remove VDAF preparation state of reports that were rejected early.
-                                if transition.report_id == state.seq[state_index].2 {
-                                    let _val = state.seq.remove(state_index);
-                                } else {
-                                    // The report ID in the Helper state and Aggregate response
-                                    // must be aligned. If not, handle as an internal error.
-                                    return Err(fatal_error!(err = "report IDs not aligned").into());
-                                }
-
-                                // NOTE(cjpatton) Unlike the Leader, the Helper filters out early
-                                // rejects after processing all of the reports. (This is an
-                                // optimization intended to reduce latency.) To avoid overcounting
-                                // rejection metrics, the latter rejections take precedence. The
-                                // Leader has the opposite behavior: Early rejections are resolved
-                                // first, so take precedence.
-                                metrics.report_inc_by(&format!("rejected_{failure}"), 1);
-                            } else {
-                                state_index += 1;
-                            }
-                        }
-
+                    DapHelperTransition::Continue(state, agg_job_resp) => {
                         if !self
                             .put_helper_state_if_not_exists(task_id, &agg_job_id, &state)
                             .await?
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 0e4c5d4..207199b 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -15,6 +15,7 @@ use crate::{
         TransitionVar,
     },
     metrics::ContextualizedDaphneMetrics,
+    roles::DapReportInitializer,
     vdaf::{
         prio2::{
             prio2_helper_prepare_finish, prio2_leader_prepare_finish, prio2_prepare_init,
@@ -38,7 +39,7 @@ use prio::{
     },
 };
 use rand::prelude::*;
-use serde::{Deserialize, Serialize};
+use serde::{Deserialize, Serialize, Serializer};
 use std::{borrow::Cow, collections::HashSet, convert::TryInto};
 
 use self::{prio2::prio2_decode_prepare_state, prio3::prio3_decode_prepare_state};
@@ -77,16 +78,28 @@ impl AsRef<[u8]> for VdafVerifyKey {
     }
 }
 
+/// Report state during aggregation initialization.
+pub trait EarlyReportState {
+    fn metadata(&self) -> &ReportMetadata;
+    fn is_ready(&self) -> bool;
+}
+
 /// Report state during aggregation initialization after consuming the report share. This involves
 /// decryption as well a few validation steps.
-#[derive(Clone, Debug, PartialEq)]
+#[derive(Clone, Serialize, Deserialize)]
+#[serde(rename_all = "snake_case")]
 pub enum EarlyReportStateConsumed<'req> {
     Ready {
         metadata: Cow<'req, ReportMetadata>,
+        #[serde(with = "serialize_bytes")]
         public_share: Cow<'req, [u8]>,
+        #[serde(with = "serialize_bytes")]
         input_share: Vec<u8>,
     },
-    Rejected(TransitionFailure),
+    Rejected {
+        metadata: Cow<'req, ReportMetadata>,
+        failure: TransitionFailure,
+    },
 }
 
 impl<'req> EarlyReportStateConsumed<'req> {
@@ -96,12 +109,15 @@ impl<'req> EarlyReportStateConsumed<'req> {
         is_leader: bool,
         task_id: &TaskId,
         task_config: &DapTaskConfig,
-        metadata: &'req ReportMetadata,
-        public_share: &'req [u8],
+        metadata: Cow<'req, ReportMetadata>,
+        public_share: Cow<'req, [u8]>,
         encrypted_input_share: &HpkeCiphertext,
     ) -> Result<EarlyReportStateConsumed<'req>, DapError> {
         if metadata.time >= task_config.expiration {
-            return Ok(Self::Rejected(TransitionFailure::TaskExpired));
+            return Ok(Self::Rejected {
+                metadata,
+                failure: TransitionFailure::TaskExpired,
+            });
         }
 
         let input_share_text = match task_config.version {
@@ -124,14 +140,14 @@ impl<'req> EarlyReportStateConsumed<'req> {
         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);
+        encode_u32_bytes(&mut aad, public_share.as_ref());
 
         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(DapError::Transition(failure)) => return Ok(Self::Rejected { metadata, failure }),
             Err(e) => return Err(e),
         };
 
@@ -145,33 +161,93 @@ impl<'req> EarlyReportStateConsumed<'req> {
             },
             DapVersion::Draft04 => match PlaintextInputShare::get_decoded(&encoded_input_share) {
                 Ok(input_share) => input_share,
-                Err(..) => return Ok(Self::Rejected(TransitionFailure::UnrecognizedMessage)),
+                Err(..) => {
+                    return Ok(Self::Rejected {
+                        metadata,
+                        failure: TransitionFailure::UnrecognizedMessage,
+                    })
+                }
             },
             _ => return Err(unimplemented_version()),
         };
 
         Ok(Self::Ready {
-            metadata: Cow::Borrowed(metadata),
-            public_share: Cow::Borrowed(public_share),
+            metadata,
+            public_share,
             input_share: input_share.payload,
         })
     }
 }
 
+impl EarlyReportState for EarlyReportStateConsumed<'_> {
+    fn metadata(&self) -> &ReportMetadata {
+        match self {
+            Self::Ready { metadata, .. } => metadata,
+            Self::Rejected { metadata, .. } => metadata,
+        }
+    }
+
+    fn is_ready(&self) -> bool {
+        matches!(self, Self::Ready { .. })
+    }
+}
+
 /// Report state during aggregation initialization after the VDAF preparation step.
-pub enum EarlyReportStateInitialized {
+#[derive(Clone, Serialize)]
+#[serde(rename_all = "snake_case")]
+pub enum EarlyReportStateInitialized<'req> {
     Ready {
+        metadata: Cow<'req, ReportMetadata>,
+        #[serde(with = "serialize_bytes")]
+        public_share: Cow<'req, [u8]>,
+        #[serde(serialize_with = "serialize_encodable")]
         state: VdafPrepState,
+        #[serde(serialize_with = "serialize_encodable")]
         message: VdafPrepMessage,
     },
-    Rejected(TransitionFailure),
+    Rejected {
+        metadata: Cow<'req, ReportMetadata>,
+        failure: TransitionFailure,
+    },
+}
+
+mod serialize_bytes {
+    use serde::{de, Deserializer, Serializer};
+    pub(super) fn serialize<S, T>(x: &T, s: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+        T: AsRef<[u8]>,
+    {
+        s.serialize_str(&hex::encode(x.as_ref()))
+    }
+
+    pub(super) fn deserialize<'de, D, T>(deserializer: D) -> Result<T, D::Error>
+    where
+        D: Deserializer<'de>,
+        T: TryFrom<Vec<u8>>,
+        <T as TryFrom<Vec<u8>>>::Error: std::fmt::Display,
+    {
+        hex::deserialize::<_, Vec<u8>>(deserializer)
+            .and_then(|bytes| bytes.try_into().map_err(de::Error::custom))
+    }
+}
+
+fn serialize_encodable<S, T>(x: &T, s: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+    T: Encode,
+{
+    s.serialize_str(&hex::encode(x.get_encoded()))
 }
 
-impl EarlyReportStateInitialized {
-    pub(crate) fn initialize(
+impl<'req> EarlyReportStateInitialized<'req> {
+    /// Initialize VDAF preparation for a report. This method is meant to be called by
+    /// [`DapReportInitializer`].
+    pub fn initialize(
         is_leader: bool,
-        task_config: &DapTaskConfig,
-        early_report_state_consumed: EarlyReportStateConsumed,
+        vdaf_verify_key: &VdafVerifyKey,
+        vdaf_config: &VdafConfig,
+        early_report_state_consumed: EarlyReportStateConsumed<'req>,
     ) -> Result<Self, DapError> {
         let (metadata, public_share, input_share) = match early_report_state_consumed {
             EarlyReportStateConsumed::Ready {
@@ -179,11 +255,13 @@ impl EarlyReportStateInitialized {
                 public_share,
                 input_share,
             } => (metadata, public_share, input_share),
-            EarlyReportStateConsumed::Rejected(failure) => return Ok(Self::Rejected(failure)),
+            EarlyReportStateConsumed::Rejected { metadata, failure } => {
+                return Ok(Self::Rejected { metadata, failure })
+            }
         };
 
         let agg_id = usize::from(!is_leader);
-        let res = match (&task_config.vdaf, &task_config.vdaf_verify_key) {
+        let res = match (vdaf_config, vdaf_verify_key) {
             (VdafConfig::Prio3(ref prio3_config), VdafVerifyKey::Prio3(ref verify_key)) => {
                 prio3_prepare_init(
                     prio3_config,
@@ -208,13 +286,34 @@ impl EarlyReportStateInitialized {
         };
 
         let early_report_state_initialized = match res {
-            Ok((state, message)) => Self::Ready { state, message },
-            Err(..) => Self::Rejected(TransitionFailure::VdafPrepError),
+            Ok((state, message)) => Self::Ready {
+                metadata,
+                public_share,
+                state,
+                message,
+            },
+            Err(..) => Self::Rejected {
+                metadata,
+                failure: TransitionFailure::VdafPrepError,
+            },
         };
         Ok(early_report_state_initialized)
     }
 }
 
+impl EarlyReportState for EarlyReportStateInitialized<'_> {
+    fn metadata(&self) -> &ReportMetadata {
+        match self {
+            Self::Ready { metadata, .. } => metadata,
+            Self::Rejected { metadata, .. } => metadata,
+        }
+    }
+
+    fn is_ready(&self) -> bool {
+        matches!(self, Self::Ready { .. })
+    }
+}
+
 /// VDAF preparation state.
 #[derive(Clone, Debug, PartialEq)]
 pub enum VdafPrepState {
@@ -233,18 +332,21 @@ impl Encode for VdafPrepState {
     }
 }
 
-impl ParameterizedDecode<VdafConfig> for VdafPrepState {
+impl<'a> ParameterizedDecode<(&'a VdafConfig, bool /* is_leader */)> for VdafPrepState {
     fn decode_with_param(
-        vdaf_config: &VdafConfig,
+        (vdaf_config, is_leader): &(&VdafConfig, bool),
         bytes: &mut std::io::Cursor<&[u8]>,
     ) -> Result<Self, CodecError> {
+        let agg_id = usize::from(!is_leader);
         match vdaf_config {
             VdafConfig::Prio3(ref prio3_config) => {
-                Ok(prio3_decode_prepare_state(prio3_config, 1, bytes)
+                Ok(prio3_decode_prepare_state(prio3_config, agg_id, bytes)
+                    .map_err(|e| CodecError::Other(Box::new(e)))?)
+            }
+            VdafConfig::Prio2 { dimension } => {
+                Ok(prio2_decode_prepare_state(*dimension, agg_id, 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)))?),
         }
     }
 }
@@ -267,6 +369,25 @@ impl Encode for VdafPrepMessage {
     }
 }
 
+impl ParameterizedDecode<VdafPrepState> for VdafPrepMessage {
+    fn decode_with_param(
+        state: &VdafPrepState,
+        bytes: &mut std::io::Cursor<&[u8]>,
+    ) -> Result<Self, CodecError> {
+        match state {
+            VdafPrepState::Prio3Field64(state) => Ok(VdafPrepMessage::Prio3ShareField64(
+                Prio3PrepareShare::decode_with_param(state, bytes)?,
+            )),
+            VdafPrepState::Prio3Field128(state) => Ok(VdafPrepMessage::Prio3ShareField128(
+                Prio3PrepareShare::decode_with_param(state, bytes)?,
+            )),
+            VdafPrepState::Prio2(state) => Ok(VdafPrepMessage::Prio2Share(
+                Prio2PrepareShare::decode_with_param(state, bytes)?,
+            )),
+        }
+    }
+}
+
 #[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub(crate) enum VdafAggregateShare {
@@ -510,6 +631,7 @@ impl VdafConfig {
     pub async fn produce_agg_job_init_req(
         &self,
         decrypter: &impl HpkeDecrypter,
+        initializer: &impl DapReportInitializer,
         task_id: &TaskId,
         task_config: &DapTaskConfig,
         agg_job_id: &MetaAggregationJobId<'_>,
@@ -520,6 +642,8 @@ impl VdafConfig {
         let mut processed = HashSet::with_capacity(reports.len());
         let mut states = Vec::with_capacity(reports.len());
         let mut seq = Vec::with_capacity(reports.len());
+        let mut consumed_reports = Vec::with_capacity(reports.len());
+        let mut helper_shares = Vec::with_capacity(reports.len());
         for report in reports.into_iter() {
             if processed.contains(&report.report_metadata.id) {
                 return Err(fatal_error!(
@@ -535,39 +659,46 @@ impl VdafConfig {
                 (it.next().unwrap(), it.next().unwrap())
             };
 
-            let early_report_state_consumed = EarlyReportStateConsumed::consume(
-                decrypter,
-                true,
-                task_id,
-                task_config,
-                &report.report_metadata,
-                &report.public_share,
-                &leader_share,
-            )
-            .await?;
+            consumed_reports.push(
+                EarlyReportStateConsumed::consume(
+                    decrypter,
+                    true,
+                    task_id,
+                    task_config,
+                    Cow::Owned(report.report_metadata),
+                    Cow::Owned(report.public_share),
+                    &leader_share,
+                )
+                .await?,
+            );
+            helper_shares.push(helper_share);
+        }
 
-            let early_report_state_initialized = EarlyReportStateInitialized::initialize(
-                true,
-                task_config,
-                early_report_state_consumed,
-            )?;
+        let initialized_reports = initializer
+            .initialize_reports(true, task_id, task_config, part_batch_sel, consumed_reports)
+            .await?;
 
-            match early_report_state_initialized {
-                EarlyReportStateInitialized::Ready { state, message } => {
-                    states.push((
-                        state,
-                        message,
-                        report.report_metadata.time,
-                        report.report_metadata.id.clone(),
-                    ));
+        assert_eq!(initialized_reports.len(), helper_shares.len());
+        for (initialized_report, helper_share) in initialized_reports
+            .into_iter()
+            .zip(helper_shares.into_iter())
+        {
+            match initialized_report {
+                EarlyReportStateInitialized::Ready {
+                    metadata,
+                    public_share,
+                    state,
+                    message,
+                } => {
+                    states.push((state, message, metadata.time, metadata.id.clone()));
                     seq.push(ReportShare {
-                        report_metadata: report.report_metadata,
-                        public_share: report.public_share,
+                        report_metadata: metadata.into_owned(),
+                        public_share: public_share.into_owned(),
                         encrypted_input_share: helper_share,
                     });
                 }
 
-                EarlyReportStateInitialized::Rejected(failure) => {
+                EarlyReportStateInitialized::Rejected { failure, .. } => {
                     // Skip report that can't be processed any further.
                     metrics.report_inc_by(&format!("rejected_{failure}"), 1);
                     continue;
@@ -615,6 +746,7 @@ impl VdafConfig {
     pub(crate) async fn handle_agg_job_init_req(
         &self,
         decrypter: &impl HpkeDecrypter,
+        initializer: &impl DapReportInitializer,
         task_id: &TaskId,
         task_config: &DapTaskConfig,
         agg_job_init_req: &AggregationJobInitReq,
@@ -624,6 +756,7 @@ impl VdafConfig {
         let mut processed = HashSet::with_capacity(num_reports);
         let mut states = Vec::with_capacity(num_reports);
         let mut transitions = Vec::with_capacity(num_reports);
+        let mut consumed_reports = Vec::with_capacity(num_reports);
         for report_share in agg_job_init_req.report_shares.iter() {
             if processed.contains(&report_share.report_metadata.id) {
                 return Err(DapAbort::UnrecognizedMessage {
@@ -636,43 +769,55 @@ impl VdafConfig {
             }
             processed.insert(report_share.report_metadata.id.clone());
 
-            let early_report_state_consumed = EarlyReportStateConsumed::consume(
-                decrypter,
+            consumed_reports.push(
+                EarlyReportStateConsumed::consume(
+                    decrypter,
+                    false,
+                    task_id,
+                    task_config,
+                    Cow::Borrowed(&report_share.report_metadata),
+                    Cow::Borrowed(&report_share.public_share),
+                    &report_share.encrypted_input_share,
+                )
+                .await?,
+            );
+        }
+
+        let initialized_reports = initializer
+            .initialize_reports(
                 false,
                 task_id,
                 task_config,
-                &report_share.report_metadata,
-                &report_share.public_share,
-                &report_share.encrypted_input_share,
+                &agg_job_init_req.part_batch_sel,
+                consumed_reports,
             )
             .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((
-                        state,
-                        report_share.report_metadata.time,
-                        report_share.report_metadata.id.clone(),
-                    ));
-                    TransitionVar::Continued(message.get_encoded())
+        for initialized_report in initialized_reports.into_iter() {
+            let transition = match initialized_report {
+                EarlyReportStateInitialized::Ready {
+                    metadata,
+                    public_share: _,
+                    state,
+                    message,
+                } => {
+                    states.push((state, metadata.time, metadata.id.clone()));
+                    Transition {
+                        report_id: metadata.into_owned().id,
+                        var: TransitionVar::Continued(message.get_encoded()),
+                    }
                 }
 
-                EarlyReportStateInitialized::Rejected(failure) => {
+                EarlyReportStateInitialized::Rejected { metadata, failure } => {
                     metrics.report_inc_by(&format!("rejected_{failure}"), 1);
-                    TransitionVar::Failed(failure)
+                    Transition {
+                        report_id: metadata.into_owned().id,
+                        var: TransitionVar::Failed(failure),
+                    }
                 }
             };
 
-            transitions.push(Transition {
-                report_id: report_share.report_metadata.id.clone(),
-                var,
-            });
+            transitions.push(transition);
         }
 
         Ok(DapHelperTransition::Continue(
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 1695c08..1c21d6d 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -32,7 +32,9 @@ use crate::{
             PendingReport, ReportsPendingResult, DURABLE_REPORTS_PENDING_GET,
             DURABLE_REPORTS_PENDING_PUT,
         },
-        reports_processed::DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED,
+        reports_processed::{
+            ReportsProcessedReq, ReportsProcessedResp, DURABLE_REPORTS_PROCESSED_INITIALIZE,
+        },
         BINDING_DAP_AGGREGATE_STORE, BINDING_DAP_HELPER_STATE_STORE,
         BINDING_DAP_LEADER_AGG_JOB_QUEUE, BINDING_DAP_LEADER_BATCH_QUEUE,
         BINDING_DAP_LEADER_COL_JOB_QUEUE, BINDING_DAP_REPORTS_PENDING,
@@ -50,16 +52,20 @@ use daphne::{
     hpke::{HpkeConfig, HpkeDecrypter},
     messages::{
         BatchId, BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeCiphertext,
-        PartialBatchSelector, Report, ReportId, ReportMetadata, TaskId, TransitionFailure,
+        PartialBatchSelector, Report, TaskId, TransitionFailure,
     },
     metrics::DaphneMetrics,
-    roles::{early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
+    roles::{
+        early_metadata_check, DapAggregator, DapAuthorizedSender, DapHelper, DapLeader,
+        DapReportInitializer,
+    },
+    vdaf::{EarlyReportState, EarlyReportStateConsumed, EarlyReportStateInitialized},
     DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
     DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapSender, DapTaskConfig, DapVersion,
     MetaAggregationJobId,
 };
 use futures::future::try_join_all;
-use prio::codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode};
+use prio::codec::{Encode, ParameterizedDecode, ParameterizedEncode};
 use std::{borrow::Cow, collections::HashMap};
 use tracing::debug;
 use worker::*;
@@ -214,6 +220,161 @@ impl DapAuthorizedSender<DaphneWorkerAuth> for DaphneWorker<'_> {
     }
 }
 
+#[async_trait(?Send)]
+impl DapReportInitializer for DaphneWorker<'_> {
+    async fn initialize_reports<'req>(
+        &self,
+        is_leader: bool,
+        task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        part_batch_sel: &PartialBatchSelector,
+        consumed_reports: Vec<EarlyReportStateConsumed<'req>>,
+    ) -> std::result::Result<Vec<EarlyReportStateInitialized<'req>>, DapError> {
+        let current_time = self.get_current_time();
+        let min_time = self.least_valid_report_time(current_time);
+        let max_time = self.greatest_valid_report_time(current_time);
+        let durable = self.durable().with_retry();
+        let task_id_hex = task_id.to_hex();
+        let span = task_config
+            .as_ref()
+            .batch_span_for_meta(part_batch_sel, consumed_reports.iter())?;
+
+        // Coalesce reports pertaining to the same ReportsProcessed or AggregateStore instance.
+        let mut reports_processed_request_data: HashMap<String, ReportsProcessedReq> =
+            HashMap::new();
+        let mut agg_store_request_name = Vec::new();
+        let mut agg_store_request_bucket = Vec::new();
+        for (bucket, consumed_reports_per_bucket) in span.iter() {
+            agg_store_request_name.push(durable_name_agg_store(
+                &task_config.version,
+                &task_id_hex,
+                bucket,
+            ));
+            agg_store_request_bucket.push(bucket);
+            for consumed_report in consumed_reports_per_bucket.iter() {
+                let durable_name = self.config().durable_name_report_store(
+                    task_config.as_ref(),
+                    &task_id_hex,
+                    consumed_report.metadata(),
+                );
+
+                reports_processed_request_data
+                    .entry(durable_name)
+                    .or_insert(ReportsProcessedReq {
+                        is_leader,
+                        vdaf_verify_key: task_config.vdaf_verify_key.clone(),
+                        vdaf_config: task_config.vdaf.clone(),
+                        consumed_reports: Vec::default(),
+                    })
+                    .consumed_reports
+                    .push((*consumed_report).clone());
+            }
+        }
+
+        // Send ReportsProcessed requests.
+        let mut reports_processed_requests = Vec::new();
+        for (durable_name, consumed_reports) in reports_processed_request_data.into_iter() {
+            reports_processed_requests.push(durable.post_with_handler(
+                BINDING_DAP_REPORTS_PROCESSED,
+                DURABLE_REPORTS_PROCESSED_INITIALIZE,
+                durable_name,
+                consumed_reports,
+                |output: ReportsProcessedResp, retried| (output, retried),
+            ));
+        }
+        let reports_processed_responses: Vec<(ReportsProcessedResp, bool)> =
+            try_join_all(reports_processed_requests)
+                .await
+                .map_err(|e| fatal_error!(err = e))?;
+
+        // Flatten the responses from ReportsProcessed into a hash map.
+        let mut initialized_reports = HashMap::new();
+        for (reports_processed_response, retried) in reports_processed_responses.into_iter() {
+            for mut initialized_report in reports_processed_response.initialized_reports.into_iter()
+            {
+                if let EarlyReportStateInitialized::Rejected {
+                    ref mut failure, ..
+                } = initialized_report
+                {
+                    // The request to ReportsProcessed is not idempotent. If it was retried, then
+                    // it's possible that one or more of the reports were erroneously marked as
+                    // replayed. This would happen if, for example, a fresh report was marked as
+                    // aggregated in the failed request, then marked as replayed by a retried
+                    // request. In this case we cannot determine definitively if the report was
+                    // replayed, so we drop it to be safe.
+                    if retried && matches!(failure, TransitionFailure::ReportReplayed) {
+                        *failure = TransitionFailure::ReportDropped;
+                    }
+                }
+                initialized_reports
+                    .insert(initialized_report.metadata().id.clone(), initialized_report);
+            }
+        }
+
+        // Send AggregateStore requests.
+        let mut agg_store_requests = Vec::new();
+        for durable_name in agg_store_request_name {
+            agg_store_requests.push(durable.get(
+                BINDING_DAP_AGGREGATE_STORE,
+                DURABLE_AGGREGATE_STORE_CHECK_COLLECTED,
+                durable_name,
+            ));
+        }
+        let agg_store_responses: Vec<bool> = try_join_all(agg_store_requests)
+            .await
+            .map_err(|e| fatal_error!(err = e))?;
+
+        // Reject reports that have been collected.
+        for (bucket, collected) in agg_store_request_bucket
+            .iter()
+            .zip(agg_store_responses.into_iter())
+        {
+            for metadata in span
+                .get(bucket)
+                .unwrap()
+                .iter()
+                .map(|report| report.metadata())
+            {
+                if let Some(initialized_report) = initialized_reports.get_mut(&metadata.id) {
+                    let processed = match initialized_report {
+                        EarlyReportStateInitialized::Ready { .. } => false,
+                        EarlyReportStateInitialized::Rejected { failure, .. }
+                            if matches!(failure, TransitionFailure::ReportReplayed) =>
+                        {
+                            true
+                        }
+                        EarlyReportStateInitialized::Rejected { .. } => {
+                            continue;
+                        }
+                    };
+
+                    if let Some(failure) =
+                        early_metadata_check(metadata, processed, collected, min_time, max_time)
+                    {
+                        *initialized_report = EarlyReportStateInitialized::Rejected {
+                            metadata: Cow::Owned(metadata.clone()),
+                            failure,
+                        };
+                    }
+                }
+            }
+        }
+
+        Ok(consumed_reports
+            .iter()
+            .map(|report| {
+                initialized_reports
+                    .remove(&report.metadata().id)
+                    .ok_or_else(|| {
+                        fatal_error!(
+                            err = "Response from ReportsProcessed does not match the request"
+                        )
+                    })
+            })
+            .collect::<std::result::Result<Vec<_>, DapError>>()?)
+    }
+}
+
 #[async_trait(?Send)]
 impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
     type WrappedDapTaskConfig<'a> = DapTaskConfigKvPair<'a>;
@@ -501,126 +662,6 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
         Ok(agg_share)
     }
 
-    async fn check_early_reject(
-        &self,
-        task_id: &TaskId,
-        part_batch_sel: &PartialBatchSelector,
-        report_meta: impl Iterator<Item = &ReportMetadata>,
-    ) -> std::result::Result<HashMap<ReportId, TransitionFailure>, DapError> {
-        let durable = self.durable().with_retry();
-        let task_config = self.try_get_task_config(task_id).await?;
-        let task_id_hex = task_id.to_hex();
-        let span = task_config
-            .as_ref()
-            .batch_span_for_meta(part_batch_sel, report_meta)?;
-
-        // Coalesce reports pertaining to the same ReportsProcessed or AggregateStore instance.
-        let mut reports_processed_request_data: HashMap<String, Vec<String>> = HashMap::new();
-        let mut agg_store_request_name = Vec::new();
-        let mut agg_store_request_bucket = Vec::new();
-        for (bucket, report_meta) in span.iter() {
-            agg_store_request_name.push(durable_name_agg_store(
-                &task_config.as_ref().version,
-                &task_id_hex,
-                bucket,
-            ));
-            agg_store_request_bucket.push(bucket);
-            for metadata in report_meta {
-                let durable_name = self.config().durable_name_report_store(
-                    task_config.as_ref(),
-                    &task_id_hex,
-                    metadata,
-                );
-                let report_id_hex = hex::encode(metadata.id.get_encoded());
-                let report_id_hex_set = reports_processed_request_data
-                    .entry(durable_name)
-                    .or_default();
-                report_id_hex_set.push(report_id_hex);
-            }
-        }
-
-        // Send ReportsProcessed requests.
-        let mut reports_processed_requests = Vec::new();
-        for (durable_name, report_id_hex_set) in reports_processed_request_data.into_iter() {
-            reports_processed_requests.push(durable.post_with_handler(
-                BINDING_DAP_REPORTS_PROCESSED,
-                DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED,
-                durable_name,
-                report_id_hex_set,
-                |output: Vec<String>, retried| (output, retried),
-            ));
-        }
-
-        // Send AggregateStore requests.
-        let mut agg_store_requests = Vec::new();
-        for durable_name in agg_store_request_name {
-            agg_store_requests.push(durable.get(
-                BINDING_DAP_AGGREGATE_STORE,
-                DURABLE_AGGREGATE_STORE_CHECK_COLLECTED,
-                durable_name,
-            ));
-        }
-
-        // Create the set of reports that have been processed.
-        let reports_processed_responses: Vec<(Vec<String>, bool)> =
-            try_join_all(reports_processed_requests)
-                .await
-                .map_err(|e| fatal_error!(err = e))?;
-        let mut early_fails = HashMap::new();
-        for (reports_processed, retried) in reports_processed_responses.into_iter() {
-            // The request to ReportsProcessed is not idempotent. If it was retried, then it's
-            // possible that one or more of the reports were erroneously as replayed. This would
-            // happen if, for example, a fresh report was marked as aggregated in the failed
-            // request, then marked as replayed by a retried request. In this case we cannot
-            // determine definitively if the report was replayed, so we drop it to be safe.
-            let failure = if retried {
-                TransitionFailure::ReportDropped
-            } else {
-                TransitionFailure::ReportReplayed
-            };
-            for report_id_hex in reports_processed.into_iter() {
-                let report_id = ReportId::get_decoded(
-                    &hex::decode(report_id_hex)
-                        .map_err(|e| DapAbort::from_hex_error(e, task_id.clone()))?,
-                )
-                .map_err(|e| DapAbort::from_codec_error(e, task_id.clone()))?;
-                early_fails.insert(report_id, failure);
-            }
-        }
-
-        let agg_store_responses: Vec<bool> = try_join_all(agg_store_requests)
-            .await
-            .map_err(|e| fatal_error!(err = e))?;
-
-        // Decide which reports to reject early. A report will be rejected here if, for example,
-        // it has been processed but not collected, or if it has not been proceessed but pertains
-        // to a batch that was previously collected, or if it is not within time bounds specified
-        // by the configuration.
-        let current_time = self.get_current_time();
-        let min_time = self.least_valid_report_time(current_time);
-        let max_time = self.greatest_valid_report_time(current_time);
-        for (bucket, collected) in agg_store_request_bucket
-            .iter()
-            .zip(agg_store_responses.into_iter())
-        {
-            for metadata in span.get(bucket).unwrap() {
-                let result = early_fails.get(&metadata.id);
-                if matches!(result, Some(TransitionFailure::ReportDropped)) {
-                    continue;
-                }
-
-                let processed = matches!(result, Some(TransitionFailure::ReportReplayed));
-                if let Some(failure) =
-                    early_metadata_check(metadata, processed, collected, min_time, max_time)
-                {
-                    early_fails.insert(metadata.id.clone(), failure);
-                }
-            }
-        }
-
-        Ok(early_fails)
-    }
-
     async fn mark_collected(
         &self,
         task_id: &TaskId,
diff --git a/daphne_worker/src/durable/reports_processed.rs b/daphne_worker/src/durable/reports_processed.rs
index 459d440..f6061fd 100644
--- a/daphne_worker/src/durable/reports_processed.rs
+++ b/daphne_worker/src/durable/reports_processed.rs
@@ -6,13 +6,23 @@ use crate::{
     durable::{state_set_if_not_exists, BINDING_DAP_REPORTS_PROCESSED},
     initialize_tracing, int_err,
 };
+use daphne::{
+    messages::{ReportId, ReportMetadata, TransitionFailure},
+    vdaf::{
+        EarlyReportState, EarlyReportStateConsumed, EarlyReportStateInitialized, VdafPrepMessage,
+        VdafPrepState, VdafVerifyKey,
+    },
+    DapError, VdafConfig,
+};
 use futures::future::try_join_all;
-use std::time::Duration;
+use prio::codec::{CodecError, Decode, ParameterizedDecode};
+use serde::{Deserialize, Serialize};
+use std::{borrow::Cow, collections::HashSet, time::Duration};
 use tracing::warn;
 use worker::*;
 
-pub(crate) const DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED: &str =
-    "/internal/do/report_store/mark_aggregated";
+pub(crate) const DURABLE_REPORTS_PROCESSED_INITIALIZE: &str =
+    "/internal/do/reports_processed/initialize";
 
 /// Durable Object (DO) for tracking which reports have been processed.
 ///
@@ -77,22 +87,71 @@ impl DurableObject for ReportsProcessed {
         );
 
         match (req.path().as_ref(), req.method()) {
-            // Mark a set of reports as aggregated. Return the set of report IDs that already
-            // exist.
+            // Initialize a report:
+            //  * Ensure the report wasn't replayed
+            //  * Ensure the report won't be included in a batch that was already collected
+            //  * Initialize VDAF preparation.
             //
             // Non-idempotent
-            // Input: `report_id_hex_set: Vec<String>` (hex-encoded report IDs)
-            // Output: `Vec<String>` (subset of the inputs that already exist).
-            (DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED, Method::Post) => {
-                let report_id_hex_set: Vec<String> = req.json().await?;
-                let mut requests = Vec::new();
-                for report_id_hex in report_id_hex_set.into_iter() {
-                    requests.push(self.to_checked(report_id_hex));
-                }
+            // Input: `ReportsProcessedReq`
+            // Output: `ReportsProcessedResp`
+            (DURABLE_REPORTS_PROCESSED_INITIALIZE, Method::Post) => {
+                let reports_processed_request: ReportsProcessedReq = req.json().await?;
+                let replayed_reports = try_join_all(
+                    reports_processed_request
+                        .consumed_reports
+                        .iter()
+                        .filter(|consumed_report| consumed_report.is_ready())
+                        .map(|consumed_report| {
+                            self.to_checked(consumed_report.metadata().id.to_hex())
+                        }),
+                )
+                .await?
+                .into_iter()
+                // The iterator here is on `Option<String>`. What we want to do is filter out
+                // values that are `None` and unwrap the `String`.
+                .flatten()
+                .map(|report_id_hex| {
+                    hex::decode(report_id_hex)
+                        .map_err(|e| format!("Failed to hex decode ReportId: {e}"))
+                        .and_then(|report_id_data| {
+                            ReportId::get_decoded(&report_id_data)
+                                .map_err(|e| format!("Failed to decode ReportId: {e}"))
+                        })
+                })
+                .collect::<std::result::Result<HashSet<ReportId>, String>>()
+                .map_err(|e| int_err(format!("ReportsProcessed: {e}")))?;
+
+                let initialized_reports = reports_processed_request
+                    .consumed_reports
+                    .into_iter()
+                    .map(|consumed_report| {
+                        if replayed_reports.contains(&consumed_report.metadata().id) {
+                            Ok(EarlyReportStateInitialized::Rejected {
+                                metadata: Cow::Owned(consumed_report.metadata().clone()),
+                                failure: TransitionFailure::ReportReplayed,
+                            })
+                        } else {
+                            EarlyReportStateInitialized::initialize(
+                                reports_processed_request.is_leader,
+                                &reports_processed_request.vdaf_verify_key,
+                                &reports_processed_request.vdaf_config,
+                                consumed_report,
+                            )
+                        }
+                    })
+                    .collect::<std::result::Result<Vec<EarlyReportStateInitialized>, DapError>>()
+                    .map_err(|e| {
+                        int_err(format!(
+                            "ReportsProcessed: failed to initialize a report: {e}"
+                        ))
+                    })?;
 
-                let responses: Vec<Option<String>> = try_join_all(requests).await?;
-                let res: Vec<String> = responses.into_iter().flatten().collect();
-                Response::from_json(&res)
+                Response::from_json(&ReportsProcessedResp {
+                    is_leader: reports_processed_request.is_leader,
+                    vdaf_config: reports_processed_request.vdaf_config,
+                    initialized_reports,
+                })
             }
 
             _ => Err(int_err(format!(
@@ -110,3 +169,89 @@ impl DurableObject for ReportsProcessed {
         Response::from_json(&())
     }
 }
+
+#[derive(Serialize, Deserialize)]
+pub(crate) struct ReportsProcessedReq<'req> {
+    pub(crate) is_leader: bool,
+    pub(crate) vdaf_verify_key: VdafVerifyKey,
+    pub(crate) vdaf_config: VdafConfig,
+    pub(crate) consumed_reports: Vec<EarlyReportStateConsumed<'req>>,
+}
+
+#[derive(Serialize, Deserialize)]
+#[serde(try_from = "ShadowReportsProcessedResp")]
+pub(crate) struct ReportsProcessedResp<'req> {
+    pub(crate) is_leader: bool,
+    pub(crate) vdaf_config: VdafConfig,
+    pub(crate) initialized_reports: Vec<EarlyReportStateInitialized<'req>>,
+}
+
+// we need this custom deserializer because VdafPrepState and VdafPrepMessage don't implement
+// Decode, only ParameterizedDecode.
+#[derive(Deserialize)]
+struct ShadowReportsProcessedResp {
+    pub(crate) is_leader: bool,
+    pub(crate) vdaf_config: VdafConfig,
+    pub(crate) initialized_reports: Vec<EarlyReportStateInitializedOwned>,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "snake_case")]
+pub enum EarlyReportStateInitializedOwned {
+    Ready {
+        metadata: ReportMetadata,
+        #[serde(with = "hex")]
+        public_share: Vec<u8>,
+        #[serde(with = "hex")]
+        state: Vec<u8>,
+        #[serde(with = "hex")]
+        message: Vec<u8>,
+    },
+    Rejected {
+        metadata: ReportMetadata,
+        failure: TransitionFailure,
+    },
+}
+
+impl TryFrom<ShadowReportsProcessedResp> for ReportsProcessedResp<'_> {
+    type Error = CodecError;
+
+    fn try_from(other: ShadowReportsProcessedResp) -> std::result::Result<Self, CodecError> {
+        let initialized_reports = other
+            .initialized_reports
+            .into_iter()
+            .map(|initialized_report| match initialized_report {
+                EarlyReportStateInitializedOwned::Ready {
+                    metadata,
+                    public_share,
+                    state,
+                    message,
+                } => {
+                    let state = VdafPrepState::get_decoded_with_param(
+                        &(&other.vdaf_config, other.is_leader),
+                        &state,
+                    )?;
+                    let message = VdafPrepMessage::get_decoded_with_param(&state, &message)?;
+
+                    Ok(EarlyReportStateInitialized::Ready {
+                        metadata: Cow::Owned(metadata),
+                        public_share: Cow::Owned(public_share),
+                        state,
+                        message,
+                    })
+                }
+                EarlyReportStateInitializedOwned::Rejected { metadata, failure } => {
+                    Ok(EarlyReportStateInitialized::Rejected {
+                        metadata: Cow::Owned(metadata),
+                        failure,
+                    })
+                }
+            })
+            .collect::<std::result::Result<Vec<EarlyReportStateInitialized>, CodecError>>()?;
+        Ok(Self {
+            is_leader: other.is_leader,
+            vdaf_config: other.vdaf_config,
+            initialized_reports,
+        })
+    }
+}


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 39ebb4a..9895666 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -16,7 +16,8 @@ use crate::{
         ReportMetadata, TaskId, Time, TransitionFailure,
     },
     metrics::DaphneMetrics,
-    roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
+    roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader, DapReportInitializer},
+    vdaf::{EarlyReportState, EarlyReportStateConsumed, EarlyReportStateInitialized},
     DapAbort, DapAggregateResult, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError,
     DapGlobalConfig, DapHelperState, DapHelperTransition, DapLeaderState, DapLeaderTransition,
     DapLeaderUncommitted, DapMeasurement, DapOutputShare, DapQueryConfig, DapRequest, DapResponse,
@@ -355,6 +356,66 @@ impl DapAuthorizedSender<BearerToken> for MockAggregator {
     }
 }
 
+#[async_trait(?Send)]
+impl DapReportInitializer for MockAggregator {
+    async fn initialize_reports<'req>(
+        &self,
+        is_leader: bool,
+        task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        part_batch_sel: &PartialBatchSelector,
+        consumed_reports: Vec<EarlyReportStateConsumed<'req>>,
+    ) -> Result<Vec<EarlyReportStateInitialized<'req>>, DapError> {
+        let span = task_config.batch_span_for_meta(
+            part_batch_sel,
+            consumed_reports.iter().filter(|report| report.is_ready()),
+        )?;
+
+        let mut early_fails = HashMap::new();
+        for (bucket, reports_consumed_per_bucket) in span.iter() {
+            for metadata in reports_consumed_per_bucket
+                .iter()
+                .map(|report| report.metadata())
+            {
+                // Check whether Report has been collected or replayed.
+                if let Some(transition_failure) = self
+                    .check_report_early_fail(task_id, &bucket.to_owned_bucket(), metadata)
+                    .await
+                {
+                    early_fails.insert(metadata.id.clone(), transition_failure);
+                };
+
+                // Mark report processed.
+                let mut guard = self
+                    .report_store
+                    .lock()
+                    .expect("report_store: failed to lock");
+                let report_store = guard.entry(task_id.clone()).or_default();
+                report_store.processed.insert(metadata.id.clone());
+            }
+        }
+
+        Ok(consumed_reports
+            .into_iter()
+            .map(|consumed| {
+                if let Some(failure) = early_fails.get(&consumed.metadata().id) {
+                    Ok(EarlyReportStateInitialized::Rejected {
+                        metadata: Cow::Owned(consumed.metadata().clone()),
+                        failure: *failure,
+                    })
+                } else {
+                    EarlyReportStateInitialized::initialize(
+                        is_leader,
+                        &task_config.vdaf_verify_key,
+                        &task_config.vdaf,
+                        consumed,
+                    )
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?)
+    }
+}
+
 #[async_trait(?Send)]
 impl DapAggregator<BearerToken> for MockAggregator {
     // The lifetimes on the traits ensure that we can return a reference to a task config stored by
@@ -508,42 +569,6 @@ impl DapAggregator<BearerToken> for MockAggregator {
         Ok(agg_share)
     }
 
-    async fn check_early_reject(
-        &self,
-        task_id: &TaskId,
-        part_batch_sel: &PartialBatchSelector,
-        report_meta: impl Iterator<Item = &ReportMetadata>,
-    ) -> Result<HashMap<ReportId, TransitionFailure>, DapError> {
-        let task_config = self
-            .get_task_config_for(Cow::Borrowed(task_id))
-            .await
-            .unwrap()
-            .expect("tasks: unrecognized task");
-        let span = task_config.batch_span_for_meta(part_batch_sel, report_meta)?;
-        let mut early_fails = HashMap::new();
-        for (bucket, report_meta) in span.iter() {
-            for metadata in report_meta.iter() {
-                // Check whether Report has been collected or replayed.
-                if let Some(transition_failure) = self
-                    .check_report_early_fail(task_id, &bucket.to_owned_bucket(), metadata)
-                    .await
-                {
-                    early_fails.insert(metadata.id.clone(), transition_failure);
-                };
-
-                // Mark report processed.
-                let mut guard = self
-                    .report_store
-                    .lock()
-                    .expect("report_store: failed to lock");
-                let report_store = guard.entry(task_id.clone()).or_default();
-                report_store.processed.insert(metadata.id.clone());
-            }
-        }
-
-        Ok(early_fails)
-    }
-
     async fn mark_collected(
         &self,
         task_id: &TaskId,
@@ -1086,6 +1111,49 @@ pub struct AggregationJobTest {
     pub(crate) prometheus_registry: prometheus::Registry,
     pub(crate) leader_metrics: DaphneMetrics,
     pub(crate) helper_metrics: DaphneMetrics,
+    pub(crate) leader_reports_processed: Arc<Mutex<HashSet<ReportId>>>,
+    pub(crate) helper_reports_processed: Arc<Mutex<HashSet<ReportId>>>,
+}
+
+// NOTE(cjpatton) This implementation of the report initializer is not feature complete. Since
+// [`AggrregationJobTest`], is only used to test the aggregation flow, features that are not
+// directly relevant to the tests aren't implemented.
+#[async_trait(?Send)]
+impl DapReportInitializer for AggregationJobTest {
+    async fn initialize_reports<'req>(
+        &self,
+        is_leader: bool,
+        _task_id: &TaskId,
+        task_config: &DapTaskConfig,
+        _part_batch_sel: &PartialBatchSelector,
+        consumed_reports: Vec<EarlyReportStateConsumed<'req>>,
+    ) -> Result<Vec<EarlyReportStateInitialized<'req>>, DapError> {
+        let mut reports_processed = if is_leader {
+            self.leader_reports_processed.lock().unwrap()
+        } else {
+            self.helper_reports_processed.lock().unwrap()
+        };
+
+        Ok(consumed_reports
+            .into_iter()
+            .map(|consumed| {
+                if reports_processed.contains(&consumed.metadata().id) {
+                    Ok(EarlyReportStateInitialized::Rejected {
+                        metadata: Cow::Owned(consumed.metadata().clone()),
+                        failure: TransitionFailure::ReportReplayed,
+                    })
+                } else {
+                    reports_processed.insert(consumed.metadata().id.clone());
+                    EarlyReportStateInitialized::initialize(
+                        is_leader,
+                        &task_config.vdaf_verify_key,
+                        &task_config.vdaf,
+                        consumed,
+                    )
+                }
+            })
+            .collect::<Result<Vec<_>, _>>()?)
+    }
 }
 
 impl AggregationJobTest {
@@ -1136,6 +1204,8 @@ impl AggregationJobTest {
             prometheus_registry,
             leader_metrics,
             helper_metrics,
+            leader_reports_processed: Default::default(),
+            helper_reports_processed: Default::default(),
         }
     }
 
@@ -1176,6 +1246,7 @@ impl AggregationJobTest {
             .vdaf
             .produce_agg_job_init_req(
                 &self.leader_hpke_receiver_config,
+                self,
                 &self.task_id,
                 &self.task_config,
                 &self.agg_job_id,
@@ -1201,6 +1272,7 @@ impl AggregationJobTest {
             .vdaf
             .handle_agg_job_init_req(
                 &self.helper_hpke_receiver_config,
+                self,
                 &self.task_id,
                 &self.task_config,
                 agg_job_init_req,
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index eb80591..231f988 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -28,7 +28,7 @@ use prio::{
     },
 };
 use rand::prelude::*;
-use std::fmt::Debug;
+use std::{borrow::Cow, fmt::Debug};
 
 use super::{EarlyReportStateConsumed, EarlyReportStateInitialized};
 
@@ -94,14 +94,14 @@ async fn roundtrip_report(version: DapVersion) {
         true, // is_leader
         &t.task_id,
         &t.task_config,
-        &report.report_metadata,
-        &report.public_share,
+        Cow::Borrowed(&report.report_metadata),
+        Cow::Borrowed(&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 {
+    let EarlyReportStateInitialized::Ready{ state: leader_step, message: leader_share, .. } =
+        EarlyReportStateInitialized::initialize(true, &t.task_config.vdaf_verify_key, &t.task_config.vdaf, early_report_state_consumed).unwrap() else {
         panic!("rejected unexpectedly");
     };
 
@@ -110,14 +110,14 @@ async fn roundtrip_report(version: DapVersion) {
         false, // is_helper
         &t.task_id,
         &t.task_config,
-        &report.report_metadata,
-        &report.public_share,
+        Cow::Borrowed(&report.report_metadata),
+        Cow::Borrowed(&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 {
+    let EarlyReportStateInitialized::Ready{ state: helper_step, message: helper_share, .. } =
+        EarlyReportStateInitialized::initialize(false, &t.task_config.vdaf_verify_key, &t.task_config.vdaf, early_report_state_consumed).unwrap() else {
         panic!("rejected unexpectedly");
     };