Repository: cloudflare/daphne
Pull request #524 (2024-03-06 04:35:07)

Improvements to test code
Stacked on #523 

In preparation for supporting heavy hitters, make some changes to `daphne::testing::MockAggregator`:
1. Make aggregate storage look more like `daphne_server::App`
2. Rename to `InMemoryAggregator`


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 8f9916d..5e0d048 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -269,6 +269,15 @@ pub struct DapAggregateSpan<T> {
     span: HashMap<DapBatchBucket, (T, Vec<(ReportId, Time)>)>,
 }
 
+impl std::fmt::Display for DapBatchBucket {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::TimeInterval { batch_window } => write!(f, "window/{batch_window}"),
+            Self::FixedSize { batch_id } => write!(f, "batch/{}", batch_id.to_hex()),
+        }
+    }
+}
+
 // We can't derive default because it will require T to be Default, which we don't need.
 impl<T> Default for DapAggregateSpan<T> {
     fn default() -> Self {
@@ -801,6 +810,20 @@ pub enum DapAggregationParam {
     Mastic(Poplar1AggregationParam),
 }
 
+#[cfg(any(test, feature = "test-utils"))]
+impl DapAggregationParam {
+    /// Return the aggregation level for the aggregation parameter. Replay protection is enforced
+    /// with respect to this value.
+    //
+    // TODO heavy hitters: Decide if we're comfortable with this level of abstraction.
+    pub(crate) fn level(&self) -> usize {
+        match self {
+            Self::Empty => 0,
+            Self::Mastic(agg_param) => agg_param.level(),
+        }
+    }
+}
+
 #[cfg(any(test, feature = "test-utils"))]
 impl deepsize::DeepSizeOf for DapAggregationParam {
     fn deep_size_of(&self) -> usize {
diff --git a/daphne/src/roles/mod.rs b/daphne/src/roles/mod.rs
index fc2f93e..f407242 100644
--- a/daphne/src/roles/mod.rs
+++ b/daphne/src/roles/mod.rs
@@ -151,7 +151,7 @@ async fn resolve_taskprov<S: Sync>(
 
 #[cfg(test)]
 mod test {
-    use super::{aggregator, helper, leader, DapAggregator, DapAuthorizedSender, DapLeader};
+    use super::{aggregator, helper, leader, DapAuthorizedSender, DapLeader};
     use crate::{
         assert_metrics_include, async_test_version, async_test_versions,
         auth::BearerToken,
@@ -166,17 +166,17 @@ mod test {
         },
         roles::leader::WorkItem,
         test_versions,
-        testing::{AggStore, MockAggregator},
+        testing::InMemoryAggregator,
         vdaf::{mastic::MasticWeight, MasticWeightConfig, Prio3Config, VdafConfig},
-        DapAbort, DapAggregateShare, DapAggregationJobState, DapAggregationParam, DapBatchBucket,
-        DapCollectionJob, DapError, DapGlobalConfig, DapLeaderAggregationJobTransition,
-        DapMeasurement, DapQueryConfig, DapRequest, DapResource, DapTaskConfig, DapTaskParameters,
-        DapVersion, MetaAggregationJobId,
+        DapAbort, DapAggregationJobState, DapAggregationParam, DapBatchBucket, DapCollectionJob,
+        DapError, DapGlobalConfig, DapLeaderAggregationJobTransition, DapMeasurement,
+        DapQueryConfig, DapRequest, DapResource, DapTaskConfig, DapTaskParameters, DapVersion,
+        MetaAggregationJobId,
     };
     use assert_matches::assert_matches;
     use matchit::Router;
     use prio::{
-        codec::{Decode, Encode, ParameterizedDecode, ParameterizedEncode},
+        codec::{Decode, Encode, ParameterizedEncode},
         idpf::IdpfInput,
         vdaf::poplar1::Poplar1AggregationParam,
     };
@@ -419,8 +419,8 @@ mod test {
             task_id
         }
 
-        pub fn new_helper(&self) -> Arc<MockAggregator> {
-            Arc::new(MockAggregator::new_helper(
+        pub fn new_helper(&self) -> Arc<InMemoryAggregator> {
+            Arc::new(InMemoryAggregator::new_helper(
                 self.tasks.clone(),
                 self.global_config
                     .gen_hpke_receiver_config_list(thread_rng().gen())
@@ -434,8 +434,8 @@ mod test {
             ))
         }
 
-        pub fn with_leader(self, helper: Arc<MockAggregator>) -> Test {
-            let leader = Arc::new(MockAggregator::new_leader(
+        pub fn with_leader(self, helper: Arc<InMemoryAggregator>) -> Test {
+            let leader = Arc::new(InMemoryAggregator::new_leader(
                 self.tasks,
                 self.global_config
                     .gen_hpke_receiver_config_list(thread_rng().gen())
@@ -469,8 +469,8 @@ mod test {
 
     pub(super) struct Test {
         now: Time,
-        leader: Arc<MockAggregator>,
-        helper: Arc<MockAggregator>,
+        leader: Arc<InMemoryAggregator>,
+        helper: Arc<InMemoryAggregator>,
         collector_token: BearerToken,
         taskprov_collector_token: BearerToken,
         time_interval_task_id: TaskId,
@@ -1108,90 +1108,40 @@ mod test {
     async fn handle_agg_job_req_failure_report_replayed(version: DapVersion) {
         let t = Test::new(version);
         let task_id = &t.time_interval_task_id;
+        let task_config = t.helper.unchecked_get_task_config(task_id).await;
 
         let report = t.gen_test_report(task_id).await;
-        let (leader_state, req) = t
-            .gen_test_agg_job_init_req(
-                task_id,
-                version,
-                DapAggregationParam::Empty,
-                vec![report.clone()],
-            )
-            .await;
+        let req = t.gen_test_upload_req(report.clone(), task_id).await;
+        leader::handle_upload_req(&*t.leader, &req).await.unwrap();
+
+        let query = task_config.query_for_current_batch_window(t.now);
+        let req = t.gen_test_coll_job_req(query, task_id).await;
+        leader::handle_coll_job_req(&*t.leader, &req).await.unwrap();
 
         // Add dummy data to report store backend. This is done in a new scope so that the lock on the
         // report store is released before running the test.
         {
-            let mut guard = t
-                .helper
-                .report_store
-                .lock()
-                .expect("report_store: failed to lock");
-            let report_store = guard.entry(*task_id).or_default();
-            report_store.insert(report.report_metadata.id);
+            let bucket = DapBatchBucket::TimeInterval {
+                batch_window: task_config.quantized_time_lower_bound(t.now),
+            };
+            let mut agg_store = t.helper.agg_store.lock().unwrap();
+            agg_store
+                .for_collection(task_id, &bucket, &DapAggregationParam::Empty)
+                .reports
+                .insert(report.report_metadata.id);
         }
 
-        // Get AggregationJobResp and then extract the transition data from inside.
-        let agg_job_resp = AggregationJobResp::get_decoded(
-            &helper::handle_agg_job_req(&*t.helper, &req)
-                .await
-                .unwrap()
-                .payload,
-        )
-        .unwrap();
-        let transitions = if version == DapVersion::Draft02 {
-            // in version 2 replays are only detected later.
-            let agg_job_id =
-                AggregationJobInitReq::get_decoded_with_param(&DapVersion::Draft02, &req.payload)
-                    .unwrap()
-                    .draft02_agg_job_id
-                    .unwrap();
-            let task_config = t.leader.unchecked_get_task_config(task_id).await;
-            let transition = task_config
-                .handle_agg_job_resp(
-                    task_id,
-                    &MetaAggregationJobId::Draft02(agg_job_id),
-                    leader_state,
-                    agg_job_resp,
-                    t.leader.metrics(),
-                )
-                .unwrap();
-            let DapLeaderAggregationJobTransition::Uncommitted(
-                _,
-                AggregationJobContinueReq { transitions, .. },
-            ) = transition
-            else {
-                panic!("expected uncommitted transition, was {transition:?}");
-            };
-            let req = t
-                .gen_test_agg_job_cont_req(
-                    task_id,
-                    &MetaAggregationJobId::Draft02(agg_job_id),
-                    transitions,
-                    version,
-                )
-                .await;
-            AggregationJobResp::get_decoded(
-                &helper::handle_agg_job_req(&*t.helper, &req)
-                    .await
-                    .unwrap()
-                    .payload,
-            )
-            .unwrap()
-            .transitions
-        } else {
-            agg_job_resp.transitions
-        };
+        leader::process(&*t.leader, "leader.com", 100)
+            .await
+            .unwrap();
 
-        // Expect failure due to report store marked as collected.
-        assert_matches!(
-            transitions[0].var,
-            TransitionVar::Failed(TransitionFailure::ReportReplayed)
-        );
+        assert_metrics_include!(t.leader_registry, {
+            r#"report_counter{env="test_leader",host="leader.com",status="rejected_report_replayed"}"#: 1,
+        });
 
         assert_metrics_include!(t.helper_registry, {
             r#"report_counter{env="test_helper",host="helper.org",status="rejected_report_replayed"}"#: 1,
-            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: if version == DapVersion::Draft02 { 2 } else { 1 },
+            r#"inbound_request_counter{env="test_helper",host="helper.org",type="aggregate"}"#: num_agg_job_reqs_for_version(version),
             r#"aggregation_job_counter{env="test_helper",host="helper.org",status="started"}"#: 1,
         });
     }
@@ -1210,22 +1160,13 @@ mod test {
         // Add mock data to the aggreagte store backend. This is done in its own scope so that the lock
         // is released before running the test. Otherwise the test will deadlock.
         {
-            let mut guard = t
-                .helper
-                .agg_store
-                .lock()
-                .expect("agg_store: failed to lock");
-            let agg_store = guard.entry(*task_id).or_default();
-
-            agg_store.insert(
-                DapBatchBucket::TimeInterval {
-                    batch_window: task_config.quantized_time_lower_bound(t.now),
-                },
-                AggStore {
-                    agg_share: DapAggregateShare::default(),
-                    collected: true,
-                },
-            );
+            let bucket = DapBatchBucket::TimeInterval {
+                batch_window: task_config.quantized_time_lower_bound(t.now),
+            };
+            let mut agg_store = t.helper.agg_store.lock().unwrap();
+            agg_store
+                .for_collection(task_id, &bucket, &DapAggregationParam::Empty)
+                .collected = true;
         }
 
         let query = task_config.query_for_current_batch_window(t.now);
diff --git a/daphne_service_utils/src/durable_requests/bindings.rs b/daphne_service_utils/src/durable_requests/bindings.rs
index f8c7eb3..143455c 100644
--- a/daphne_service_utils/src/durable_requests/bindings.rs
+++ b/daphne_service_utils/src/durable_requests/bindings.rs
@@ -98,14 +98,7 @@ define_do_binding! {
 
     fn name((version, task_id_hex, bucket): (DapVersion, &'n str, &'n DapBatchBucket)) -> ObjectIdFrom {
         fn durable_name_bucket(bucket: &DapBatchBucket) -> String {
-            match bucket {
-                DapBatchBucket::TimeInterval { batch_window } => {
-                    format!("window/{batch_window}")
-                }
-                DapBatchBucket::FixedSize { batch_id } => {
-                    format!("batch/{}", batch_id.to_hex())
-                }
-            }
+            format!("{bucket}")
         }
         ObjectIdFrom::Name(format!(
             "{}/{}",
@@ -165,3 +158,27 @@ define_do_binding! {
     }
 
 }
+
+#[cfg(test)]
+mod tests {
+    use daphne::{messages::BatchId, DapBatchBucket};
+
+    // We use `std::fmt::Display` for `DapBatchBucket` to format names for DO instances. Ensure
+    // that they are formatted the way we expect.
+    #[test]
+    fn bucket_display() {
+        assert_eq!(
+            "batch/1111111111111111111111111111111111111111111111111111111111111111",
+            format!(
+                "{}",
+                DapBatchBucket::FixedSize {
+                    batch_id: BatchId([17; 32])
+                }
+            )
+        );
+        assert_eq!(
+            "window/1337",
+            format!("{}", DapBatchBucket::TimeInterval { batch_window: 1337 })
+        );
+    }
+}


Test patch:
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 783827f..187eb41 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -567,16 +567,51 @@ impl AuditLog for MockAuditLog {
     }
 }
 
-pub struct MockAggregator {
+#[derive(Default)]
+#[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
+pub(crate) struct AggregateStoreForCollection {
+    /// The aggregate share.
+    pub(crate) agg_share: DapAggregateShare,
+    /// Indication of whether the aggregate share has been collected.
+    pub(crate) collected: bool,
+    /// The reports included in the aggregate share.
+    pub(crate) reports: HashSet<ReportId>,
+}
+
+#[derive(Default)]
+#[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
+pub(crate) struct InMemoryAggregateStore(HashMap<String, AggregateStoreForCollection>);
+
+impl InMemoryAggregateStore {
+    /// Return the aggregate store for the given "collection". A collection is the result of
+    /// aggregating a bucket of reports with a particular aggregation parameter.
+    pub(crate) fn for_collection(
+        &mut self,
+        task_id: &TaskId,
+        bucket: &DapBatchBucket,
+        agg_param: &DapAggregationParam,
+    ) -> &mut AggregateStoreForCollection {
+        let agg_level = agg_param.level();
+        // NOTE(cjpatton) `daphne_server::App` will use a similar naming scheme for instances of
+        // the `AggregateStore` DO. However, for backwards compatibility, if `agg_level == 0`, then
+        // we won't include it in the DO name.
+        self.0
+            .entry(format!("{task_id}/{bucket}/{agg_level}"))
+            .or_default()
+    }
+}
+
+/// An implementation of a DAP Aggregator without long-term storage. This is intended to be used
+/// for testing purposes only.
+pub struct InMemoryAggregator {
     pub global_config: DapGlobalConfig,
     pub(crate) tasks: Arc<Mutex<HashMap<TaskId, DapTaskConfig>>>,
     pub hpke_receiver_config_list: Vec<HpkeReceiverConfig>,
     pub leader_token: BearerToken,
     pub collector_token: Option<BearerToken>, // Not set by Helper
-    pub(crate) report_store: Arc<Mutex<HashMap<TaskId, HashSet<ReportId>>>>,
     pub(crate) leader_state_store: Arc<Mutex<InMemoryLeaderState>>,
-    pub(crate) helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapAggregationJobState>>>,
-    pub(crate) agg_store: Arc<Mutex<HashMap<TaskId, HashMap<DapBatchBucket, AggStore>>>>,
+    helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapAggregationJobState>>>,
+    pub(crate) agg_store: Arc<Mutex<InMemoryAggregateStore>>,
     pub collector_hpke_config: HpkeConfig,
     pub metrics: DaphnePromMetrics,
     pub(crate) audit_log: MockAuditLog,
@@ -587,11 +622,11 @@ pub struct MockAggregator {
     pub taskprov_collector_token: Option<BearerToken>, // Not set by Helper
 
     // Leader: Reference to peer. Used to simulate HTTP requests from Leader to Helper, i.e.,
-    // implement `DapLeader::send_http_post()` for `MockAggregator`. Not set by the Helper.
-    pub peer: Option<Arc<MockAggregator>>,
+    // implement `DapLeader::send_http_post()` for `InMemoryAggregator`. Not set by the Helper.
+    pub peer: Option<Arc<InMemoryAggregator>>,
 }
 
-impl DeepSizeOf for MockAggregator {
+impl DeepSizeOf for InMemoryAggregator {
     fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
         self.global_config.deep_size_of_children(context)
                 + self.tasks.deep_size_of_children(context)
@@ -600,7 +635,6 @@ impl DeepSizeOf for MockAggregator {
                     .deep_size_of_children(context)
                 + self.leader_token.deep_size_of_children(context)
                 + self.collector_token.deep_size_of_children(context)
-                + self.report_store.deep_size_of_children(context)
                 + self.helper_state_store.deep_size_of_children(context)
                 + self.agg_store.deep_size_of_children(context)
                 + self.collector_hpke_config.deep_size_of_children(context)
@@ -615,7 +649,7 @@ impl DeepSizeOf for MockAggregator {
     }
 }
 
-impl MockAggregator {
+impl InMemoryAggregator {
     #[allow(clippy::too_many_arguments)]
     pub fn new_helper(
         tasks: impl IntoIterator<Item = (TaskId, DapTaskConfig)>,
@@ -633,7 +667,6 @@ impl MockAggregator {
             hpke_receiver_config_list: hpke_receiver_config_list.into_iter().collect(),
             leader_token,
             collector_token: None,
-            report_store: Default::default(),
             leader_state_store: Default::default(),
             helper_state_store: Default::default(),
             agg_store: Default::default(),
@@ -667,7 +700,6 @@ impl MockAggregator {
             hpke_receiver_config_list: hpke_receiver_config_list.into_iter().collect(),
             leader_token,
             collector_token: collector_token.into(),
-            report_store: Default::default(),
             leader_state_store: Default::default(),
             helper_state_store: Default::default(),
             agg_store: Default::default(),
@@ -700,7 +732,7 @@ impl MockAggregator {
 }
 
 #[async_trait]
-impl BearerTokenProvider for MockAggregator {
+impl BearerTokenProvider for InMemoryAggregator {
     type WrappedBearerToken<'a> = &'a BearerToken;
 
     async fn get_leader_bearer_token_for<'s>(
@@ -722,18 +754,18 @@ impl BearerTokenProvider for MockAggregator {
     ) -> Result<Option<Self::WrappedBearerToken<'s>>, DapError> {
         if task_config.method_is_taskprov() {
             Ok(Some(self.taskprov_collector_token.as_ref().expect(
-                "MockAggregator not configured with taskprov collector token",
+                "InMemoryAggregator not configured with taskprov collector token",
             )))
         } else {
             Ok(Some(self.collector_token.as_ref().expect(
-                "MockAggregator not configured with collector token",
+                "InMemoryAggregator not configured with collector token",
             )))
         }
     }
 }
 
 #[async_trait]
-impl HpkeDecrypter for MockAggregator {
+impl HpkeDecrypter for InMemoryAggregator {
     type WrappedHpkeConfig<'a> = &'a HpkeConfig;
 
     async fn get_hpke_config_for<'s>(
@@ -746,9 +778,9 @@ impl HpkeDecrypter for MockAggregator {
         }
 
         // Aggregators MAY abort if the HPKE config request does not specify a task ID. While not
-        // required for MockAggregator, we simulate this behavior for testing purposes.
+        // required for InMemoryAggregator, we simulate this behavior for testing purposes.
         //
-        // TODO(cjpatton) To make this clearer, have MockAggregator store a map from task IDs to
+        // TODO(cjpatton) To make this clearer, have InMemoryAggregator store a map from task IDs to
         // HPKE receiver configs.
         if task_id.is_none() {
             return Err(DapError::Abort(DapAbort::MissingTaskId));
@@ -779,7 +811,7 @@ impl HpkeDecrypter for MockAggregator {
 }
 
 #[async_trait]
-impl DapAuthorizedSender<BearerToken> for MockAggregator {
+impl DapAuthorizedSender<BearerToken> for InMemoryAggregator {
     async fn authorize(
         &self,
         task_id: &TaskId,
@@ -795,7 +827,7 @@ impl DapAuthorizedSender<BearerToken> for MockAggregator {
 }
 
 #[async_trait]
-impl DapReportInitializer for MockAggregator {
+impl DapReportInitializer for InMemoryAggregator {
     async fn initialize_reports(
         &self,
         is_leader: bool,
@@ -808,9 +840,9 @@ impl DapReportInitializer for MockAggregator {
 }
 
 #[async_trait]
-impl DapAggregator<BearerToken> for MockAggregator {
+impl DapAggregator<BearerToken> for InMemoryAggregator {
     // The lifetimes on the traits ensure that we can return a reference to a task config stored by
-    // the DapAggregator. (See DaphneWorkerConfig for an example.) For simplicity, MockAggregator
+    // the DapAggregator. (See DaphneWorkerConfig for an example.) For simplicity, InMemoryAggregator
     // clones the task config as needed.
     type WrappedDapTaskConfig<'a> = DapTaskConfig;
 
@@ -877,17 +909,16 @@ impl DapAggregator<BearerToken> for MockAggregator {
             .get_task_config_for(task_id)
             .await?
             .ok_or(DapError::Abort(DapAbort::UnrecognizedTask))?;
-
-        let agg_store = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
-        let Some(agg_store_per_task) = agg_store.get(task_id) else {
-            return Ok(false);
-        };
+        let mut agg_store = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
+        // TODO heavy hitters: Replace this with the agg param specified by the Collector.
+        let agg_param = DapAggregationParam::Empty;
 
         for bucket in task_config.batch_span_for_sel(batch_sel)? {
-            if let Some(inner_agg_store) = agg_store_per_task.get(&bucket) {
-                if inner_agg_store.collected {
-                    return Ok(true);
-                }
+            if agg_store
+                .for_collection(task_id, &bucket, &agg_param)
+                .collected
+            {
+                return Ok(true);
             }
         }
 
@@ -900,11 +931,13 @@ impl DapAggregator<BearerToken> for MockAggregator {
         };
 
         let aggregated = {
-            let agg_store_guard = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
-            agg_store_guard
-                .get(task_id)
-                .map(|agg_store| agg_store.get(&bucket))
-                .is_some()
+            let mut agg_store = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
+            // TODO heavy hitters: Replace this with the agg param specified by the Collector.
+            let agg_param = DapAggregationParam::Empty;
+            !agg_store
+                .for_collection(task_id, &bucket, &agg_param)
+                .agg_share
+                .empty()
         };
 
         let uploaded = {
@@ -924,31 +957,31 @@ impl DapAggregator<BearerToken> for MockAggregator {
         _task_config: &DapTaskConfig,
         agg_agg_span: DapAggregateSpan<DapAggregateShare>,
     ) -> DapAggregateSpan<Result<(), MergeAggShareError>> {
-        let mut report_store_guard = self
-            .report_store
-            .lock()
-            .expect("report_store: failed to lock");
-        let report_store = report_store_guard.entry(*task_id).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).or_default();
+        let mut agg_store = self.agg_store.lock().unwrap();
+        // TODO heavy hitters: Replace this with the agg param specified by the Collector.
+        let agg_param = DapAggregationParam::Empty;
 
         agg_agg_span
             .into_iter()
             .map(|(bucket, (agg_share_delta, report_metadatas))| {
+                let agg_store_for_collection =
+                    agg_store.for_collection(task_id, &bucket, &agg_param);
+
                 let replayed = report_metadatas
                     .iter()
                     .map(|(id, _)| *id)
-                    .filter(|id| report_store.contains(id))
+                    .filter(|id| agg_store_for_collection.reports.contains(id))
                     .collect::<HashSet<_>>();
 
                 let result = if replayed.is_empty() {
-                    report_store.extend(report_metadatas.iter().map(|(id, _)| *id));
+                    agg_store_for_collection
+                        .reports
+                        .extend(report_metadatas.iter().map(|(id, _)| *id));
                     // Add to aggregate share.
-                    let agg_share = agg_store.entry(bucket.clone()).or_default();
-                    if agg_share.collected {
+                    if agg_store_for_collection.collected {
                         Err(MergeAggShareError::AlreadyCollected)
                     } else {
-                        agg_share
+                        agg_store_for_collection
                             .agg_share
                             .merge(agg_share_delta.clone())
                             .map_err(MergeAggShareError::Other)
@@ -971,18 +1004,18 @@ impl DapAggregator<BearerToken> for MockAggregator {
             .await
             .unwrap()
             .expect("tasks: unrecognized task");
-        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
-        let agg_store = guard.entry(*task_id).or_default();
+        let mut agg_store = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
+        // TODO heavy hitters: Replace this with the agg param specified by the Collector.
+        let agg_param = DapAggregationParam::Empty;
 
         // Fetch aggregate shares.
         let mut agg_share = DapAggregateShare::default();
         for bucket in task_config.batch_span_for_sel(batch_sel)? {
-            if let Some(inner_agg_store) = agg_store.get(&bucket) {
-                if inner_agg_store.collected {
-                    return Err(DapError::Abort(DapAbort::batch_overlap(task_id, batch_sel)));
-                }
-                agg_share.merge(inner_agg_store.agg_share.clone())?;
+            let agg_store_for_collection = agg_store.for_collection(task_id, &bucket, &agg_param);
+            if agg_store_for_collection.collected {
+                return Err(DapError::Abort(DapAbort::batch_overlap(task_id, batch_sel)));
             }
+            agg_share.merge(agg_store_for_collection.agg_share.clone())?;
         }
 
         Ok(agg_share)
@@ -994,13 +1027,14 @@ impl DapAggregator<BearerToken> for MockAggregator {
         batch_sel: &BatchSelector,
     ) -> Result<(), DapError> {
         let task_config = self.unchecked_get_task_config(task_id).await;
-        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
-        let agg_store = guard.entry(*task_id).or_default();
+        let mut agg_store = self.agg_store.lock().map_err(|e| fatal_error!(err = ?e))?;
+        // TODO heavy hitters: Replace this with the agg param specified by the Collector.
+        let agg_param = DapAggregationParam::Empty;
 
         for bucket in task_config.batch_span_for_sel(batch_sel)? {
-            if let Some(inner_agg_store) = agg_store.get_mut(&bucket) {
-                inner_agg_store.collected = true;
-            }
+            agg_store
+                .for_collection(task_id, &bucket, &agg_param)
+                .collected = true;
         }
 
         Ok(())
@@ -1020,7 +1054,7 @@ impl DapAggregator<BearerToken> for MockAggregator {
 }
 
 #[async_trait]
-impl DapHelper<BearerToken> for MockAggregator {
+impl DapHelper<BearerToken> for InMemoryAggregator {
     async fn put_helper_state_if_not_exists<Id>(
         &self,
         task_id: &TaskId,
@@ -1076,7 +1110,7 @@ impl DapHelper<BearerToken> for MockAggregator {
 }
 
 #[async_trait]
-impl DapLeader<BearerToken> for MockAggregator {
+impl DapLeader<BearerToken> for InMemoryAggregator {
     async fn put_report(&self, report: &Report, task_id: &TaskId) -> Result<(), DapError> {
         let task_config = self
             .get_task_config_for(task_id)
@@ -1207,21 +1241,11 @@ impl DapLeader<BearerToken> for MockAggregator {
 /// Information associated to a certain helper state for a given task ID and aggregate job ID.
 #[derive(Clone, Eq, Hash, PartialEq, Deserialize, Serialize)]
 #[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
-pub struct HelperStateInfo {
+struct HelperStateInfo {
     task_id: TaskId,
     agg_job_id_owned: MetaAggregationJobId,
 }
 
-/// `AggStore` keeps track of the following:
-/// * Aggregate share
-/// * Whether this aggregate share has been collected
-#[derive(Default)]
-#[cfg_attr(any(test, feature = "test-utils"), derive(deepsize::DeepSizeOf))]
-pub struct AggStore {
-    pub(crate) agg_share: DapAggregateShare,
-    pub(crate) collected: bool,
-}
-
 /// Helper macro used by `assert_metrics_include`.
 #[macro_export]
 macro_rules! assert_metrics_include_auxiliary_function {