Repository: cloudflare/daphne
Pull request #374 (2023-07-21 19:04:56)

Wait to mark reports as aggregated until just before committing
Loads that require a lot of CPU time can trigger the Workers runtime to reset itself, causing a 500 error. This is most likely to happen when processing the AggregationJobInitReq, as this involves the expensive, VDAF prep initialization step. Some reports may still get marked as aggregated; if the peer retries the request, those reports will get rejected.

This problem is so common that we need to make this request idempotent. As a first step, wait to mark the reports as aggregated until just before we have committed to aggregating them. In particular, instead of doing this in `DapReportInitializer::initialize_reports()`, wait until `DapAggregator::put_out_shares()`.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index be7a37c..8791724 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -310,20 +310,19 @@ impl DapTaskConfig {
         self.quantized_time_lower_bound(time) + self.time_precision
     }
 
-    /// Compute the "batch span" of a set of output shares and, for each buckent in the span,
-    /// aggregate the output shares into an aggregate share.
+    /// Compute the "batch span" of a set of output shares.
     pub fn batch_span_for_out_shares<'a>(
         &self,
         part_batch_sel: &'a PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
-    ) -> Result<HashMap<DapBatchBucket<'a>, DapAggregateShare>, DapError> {
+    ) -> Result<HashMap<DapBatchBucket<'a>, Vec<DapOutputShare>>, DapError> {
         if !self.query.is_valid_part_batch_sel(part_batch_sel) {
             return Err(fatal_error!(
                 err = "partial batch selector not compatible with task",
             ));
         }
 
-        let mut span: HashMap<DapBatchBucket<'a>, DapAggregateShare> = HashMap::new();
+        let mut span: HashMap<DapBatchBucket<'a>, Vec<DapOutputShare>> = HashMap::new();
         for out_share in out_shares.into_iter() {
             let bucket = match part_batch_sel {
                 PartialBatchSelector::TimeInterval => DapBatchBucket::TimeInterval {
@@ -334,14 +333,7 @@ impl DapTaskConfig {
                 }
             };
 
-            let agg_share = span.entry(bucket).or_default();
-            agg_share.merge(DapAggregateShare {
-                report_count: 1,
-                min_time: out_share.time,
-                max_time: out_share.time,
-                checksum: out_share.checksum,
-                data: Some(out_share.data),
-            })?;
+            span.entry(bucket).or_default().push(out_share);
         }
 
         Ok(span)
@@ -515,7 +507,8 @@ impl DapHelperState {
 #[derive(Debug)]
 /// An ouptut share produced by an Aggregator for a single report.
 pub struct DapOutputShare {
-    pub(crate) time: u64, // Value from the report
+    pub report_id: ReportId, // Value from report
+    pub time: u64,           // Value from the report
     pub(crate) checksum: [u8; 32],
     pub(crate) data: VdafAggregateShare,
 }
diff --git a/daphne/src/roles/aggregator.rs b/daphne/src/roles/aggregator.rs
index fc9cdcd..e367f8c 100644
--- a/daphne/src/roles/aggregator.rs
+++ b/daphne/src/roles/aggregator.rs
@@ -1,7 +1,7 @@
 // Copyright (c) 2023 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
-use std::borrow::Cow;
+use std::{borrow::Cow, collections::HashSet};
 
 use async_trait::async_trait;
 use prio::codec::Encode;
@@ -12,8 +12,8 @@ use crate::{
     error::DapAbort,
     hpke::{HpkeConfig, HpkeDecrypter},
     messages::{
-        decode_base64url, BatchId, BatchSelector, HpkeConfigList, PartialBatchSelector, TaskId,
-        Time,
+        decode_base64url, BatchId, BatchSelector, HpkeConfigList, PartialBatchSelector, ReportId,
+        TaskId, Time,
     },
     metrics::{DaphneMetrics, DaphneRequestType},
     vdaf::{EarlyReportStateConsumed, EarlyReportStateInitialized},
@@ -102,13 +102,17 @@ pub trait DapAggregator<S>: HpkeDecrypter + DapReportInitializer + Sized {
     /// (resp. Helper) in response to a CollectReq (resp. AggregateShareReq) for fixed-size tasks.
     async fn batch_exists(&self, task_id: &TaskId, batch_id: &BatchId) -> Result<bool, DapError>;
 
-    /// Store a set of output shares.
+    /// Store a set of output shares and mark the corresponding reports as aggregated. Any reports
+    /// that were already aggregated are not committed.
+    ///
+    /// TODO spec: Ensure the spec allows rejecting due to replay at this stage.
     async fn put_out_shares(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
-    ) -> Result<(), DapError>;
+    ) -> Result<HashSet<ReportId>, DapError>;
 
     /// Fetch the aggregate share for the given batch.
     async fn get_agg_share(
diff --git a/daphne/src/roles/helper.rs b/daphne/src/roles/helper.rs
index c5d5afc..5e51a85 100644
--- a/daphne/src/roles/helper.rs
+++ b/daphne/src/roles/helper.rs
@@ -16,7 +16,8 @@ use crate::{
     fatal_error,
     messages::{
         constant_time_eq, AggregateShare, AggregateShareReq, AggregationJobContinueReq,
-        AggregationJobInitReq, PartialBatchSelector, ReportMetadata, TaskId,
+        AggregationJobInitReq, PartialBatchSelector, ReportMetadata, TaskId, TransitionFailure,
+        TransitionVar,
     },
     metrics::DaphneRequestType,
     DapError, DapHelperState, DapHelperTransition, DapRequest, DapResource, DapResponse,
@@ -245,10 +246,24 @@ pub trait DapHelper<S>: DapAggregator<S> {
                     DapHelperTransition::Continue(..) => {
                         return Err(fatal_error!(err = "unexpected transition (continued)").into());
                     }
-                    DapHelperTransition::Finish(out_shares, agg_job_resp) => {
+                    DapHelperTransition::Finish(out_shares, mut agg_job_resp) => {
                         let out_shares_count = u64::try_from(out_shares.len()).unwrap();
-                        self.put_out_shares(task_id, &part_batch_sel, out_shares)
+                        let replayed = self
+                            .put_out_shares(task_id, task_config, &part_batch_sel, out_shares)
                             .await?;
+
+                        // If there are multiple aggregation jobs in flight that contain the
+                        // same report, then we may need to reject the report at this late stage.
+                        if !replayed.is_empty() {
+                            for transition in agg_job_resp.transitions.iter_mut() {
+                                if replayed.contains(&transition.report_id) {
+                                    let failure = TransitionFailure::ReportReplayed;
+                                    transition.var = TransitionVar::Failed(failure);
+                                    metrics.report_inc_by(&format!("rejected_{failure}"), 1);
+                                }
+                            }
+                        }
+
                         (agg_job_resp, out_shares_count)
                     }
                 };
diff --git a/daphne/src/roles/leader.rs b/daphne/src/roles/leader.rs
index c153d1d..e524ebf 100644
--- a/daphne/src/roles/leader.rs
+++ b/daphne/src/roles/leader.rs
@@ -392,7 +392,13 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
                 .vdaf
                 .handle_final_agg_job_resp(uncommited, agg_job_resp, &metrics)?;
         let out_shares_count = out_shares.len() as u64;
-        self.put_out_shares(task_id, part_batch_sel, out_shares)
+
+        // At this point we're committed to aggregating the reports: if we do detect a report was
+        // replayed at this stage, then we may end up with a batch mismatch. However, this should
+        // only happen if there are multiple aggregation jobs in-flight that include the same
+        // report.
+        let _ = self
+            .put_out_shares(task_id, task_config, part_batch_sel, out_shares)
             .await?;
 
         metrics.report_inc_by("aggregated", out_shares_count);
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 858631a..58f759e 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -913,6 +913,7 @@ impl VdafConfig {
 
                     states.push((
                         DapOutputShare {
+                            report_id: leader_report_id.clone(),
                             time: leader_time,
                             checksum: checksum.as_ref().try_into().unwrap(),
                             data,
@@ -1063,6 +1064,7 @@ impl VdafConfig {
                         );
 
                         out_shares.push(DapOutputShare {
+                            report_id: helper_report_id.clone(),
                             time: helper_time,
                             checksum: checksum.as_ref().try_into().unwrap(),
                             data,
diff --git a/daphne_worker/src/config.rs b/daphne_worker/src/config.rs
index 1b8eabc..3613da6 100644
--- a/daphne_worker/src/config.rs
+++ b/daphne_worker/src/config.rs
@@ -26,7 +26,7 @@ use daphne::{
     fatal_error,
     hpke::{HpkeConfig, HpkeReceiverConfig},
     messages::{
-        decode_base64url_vec, AggregationJobId, BatchId, CollectionJobId, ReportMetadata, TaskId,
+        decode_base64url_vec, AggregationJobId, BatchId, CollectionJobId, ReportId, TaskId, Time,
     },
     DapError, DapGlobalConfig, DapQueryConfig, DapRequest, DapResource, DapResponse, DapTaskConfig,
     DapVersion, Prio3Config, VdafConfig,
@@ -345,17 +345,14 @@ impl DaphneWorkerConfig {
         &self,
         task_config: &DapTaskConfig,
         task_id_hex: &str,
-        metadata: &ReportMetadata,
+        report_id: &ReportId,
+        report_time: Time,
     ) -> String {
         let mut shard_seed = [0; 8];
-        PrgSha3::seed_stream(
-            &self.report_shard_key,
-            b"report shard",
-            metadata.id.as_ref(),
-        )
-        .fill(&mut shard_seed);
+        PrgSha3::seed_stream(&self.report_shard_key, b"report shard", report_id.as_ref())
+            .fill(&mut shard_seed);
         let shard = u64::from_be_bytes(shard_seed) % self.report_shard_count;
-        let epoch = metadata.time - (metadata.time % self.global.report_storage_epoch_duration);
+        let epoch = report_time - (report_time % self.global.report_storage_epoch_duration);
         durable_name_report_store(&task_config.version, task_id_hex, epoch, shard)
     }
 }
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 8357f29..09eb778 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -34,6 +34,7 @@ use crate::{
         },
         reports_processed::{
             ReportsProcessedReq, ReportsProcessedResp, DURABLE_REPORTS_PROCESSED_INITIALIZE,
+            DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED,
         },
         BINDING_DAP_AGGREGATE_STORE, BINDING_DAP_HELPER_STATE_STORE,
         BINDING_DAP_LEADER_AGG_JOB_QUEUE, BINDING_DAP_LEADER_BATCH_QUEUE,
@@ -52,7 +53,7 @@ use daphne::{
     hpke::{HpkeConfig, HpkeDecrypter},
     messages::{
         BatchId, BatchSelector, Collection, CollectionJobId, CollectionReq, HpkeCiphertext,
-        PartialBatchSelector, Report, TaskId, TransitionFailure,
+        PartialBatchSelector, Report, ReportId, TaskId, TransitionFailure,
     },
     metrics::DaphneMetrics,
     roles::{
@@ -66,7 +67,10 @@ use daphne::{
 };
 use futures::future::try_join_all;
 use prio::codec::{Encode, ParameterizedDecode, ParameterizedEncode};
-use std::{borrow::Cow, collections::HashMap};
+use std::{
+    borrow::Cow,
+    collections::{HashMap, HashSet},
+};
 use tracing::debug;
 use worker::*;
 
@@ -255,7 +259,8 @@ impl DapReportInitializer for DaphneWorker<'_> {
                 let durable_name = self.config().durable_name_report_store(
                     task_config.as_ref(),
                     &task_id_hex,
-                    consumed_report.metadata(),
+                    &consumed_report.metadata().id,
+                    consumed_report.metadata().time,
                 );
 
                 reports_processed_request_data
@@ -274,38 +279,22 @@ impl DapReportInitializer for DaphneWorker<'_> {
         // 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(
+            reports_processed_requests.push(durable.post(
                 BINDING_DAP_REPORTS_PROCESSED,
                 DURABLE_REPORTS_PROCESSED_INITIALIZE,
                 durable_name,
                 consumed_reports,
-                |output: ReportsProcessedResp, retried| (output, retried),
             ));
         }
-        let reports_processed_responses: Vec<(ReportsProcessedResp, bool)> =
+        let reports_processed_responses: Vec<ReportsProcessedResp> =
             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;
-                    }
-                }
+        for reports_processed_response in reports_processed_responses.into_iter() {
+            for initialized_report in reports_processed_response.initialized_reports.into_iter() {
                 initialized_reports
                     .insert(initialized_report.metadata().id.clone(), initialized_report);
             }
@@ -607,30 +596,82 @@ impl<'srv> DapAggregator<DaphneWorkerAuth> for DaphneWorker<'srv> {
     async fn put_out_shares(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
-    ) -> std::result::Result<(), DapError> {
-        let task_config = self.try_get_task_config(task_id).await?;
-
+    ) -> std::result::Result<HashSet<ReportId>, DapError> {
+        let task_id_hex = task_id.to_hex();
         let durable = self.durable();
-        let mut requests = Vec::new();
-        for (bucket, agg_share) in task_config
+        let mut agg_store_request_data: HashMap<String, Vec<DapOutputShare>> = HashMap::new();
+        let mut reports_processed_request_data: HashMap<String, Vec<ReportId>> = HashMap::new();
+        for (bucket, out_shares) in task_config
             .as_ref()
             .batch_span_for_out_shares(part_batch_sel, out_shares)?
         {
-            let durable_name =
-                durable_name_agg_store(&task_config.as_ref().version, &task_id.to_hex(), &bucket);
-            requests.push(durable.post::<_, ()>(
-                BINDING_DAP_AGGREGATE_STORE,
-                DURABLE_AGGREGATE_STORE_MERGE,
-                durable_name,
-                agg_share,
-            ));
+            for out_share in out_shares.into_iter() {
+                let reports_processed_name = self.config().durable_name_report_store(
+                    task_config.as_ref(),
+                    &task_id_hex,
+                    &out_share.report_id,
+                    out_share.time,
+                );
+                reports_processed_request_data
+                    .entry(reports_processed_name)
+                    .or_default()
+                    .push(out_share.report_id.clone());
+
+                let agg_store_name =
+                    durable_name_agg_store(&task_config.version, &task_id_hex, &bucket);
+                agg_store_request_data
+                    .entry(agg_store_name)
+                    .or_default()
+                    .push(out_share);
+            }
         }
-        try_join_all(requests)
-            .await
-            .map_err(|e| fatal_error!(err = e))?;
-        Ok(())
+
+        let replayed = try_join_all(reports_processed_request_data.into_iter().map(
+            |(durable_name, report_ids)| async {
+                durable
+                    .post::<_, Vec<ReportId>>(
+                        BINDING_DAP_REPORTS_PROCESSED,
+                        DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED,
+                        durable_name,
+                        report_ids,
+                    )
+                    .await
+            },
+        ))
+        .await
+        .map_err(|e| fatal_error!(err = e))?
+        .into_iter()
+        .flatten()
+        .collect::<HashSet<ReportId>>();
+
+        try_join_all(agg_store_request_data.into_iter().map(
+            |(agg_store_name, out_shares)| async {
+                // Only aggregate the output shares that haven't been replayed.
+                let agg_share = DapAggregateShare::try_from_out_shares(
+                    out_shares
+                        .into_iter()
+                        .filter(|out_share| !replayed.contains(&out_share.report_id)),
+                )?;
+
+                std::result::Result::<_, DapError>::Ok(
+                    durable
+                        .post::<_, ()>(
+                            BINDING_DAP_AGGREGATE_STORE,
+                            DURABLE_AGGREGATE_STORE_MERGE,
+                            agg_store_name,
+                            agg_share,
+                        )
+                        .await,
+                )
+            },
+        ))
+        .await
+        .map_err(|e| fatal_error!(err = e))?;
+
+        Ok(replayed)
     }
 
     async fn get_agg_share(
@@ -726,7 +767,8 @@ impl<'srv> DapLeader<DaphneWorkerAuth> for DaphneWorker<'srv> {
                 self.config().durable_name_report_store(
                     task_config.as_ref(),
                     &task_id_hex,
-                    &report.report_metadata,
+                    &report.report_metadata.id,
+                    report.report_metadata.time,
                 ),
                 &pending_report,
             )
diff --git a/daphne_worker/src/durable/reports_processed.rs b/daphne_worker/src/durable/reports_processed.rs
index f6061fd..9c28eca 100644
--- a/daphne_worker/src/durable/reports_processed.rs
+++ b/daphne_worker/src/durable/reports_processed.rs
@@ -3,7 +3,7 @@
 
 use crate::{
     config::DaphneWorkerConfig,
-    durable::{state_set_if_not_exists, BINDING_DAP_REPORTS_PROCESSED},
+    durable::{state_get, state_set_if_not_exists, BINDING_DAP_REPORTS_PROCESSED},
     initialize_tracing, int_err,
 };
 use daphne::{
@@ -15,7 +15,7 @@ use daphne::{
     DapError, VdafConfig,
 };
 use futures::future::try_join_all;
-use prio::codec::{CodecError, Decode, ParameterizedDecode};
+use prio::codec::{CodecError, ParameterizedDecode};
 use serde::{Deserialize, Serialize};
 use std::{borrow::Cow, collections::HashSet, time::Duration};
 use tracing::warn;
@@ -23,6 +23,8 @@ use worker::*;
 
 pub(crate) const DURABLE_REPORTS_PROCESSED_INITIALIZE: &str =
     "/internal/do/reports_processed/initialize";
+pub(crate) const DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED: &str =
+    "/internal/do/reports_processed/mark_aggregated";
 
 /// Durable Object (DO) for tracking which reports have been processed.
 ///
@@ -49,13 +51,13 @@ pub struct ReportsProcessed {
 
 impl ReportsProcessed {
     /// Check if the report has been processed. If not, return None; otherwise, return the ID.
-    async fn to_checked(&self, report_id_hex: String) -> Result<Option<String>> {
-        let key = format!("processed/{report_id_hex}");
+    async fn to_checked(&self, report_id: ReportId) -> Result<Option<ReportId>> {
+        let key = format!("processed/{}", report_id.to_hex());
         let processed: bool = state_set_if_not_exists(&self.state, &key, &true)
             .await?
             .unwrap_or(false);
         if processed {
-            Ok(Some(report_id_hex))
+            Ok(Some(report_id))
         } else {
             Ok(None)
         }
@@ -92,35 +94,32 @@ impl DurableObject for ReportsProcessed {
             //  * Ensure the report won't be included in a batch that was already collected
             //  * Initialize VDAF preparation.
             //
-            // Non-idempotent
+            // Idempotent
             // Input: `ReportsProcessedReq`
             // Output: `ReportsProcessedResp`
             (DURABLE_REPORTS_PROCESSED_INITIALIZE, Method::Post) => {
                 let reports_processed_request: ReportsProcessedReq = req.json().await?;
-                let replayed_reports = try_join_all(
+                let result = 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())
+                        .map(|consumed_report| async {
+                            if let Some(replayed) = state_get::<bool>(
+                                &self.state,
+                                &format!("processed/{}", consumed_report.metadata().id.to_hex()),
+                            )
+                            .await?
+                            {
+                                if replayed {
+                                    return Result::Ok(Some(consumed_report.metadata().id.clone()));
+                                }
+                            }
+                            Ok(None)
                         }),
                 )
-                .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}")))?;
+                .await?;
+                let replayed_reports = result.into_iter().flatten().collect::<HashSet<ReportId>>();
 
                 let initialized_reports = reports_processed_request
                     .consumed_reports
@@ -154,6 +153,26 @@ impl DurableObject for ReportsProcessed {
                 })
             }
 
+            // Mark reports as aggregated. Return the subset that were already aggregated.
+            //
+            // Non-idempotent
+            // Input: `Vec<ReportId>`
+            // Output: `Vec<ReportId>`
+            (DURABLE_REPORTS_PROCESSED_MARK_AGGREGATED, Method::Post) => {
+                let report_ids: Vec<ReportId> = req.json().await?;
+                let replayed_reports = try_join_all(
+                    report_ids
+                        .into_iter()
+                        .map(|report_id| self.to_checked(report_id)),
+                )
+                .await?
+                .into_iter()
+                .flatten()
+                .collect::<Vec<ReportId>>();
+
+                Response::from_json(&replayed_reports)
+            }
+
             _ => Err(int_err(format!(
                 "ReportsProcessed: unexpected request: method={:?}; path={:?}",
                 req.method(),


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index a0e0f4c..bd12d6e 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -384,14 +384,6 @@ impl DapReportInitializer for MockAggregator {
                 {
                     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());
             }
         }
 
@@ -520,25 +512,41 @@ impl DapAggregator<BearerToken> for MockAggregator {
     async fn put_out_shares(
         &self,
         task_id: &TaskId,
+        task_config: &DapTaskConfig,
         part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
-    ) -> Result<(), DapError> {
-        let task_config = self
-            .get_task_config_for(Cow::Borrowed(task_id))
-            .await?
-            .ok_or_else(|| fatal_error!(err = "task not found"))?;
+    ) -> Result<HashSet<ReportId>, DapError> {
+        let mut report_store_guard = self
+            .report_store
+            .lock()
+            .expect("report_store: failed to lock");
+        let report_store = report_store_guard.entry(task_id.clone()).or_default();
+        let mut agg_store_guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        let agg_store = agg_store_guard.entry(task_id.clone()).or_default();
 
-        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
-        let agg_store = guard.entry(task_id.clone()).or_default();
-        for (bucket, agg_share_delta) in task_config
+        let mut replayed = HashSet::new();
+        for (bucket, out_shares) in task_config
             .batch_span_for_out_shares(part_batch_sel, out_shares)?
             .into_iter()
         {
-            let inner_agg_store = agg_store.entry(bucket.to_owned_bucket()).or_default();
-            inner_agg_store.agg_share.merge(agg_share_delta)?;
+            for out_share in out_shares.into_iter() {
+                if !report_store.processed.contains(&out_share.report_id) {
+                    // Mark report processed.
+                    report_store.processed.insert(out_share.report_id.clone());
+
+                    // Add to aggregate share.
+                    agg_store
+                        .entry(bucket.to_owned_bucket())
+                        .or_default()
+                        .agg_share
+                        .merge(DapAggregateShare::try_from_out_shares([out_share])?)?;
+                } else {
+                    replayed.insert(out_share.report_id);
+                }
+            }
         }
 
-        Ok(())
+        Ok(replayed)
     }
 
     async fn get_agg_share(