Repository: cloudflare/daphne
Pull request #133 (2022-10-03 16:44:48)

daphne: Implement fixed-size queries
Partially addresses #100.
Based on #132 (merge that first).

Complete the implementation of fixed-size queries and add tests for MockAggregator. While at it, add the PartialBatchSelector struct to CollectResp.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 1cb3415..f070890 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -30,7 +30,10 @@
 
 use crate::{
     hpke::HpkeReceiverConfig,
-    messages::{CollectResp, Duration, HpkeConfig, Nonce, Time, TransitionFailure},
+    messages::{
+        BatchSelector, CollectResp, Duration, HpkeConfig, Id, Interval, Nonce,
+        PartialBatchSelector, ReportMetadata, Time, TransitionFailure,
+    },
     vdaf::{
         prio2::prio2_decode_prepare_state,
         prio3::{prio3_append_prepare_state, prio3_decode_prepare_state},
@@ -364,6 +367,29 @@ impl DapQueryConfig {
         Ok(report_count >= self.min_batch_size())
     }
 
+    pub(crate) fn is_valid_part_batch_sel(&self, part_batch_sel: &PartialBatchSelector) -> bool {
+        matches!(
+            (&self, part_batch_sel),
+            (
+                Self::TimeInterval { .. },
+                PartialBatchSelector::TimeInterval
+            ) | (
+                Self::FixedSize { .. },
+                PartialBatchSelector::FixedSize { .. }
+            )
+        )
+    }
+
+    pub(crate) fn is_valid_batch_sel(&self, batch_sel: &BatchSelector) -> bool {
+        matches!(
+            (&self, batch_sel),
+            (
+                Self::TimeInterval { .. },
+                BatchSelector::TimeInterval { .. }
+            ) | (Self::FixedSize { .. }, BatchSelector::FixedSize { .. })
+        )
+    }
+
     /// Return the minimum batch size determined by the query configuration.
     pub fn min_batch_size(&self) -> u64 {
         match self {
@@ -373,6 +399,29 @@ impl DapQueryConfig {
     }
 }
 
+/// A batch bucket.
+///
+/// A bucket is the smallest, disjoint set of reports that can be queried: For time-interval
+/// queries, the bucket to which a report is assigned is determined by truncating its timestamp by
+/// the task's `time_precision` parameter; for fixed-size queries, the span consists of a single
+/// bucket, which is the batch determined by the batch ID (i.e., the partial batch selector).
+#[derive(Clone, Eq, Hash, PartialEq)]
+pub enum DapBatchBucket<'a> {
+    FixedSize { batch_id: &'a Id },
+    TimeInterval { batch_window: Time },
+}
+
+impl<'a> DapBatchBucket<'a> {
+    /// TODO(issue #100) Deprecate this once "fixed_size" support is added.
+    pub fn unwrap_window(&self) -> Time {
+        if let Self::TimeInterval { batch_window } = self {
+            *batch_window
+        } else {
+            panic!("TODO(issue #100)");
+        }
+    }
+}
+
 /// Per-task DAP parameters.
 #[derive(Clone, Deserialize, Serialize)]
 #[serde(try_from = "ShadowDapTaskConfig")]
@@ -421,6 +470,87 @@ impl DapTaskConfig {
     pub(crate) fn truncate_time(&self, time: Time) -> Time {
         time - (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.
+    pub fn batch_span_for_out_shares<'a>(
+        &self,
+        part_batch_sel: &'a PartialBatchSelector,
+        out_shares: Vec<DapOutputShare>,
+    ) -> Result<HashMap<DapBatchBucket<'a>, DapAggregateShare>, DapError> {
+        if !self.query.is_valid_part_batch_sel(part_batch_sel) {
+            return Err(DapError::fatal(
+                "partial batch selector not compatible with task",
+            ));
+        }
+
+        let mut agg_shares: HashMap<DapBatchBucket<'a>, DapAggregateShare> = HashMap::new();
+        for out_share in out_shares.into_iter() {
+            let bucket = match part_batch_sel {
+                PartialBatchSelector::TimeInterval => DapBatchBucket::TimeInterval {
+                    batch_window: self.truncate_time(out_share.time),
+                },
+                PartialBatchSelector::FixedSize { batch_id } => {
+                    DapBatchBucket::FixedSize { batch_id }
+                }
+            };
+
+            let agg_share = agg_shares.entry(bucket).or_default();
+            agg_share.merge(DapAggregateShare {
+                report_count: 1,
+                checksum: out_share.checksum,
+                data: Some(out_share.data),
+            })?;
+        }
+
+        Ok(agg_shares)
+    }
+
+    /// Return the batch span determined by the given batch selector.
+    pub fn batch_span_for_sel<'a>(
+        &self,
+        batch_sel: &'a BatchSelector,
+    ) -> Result<Vec<DapBatchBucket<'a>>, DapError> {
+        if !self.query.is_valid_batch_sel(batch_sel) {
+            return Err(DapError::fatal("batch selector not compatible with task"));
+        }
+
+        match batch_sel {
+            BatchSelector::TimeInterval {
+                batch_interval: Interval { start, duration },
+            } => {
+                let windows = duration / self.time_precision;
+                Ok((0..windows)
+                    .map(|i| DapBatchBucket::TimeInterval {
+                        batch_window: start + i * self.time_precision,
+                    })
+                    .collect())
+            }
+            BatchSelector::FixedSize { batch_id } => {
+                Ok(vec![DapBatchBucket::FixedSize { batch_id }])
+            }
+        }
+    }
+
+    /// Return the batch bucket for a report with the given metadata.
+    pub fn batch_bucket_for_meta<'a>(
+        &self,
+        metadata: &'a ReportMetadata,
+        part_batch_sel: &'a PartialBatchSelector,
+    ) -> Result<DapBatchBucket<'a>, DapError> {
+        if !self.query.is_valid_part_batch_sel(part_batch_sel) {
+            return Err(DapError::fatal(
+                "partial batch selector not compatible with task",
+            ));
+        }
+
+        Ok(match part_batch_sel {
+            PartialBatchSelector::FixedSize { batch_id } => DapBatchBucket::FixedSize { batch_id },
+            PartialBatchSelector::TimeInterval => DapBatchBucket::TimeInterval {
+                batch_window: self.truncate_time(metadata.time),
+            },
+        })
+    }
 }
 
 impl AsRef<DapTaskConfig> for DapTaskConfig {
@@ -496,6 +626,7 @@ pub struct DapLeaderUncommitted {
 /// The Helper's state during the aggregation flow.
 #[derive(Clone, Debug, PartialEq)]
 pub struct DapHelperState {
+    pub(crate) part_batch_sel: PartialBatchSelector,
     pub(crate) seq: Vec<(VdafState, Time, Nonce)>,
 }
 
@@ -509,6 +640,7 @@ impl DapHelperState {
     /// Note that the encoding format is not specified by the DAP standard.
     pub fn get_encoded(&self, vdaf_config: &VdafConfig) -> Result<Vec<u8>, DapError> {
         let mut bytes = vec![];
+        self.part_batch_sel.encode(&mut bytes);
         for (state, time, nonce) in self.seq.iter() {
             match (vdaf_config, state) {
                 (VdafConfig::Prio3(prio3_config), _) => {
@@ -527,8 +659,9 @@ impl DapHelperState {
 
     /// Decode the Helper state from a byte string.
     pub fn get_decoded(vdaf_config: &VdafConfig, data: &[u8]) -> Result<Self, DapError> {
-        let mut seq = vec![];
         let mut r = std::io::Cursor::new(data);
+        let part_batch_sel = PartialBatchSelector::decode(&mut r)?;
+        let mut seq = vec![];
         while (r.position() as usize) < data.len() {
             let state = match vdaf_config {
                 VdafConfig::Prio3(ref prio3_config) => {
@@ -543,7 +676,10 @@ impl DapHelperState {
             seq.push((state, time, nonce))
         }
 
-        Ok(DapHelperState { seq })
+        Ok(DapHelperState {
+            part_batch_sel,
+            seq,
+        })
     }
 }
 
@@ -603,34 +739,27 @@ impl DapAggregateShare {
         Ok(())
     }
 
-    /// Transform a sequence of output shares into a map from batch intervals to aggregate shares.
-    pub fn batches_from_out_shares(
-        out_shares: Vec<DapOutputShare>,
-        min_batch_interval: u64,
-    ) -> Result<HashMap<u64, Self>, DapError> {
-        let mut agg_shares: HashMap<u64, Self> = HashMap::new();
-        for out_share in out_shares.into_iter() {
-            let interval = out_share.time - (out_share.time % min_batch_interval);
-            let agg_share_delta = DapAggregateShare {
-                report_count: 1,
-                checksum: out_share.checksum,
-                data: Some(out_share.data),
-            };
-            if let Some(agg_share) = agg_shares.get_mut(&interval) {
-                agg_share.merge(agg_share_delta)?;
-            } else {
-                agg_shares.insert(interval, agg_share_delta);
-            }
-        }
-        Ok(agg_shares)
-    }
-
     /// Set the aggregate share to zero.
     pub fn reset(&mut self) {
         self.report_count = 0;
         self.checksum = [0; 32];
         self.data = None;
     }
+
+    #[cfg(test)]
+    pub(crate) fn try_from_out_shares(
+        out_shares: impl IntoIterator<Item = DapOutputShare>,
+    ) -> Result<Self, DapError> {
+        let mut agg_share = Self::default();
+        for out_share in out_shares.into_iter() {
+            agg_share.merge(DapAggregateShare {
+                report_count: 1,
+                checksum: out_share.checksum,
+                data: Some(out_share.data),
+            })?;
+        }
+        Ok(agg_share)
+    }
 }
 
 /// Leader state transition during the aggregation flow.
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index 531ecdf..b738a50 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -201,13 +201,22 @@ impl Decode for ReportShare {
 
 /// Batch parameter conveyed to the Helper by the Leader in the aggreagtion sub-protocol. Used to
 /// identify which batch the reports in the [`AggregateInitializeReq`] are intended for.
-#[derive(Clone, Debug, PartialEq, Eq)]
-pub enum BatchParameter {
+#[derive(Clone, Debug, Eq, Deserialize, PartialEq, Serialize)]
+pub enum PartialBatchSelector {
     TimeInterval,
     FixedSize { batch_id: Id },
 }
 
-impl Encode for BatchParameter {
+impl From<BatchSelector> for PartialBatchSelector {
+    fn from(batch_sel: BatchSelector) -> Self {
+        match batch_sel {
+            BatchSelector::TimeInterval { .. } => Self::TimeInterval,
+            BatchSelector::FixedSize { batch_id } => Self::FixedSize { batch_id },
+        }
+    }
+}
+
+impl Encode for PartialBatchSelector {
     fn encode(&self, bytes: &mut Vec<u8>) {
         match self {
             Self::TimeInterval => QUERY_TYPE_TIME_INTERVAL.encode(bytes),
@@ -219,7 +228,7 @@ impl Encode for BatchParameter {
     }
 }
 
-impl Decode for BatchParameter {
+impl Decode for PartialBatchSelector {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         match u16::decode(bytes)? {
             QUERY_TYPE_TIME_INTERVAL => Ok(Self::TimeInterval),
@@ -237,7 +246,7 @@ pub struct AggregateInitializeReq {
     pub task_id: Id,
     pub agg_job_id: Id,
     pub agg_param: Vec<u8>,
-    pub batch_param: BatchParameter,
+    pub part_batch_sel: PartialBatchSelector,
     pub report_shares: Vec<ReportShare>,
 }
 
@@ -246,7 +255,7 @@ impl Encode for AggregateInitializeReq {
         self.task_id.encode(bytes);
         self.agg_job_id.encode(bytes);
         encode_u16_bytes(bytes, &self.agg_param);
-        self.batch_param.encode(bytes);
+        self.part_batch_sel.encode(bytes);
         encode_u32_items(bytes, &(), &self.report_shares);
     }
 }
@@ -257,7 +266,7 @@ impl Decode for AggregateInitializeReq {
             task_id: Id::decode(bytes)?,
             agg_job_id: Id::decode(bytes)?,
             agg_param: decode_u16_bytes(bytes)?,
-            batch_param: BatchParameter::decode(bytes)?,
+            part_batch_sel: PartialBatchSelector::decode(bytes)?,
             report_shares: decode_u32_items(&(), bytes)?,
         })
     }
@@ -552,12 +561,14 @@ impl Decode for CollectReq {
 // TODO Add serialization tests.
 #[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
 pub struct CollectResp {
+    pub part_batch_sel: PartialBatchSelector,
     pub report_count: u64,
     pub encrypted_agg_shares: Vec<HpkeCiphertext>,
 }
 
 impl Encode for CollectResp {
     fn encode(&self, bytes: &mut Vec<u8>) {
+        self.part_batch_sel.encode(bytes);
         self.report_count.encode(bytes);
         encode_u32_items(bytes, &(), &self.encrypted_agg_shares);
     }
@@ -566,6 +577,7 @@ impl Encode for CollectResp {
 impl Decode for CollectResp {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
+            part_batch_sel: PartialBatchSelector::decode(bytes)?,
             report_count: u64::decode(bytes)?,
             encrypted_agg_shares: decode_u32_items(&(), bytes)?,
         })
@@ -585,7 +597,7 @@ pub type BatchSelector = Query;
 #[derive(Clone, Debug, Default)]
 pub struct AggregateShareReq {
     pub task_id: Id,
-    pub batch_selector: BatchSelector,
+    pub batch_sel: BatchSelector,
     pub agg_param: Vec<u8>,
     pub report_count: u64,
     pub checksum: [u8; 32],
@@ -594,7 +606,7 @@ pub struct AggregateShareReq {
 impl Encode for AggregateShareReq {
     fn encode(&self, bytes: &mut Vec<u8>) {
         self.task_id.encode(bytes);
-        self.batch_selector.encode(bytes);
+        self.batch_sel.encode(bytes);
         encode_u16_bytes(bytes, &self.agg_param);
         self.report_count.encode(bytes);
         bytes.extend_from_slice(&self.checksum);
@@ -605,7 +617,7 @@ impl Decode for AggregateShareReq {
     fn decode(bytes: &mut Cursor<&[u8]>) -> Result<Self, CodecError> {
         Ok(Self {
             task_id: Id::decode(bytes)?,
-            batch_selector: BatchSelector::decode(bytes)?,
+            batch_sel: BatchSelector::decode(bytes)?,
             agg_param: decode_u16_bytes(bytes)?,
             report_count: u64::decode(bytes)?,
             checksum: {
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 11cef7b..e55a067 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -12,8 +12,8 @@ use crate::{
     hpke::HpkeDecrypter,
     messages::{
         constant_time_eq, AggregateContinueReq, AggregateInitializeReq, AggregateResp,
-        AggregateShareReq, AggregateShareResp, BatchParameter, BatchSelector, CollectReq,
-        CollectResp, Id, Nonce, Report, ReportShare, Time, TransitionFailure, TransitionVar,
+        AggregateShareReq, AggregateShareResp, BatchSelector, CollectReq, CollectResp, Id, Nonce,
+        PartialBatchSelector, Report, ReportShare, Time, TransitionFailure, TransitionVar,
     },
     DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
     DapHelperTransition, DapLeaderProcessTelemetry, DapLeaderTransition, DapOutputShare,
@@ -56,20 +56,25 @@ pub trait DapAggregator<'a, S>: HpkeDecrypter<'a> + Sized {
     ) -> Result<Option<Self::WrappedDapTaskConfig>, DapError>;
 
     /// Get the current time (number of seconds since the beginning of UNIX time).
-    fn get_current_time(&self) -> u64;
+    fn get_current_time(&self) -> Time;
 
     /// Check whether the batch determined by the collect request would overlap with a previous
     /// batch.
     async fn is_batch_overlapping(
         &self,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
     ) -> Result<bool, DapError>;
 
+    /// Check whether the given batch ID has been observed before. This is called by the Leader
+    /// (resp. Helper) in response to a CollectReq (resp. AggregateShareReq) for fixed-size tasks.
+    async fn batch_exists(&self, task_id: &Id, batch_id: &Id) -> Result<bool, DapError>;
+
     /// Store a set of output shares.
     async fn put_out_shares(
         &self,
         task_id: &Id,
+        part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
     ) -> Result<(), DapError>;
 
@@ -77,15 +82,12 @@ pub trait DapAggregator<'a, S>: HpkeDecrypter<'a> + Sized {
     async fn get_agg_share(
         &self,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
     ) -> Result<DapAggregateShare, DapError>;
 
     /// Mark a batch as collected.
-    async fn mark_collected(
-        &self,
-        task_id: &Id,
-        batch_selector: &BatchSelector,
-    ) -> Result<(), DapError>;
+    async fn mark_collected(&self, task_id: &Id, batch_sel: &BatchSelector)
+        -> Result<(), DapError>;
 
     /// Handle HTTP GET to `/hpke_config?task_id=<task_id>`.
     async fn http_get_hpke_config(&'a self, req: &DapRequest<S>) -> Result<DapResponse, DapAbort> {
@@ -169,7 +171,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
     async fn get_reports(
         &self,
         selector: &Self::ReportSelector,
-    ) -> Result<HashMap<Id, Vec<Report>>, DapError>;
+    ) -> Result<HashMap<Id, (PartialBatchSelector, Vec<Report>)>, DapError>;
 
     /// Create a collect job.
     //
@@ -265,8 +267,6 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
 
         // Ensure the batch boundaries are valid and that the batch doesn't overlap with previosuly
         // collected batches.
-        //
-        // TODO(issue #100) For "fixed_size" we need to check if we recognize the batch ID.
         check_batch(
             self,
             task_config,
@@ -291,6 +291,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
         &self,
         task_id: &Id,
         task_config: &DapTaskConfig,
+        part_batch_sel: PartialBatchSelector,
         reports: Vec<Report>,
     ) -> Result<u64, DapAbort> {
         let mut rng = thread_rng();
@@ -304,6 +305,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
                 &task_config.vdaf_verify_key,
                 task_id,
                 &agg_job_id,
+                &part_batch_sel,
                 reports,
             )
             .await?;
@@ -356,7 +358,8 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
             .vdaf
             .handle_final_agg_resp(uncommited, agg_resp)?;
         let out_shares_count = out_shares.len() as u64;
-        self.put_out_shares(task_id, out_shares).await?;
+        self.put_out_shares(task_id, &part_batch_sel, out_shares)
+            .await?;
         Ok(out_shares_count)
     }
 
@@ -394,7 +397,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
         // Prepare AggregateShareReq.
         let agg_share_req = AggregateShareReq {
             task_id: collect_req.task_id.clone(),
-            batch_selector: collect_req.query.clone(),
+            batch_sel: collect_req.query.clone(),
             agg_param: collect_req.agg_param.clone(),
             report_count: leader_agg_share.report_count,
             checksum: leader_agg_share.checksum,
@@ -413,6 +416,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
 
         // Complete the collect job.
         let collect_resp = CollectResp {
+            part_batch_sel: collect_req.query.clone().into(),
             report_count: leader_agg_share.report_count,
             encrypted_agg_shares: vec![leader_enc_agg_share, agg_share_resp.encrypted_agg_share],
         };
@@ -420,7 +424,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
             .await?;
 
         // Mark reports as collected.
-        self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_selector)
+        self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_sel)
             .await?;
 
         Ok(agg_share_req.report_count)
@@ -443,7 +447,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
         // Fetch reports and run an aggregation job for each task.
         //
         // TODO Handle tasks in parallel.
-        for (task_id, reports) in self.get_reports(selector).await?.into_iter() {
+        for (task_id, (part_batch_sel, reports)) in self.get_reports(selector).await?.into_iter() {
             let task_config = self
                 .get_task_config_for(&task_id)
                 .await?
@@ -452,7 +456,7 @@ pub trait DapLeader<'a, S>: DapAuthorizedSender<S> + DapAggregator<'a, S> {
             telem.reports_processed += reports.len() as u64;
             if !reports.is_empty() {
                 telem.reports_aggregated += self
-                    .run_agg_job(&task_id, task_config.as_ref(), reports)
+                    .run_agg_job(&task_id, task_config.as_ref(), part_batch_sel, reports)
                     .await?;
             }
         }
@@ -487,6 +491,7 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
     async fn mark_aggregated(
         &self,
         task_id: &Id,
+        part_batch_sel: &PartialBatchSelector,
         report_shares: &[ReportShare],
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError>;
 
@@ -539,13 +544,16 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
                 // Ensure we know which batch the request pertains to.
                 check_part_batch(
                     task_config,
-                    &agg_init_req.batch_param,
+                    &agg_init_req.part_batch_sel,
                     &agg_init_req.agg_param,
                 )?;
 
                 // Run mark_aggregated and handle_agg_init_req concurrently.
-                let early_rejects_future =
-                    self.mark_aggregated(&agg_init_req.task_id, &agg_init_req.report_shares);
+                let early_rejects_future = self.mark_aggregated(
+                    &agg_init_req.task_id,
+                    &agg_init_req.part_batch_sel,
+                    &agg_init_req.report_shares,
+                );
 
                 let transition = task_config
                     .vdaf
@@ -619,6 +627,7 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
                     .get_helper_state(&agg_cont_req.task_id, &agg_cont_req.agg_job_id)
                     .await?
                     .ok_or(DapAbort::UnrecognizedAggregationJob)?;
+                let part_batch_sel = state.part_batch_sel.clone();
                 let transition = task_config.vdaf.handle_agg_cont_req(state, &agg_cont_req)?;
 
                 let agg_resp = match transition {
@@ -626,7 +635,7 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
                         return Err(DapError::fatal("unexpected transition (continued)").into());
                     }
                     DapHelperTransition::Finish(out_shares, agg_resp) => {
-                        self.put_out_shares(&agg_cont_req.task_id, out_shares)
+                        self.put_out_shares(&agg_cont_req.task_id, &part_batch_sel, out_shares)
                             .await?;
                         agg_resp
                     }
@@ -675,20 +684,18 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
 
         // Ensure the batch boundaries are valid and that the batch doesn't overlap with previosuly
         // collected batches.
-        //
-        // TODO(issue #100) For "fixed_size" we need to check if we recognize the batch ID.
         check_batch(
             self,
             task_config,
             &agg_share_req.task_id,
-            &agg_share_req.batch_selector,
+            &agg_share_req.batch_sel,
             &agg_share_req.agg_param,
             now,
         )
         .await?;
 
         let agg_share = self
-            .get_agg_share(&agg_share_req.task_id, &agg_share_req.batch_selector)
+            .get_agg_share(&agg_share_req.task_id, &agg_share_req.batch_sel)
             .await?;
 
         // Check that we have aggreagted the same set of reports as the leader.
@@ -708,13 +715,13 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
         }
 
         // Mark each aggregated report as collected.
-        self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_selector)
+        self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_sel)
             .await?;
 
         let encrypted_agg_share = task_config.vdaf.produce_helper_encrypted_agg_share(
             &task_config.collector_hpke_config,
             &agg_share_req.task_id,
-            &agg_share_req.batch_selector,
+            &agg_share_req.batch_sel,
             &agg_share,
         )?;
 
@@ -731,14 +738,12 @@ pub trait DapHelper<'a, S>: DapAggregator<'a, S> {
 
 fn check_part_batch(
     task_config: &DapTaskConfig,
-    batch_param: &BatchParameter,
+    part_batch_sel: &PartialBatchSelector,
     agg_param: &[u8],
 ) -> Result<(), DapAbort> {
-    match (&task_config.query, batch_param) {
-        (DapQueryConfig::TimeInterval { .. }, BatchParameter::TimeInterval)
-        | (DapQueryConfig::FixedSize { .. }, BatchParameter::FixedSize { .. }) => (),
-        _ => return Err(DapAbort::QueryMismatch),
-    };
+    if !task_config.query.is_valid_part_batch_sel(part_batch_sel) {
+        return Err(DapAbort::QueryMismatch);
+    }
 
     // Check that the aggreation parameter is suitable for the given VDAF.
     if !task_config.vdaf.is_valid_agg_param(agg_param) {
@@ -792,7 +797,17 @@ async fn check_batch<'a, S>(
                 ));
             }
         }
-        (DapQueryConfig::FixedSize { .. }, BatchSelector::FixedSize { .. }) => (),
+        (DapQueryConfig::FixedSize { .. }, BatchSelector::FixedSize { batch_id }) => {
+            // TODO(cjpatton) The Helper can avoid this callback by first fetching the aggregate
+            // share and aborting with "batchInvalid" if the report count is 0. Depending on how we
+            // resolve https://github.com/ietf-wg-ppm/draft-ietf-ppm-dap/issues/342, this check may
+            // become unnecessary for the Leader.
+            //
+            // Consider removing this callback once we resolve DAP issue #342.
+            if !agg.batch_exists(task_id, batch_id).await? {
+                return Err(DapAbort::BatchInvalid);
+            }
+        }
         _ => return Err(DapAbort::QueryMismatch),
     };
 
diff --git a/daphne/src/vdaf/mod.rs b/daphne/src/vdaf/mod.rs
index 6bb44ba..d4752b7 100644
--- a/daphne/src/vdaf/mod.rs
+++ b/daphne/src/vdaf/mod.rs
@@ -8,7 +8,7 @@ use crate::{
     hpke::HpkeDecrypter,
     messages::{
         encode_u32_bytes, AggregateContinueReq, AggregateInitializeReq, AggregateResp,
-        BatchParameter, BatchSelector, HpkeCiphertext, HpkeConfig, Id, Nonce, Report,
+        BatchSelector, HpkeCiphertext, HpkeConfig, Id, Nonce, PartialBatchSelector, Report,
         ReportMetadata, ReportShare, Time, Transition, TransitionFailure, TransitionVar,
     },
     vdaf::{
@@ -150,7 +150,6 @@ impl VdafConfig {
             extensions: Vec::new(),
         };
 
-        // TODO(issue #100) Set this to the output of `shard()`.
         let public_share = Vec::new();
         let encoded_input_shares = match self {
             Self::Prio3(prio3_config) => prio3_shard(prio3_config, measurement)?,
@@ -291,6 +290,7 @@ impl VdafConfig {
         verify_key: &VdafVerifyKey,
         task_id: &Id,
         agg_job_id: &Id,
+        part_batch_sel: &PartialBatchSelector,
         reports: Vec<Report>,
     ) -> Result<DapLeaderTransition<AggregateInitializeReq>, DapAbort> {
         let mut processed = HashSet::with_capacity(reports.len());
@@ -361,7 +361,7 @@ impl VdafConfig {
                 task_id: task_id.clone(),
                 agg_job_id: agg_job_id.clone(),
                 agg_param: Vec::default(),
-                batch_param: BatchParameter::TimeInterval, // TODO(issue #100)
+                part_batch_sel: part_batch_sel.clone(),
                 report_shares: seq,
             },
         ))
@@ -443,7 +443,10 @@ impl VdafConfig {
         }
 
         Ok(DapHelperTransition::Continue(
-            DapHelperState { seq: states },
+            DapHelperState {
+                part_batch_sel: agg_init_req.part_batch_sel.clone(),
+                seq: states,
+            },
             AggregateResp { transitions },
         ))
     }
@@ -711,10 +714,10 @@ impl VdafConfig {
         &self,
         hpke_config: &HpkeConfig,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
         agg_share: &DapAggregateShare,
     ) -> Result<HpkeCiphertext, DapAbort> {
-        produce_encrypted_agg_share(true, hpke_config, task_id, batch_selector, agg_share)
+        produce_encrypted_agg_share(true, hpke_config, task_id, batch_sel, agg_share)
     }
 
     /// Like [`produce_leader_encrypted_agg_share`] but run by the Helper in response to an
@@ -723,10 +726,10 @@ impl VdafConfig {
         &self,
         hpke_config: &HpkeConfig,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
         agg_share: &DapAggregateShare,
     ) -> Result<HpkeCiphertext, DapAbort> {
-        produce_encrypted_agg_share(false, hpke_config, task_id, batch_selector, agg_share)
+        produce_encrypted_agg_share(false, hpke_config, task_id, batch_sel, agg_share)
     }
 
     /// Decrypt and unshard a sequence of aggregate shares. This method is run by the Collector
@@ -748,7 +751,7 @@ impl VdafConfig {
         &self,
         decrypter: &impl HpkeDecrypter<'_>,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
         report_count: u64,
         encrypted_agg_shares: Vec<HpkeCiphertext>,
     ) -> Result<DapAggregateResult, DapError> {
@@ -759,7 +762,7 @@ impl VdafConfig {
 
         let mut aad = Vec::with_capacity(40);
         task_id.encode(&mut aad);
-        batch_selector.encode(&mut aad);
+        batch_sel.encode(&mut aad);
 
         let mut agg_shares = Vec::with_capacity(encrypted_agg_shares.len());
         for (i, agg_share_ciphertext) in encrypted_agg_shares.iter().enumerate() {
@@ -797,7 +800,7 @@ fn produce_encrypted_agg_share(
     is_leader: bool,
     hpke_config: &HpkeConfig,
     task_id: &Id,
-    batch_selector: &BatchSelector,
+    batch_sel: &BatchSelector,
     agg_share: &DapAggregateShare,
 ) -> Result<HpkeCiphertext, DapAbort> {
     let agg_share_data = agg_share
@@ -819,7 +822,7 @@ fn produce_encrypted_agg_share(
     // TODO spec: Consider adding agg param to AAD.
     let mut aad = Vec::with_capacity(40);
     task_id.encode(&mut aad);
-    batch_selector.encode(&mut aad);
+    batch_sel.encode(&mut aad);
 
     let (enc, payload) = hpke_config
         .encrypt(&info, &aad, &agg_share_data)
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 25d360d..44cd97a 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -43,8 +43,8 @@ use daphne::{
     constants,
     hpke::HpkeDecrypter,
     messages::{
-        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, Id, Nonce, Report, ReportShare,
-        TransitionFailure,
+        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, Id, Nonce, PartialBatchSelector,
+        Report, ReportShare, TransitionFailure,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState, DapOutputShare,
@@ -254,7 +254,7 @@ impl<'a, D> DapAggregator<'a, BearerToken> for DaphneWorkerConfig<D> {
         batch_selector: &BatchSelector,
     ) -> std::result::Result<bool, DapError> {
         let task_config = self.try_get_task_config_for(task_id)?;
-        let batch_interval = batch_selector.unwrap_interval();
+        let batch_interval = batch_selector.unwrap_interval(); // TODO(issue #100)
 
         // Check whether the request overlaps with previous requests. This is done by
         // checking the AggregateStore and seeing whether it requests for aggregate
@@ -290,19 +290,28 @@ impl<'a, D> DapAggregator<'a, BearerToken> for DaphneWorkerConfig<D> {
         Ok(false)
     }
 
+    async fn batch_exists(
+        &self,
+        _task_id: &Id,
+        _batch_id: &Id,
+    ) -> std::result::Result<bool, DapError> {
+        panic!("TODO(issue #100)");
+    }
+
     async fn put_out_shares(
         &self,
         task_id: &Id,
+        _part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
     ) -> std::result::Result<(), DapError> {
         let task_config = self.try_get_task_config_for(task_id)?;
-
-        let agg_shares =
-            DapAggregateShare::batches_from_out_shares(out_shares, task_config.time_precision)?;
+        let part_batch_sel = PartialBatchSelector::TimeInterval; // TODO(issue #100)
+        let agg_shares = task_config.batch_span_for_out_shares(&part_batch_sel, out_shares)?;
 
         let durable = self.durable();
         let mut requests = Vec::new();
-        for (window, agg_share) in agg_shares.into_iter() {
+        for (bucket, agg_share) in agg_shares.into_iter() {
+            let window = bucket.unwrap_window(); // TODO(issue #100)
             let durable_name =
                 durable_agg_store_name(&task_config.version, &task_id.to_hex(), window);
             requests.push(durable.post::<_, ()>(
@@ -322,7 +331,7 @@ impl<'a, D> DapAggregator<'a, BearerToken> for DaphneWorkerConfig<D> {
         batch_selector: &BatchSelector,
     ) -> std::result::Result<DapAggregateShare, DapError> {
         let task_config = self.try_get_task_config_for(task_id)?;
-        let batch_interval = batch_selector.unwrap_interval();
+        let batch_interval = batch_selector.unwrap_interval(); // TODO(issue #100)
 
         let durable = self.durable();
         let mut requests = Vec::new();
@@ -358,7 +367,7 @@ impl<'a, D> DapAggregator<'a, BearerToken> for DaphneWorkerConfig<D> {
         batch_selector: &BatchSelector,
     ) -> std::result::Result<(), DapError> {
         let task_config = self.try_get_task_config_for(task_id)?;
-        let batch_interval = batch_selector.unwrap_interval();
+        let batch_interval = batch_selector.unwrap_interval(); // TODO(issue #100)
 
         // Mark reports collected.
         let durable = self.durable();
@@ -427,7 +436,7 @@ impl<'a, D> DapLeader<'a, BearerToken> for DaphneWorkerConfig<D> {
     async fn get_reports(
         &self,
         selector: &InternalAggregateInfo,
-    ) -> std::result::Result<HashMap<Id, Vec<Report>>, DapError> {
+    ) -> std::result::Result<HashMap<Id, (PartialBatchSelector, Vec<Report>)>, DapError> {
         // Read at most `selector.max_buckets` buckets from the agg job queue. The result is ordered
         // from oldest to newest.
         //
@@ -447,7 +456,7 @@ impl<'a, D> DapLeader<'a, BearerToken> for DaphneWorkerConfig<D> {
         // Drain at most `selector.max_reports` from each bucket.
         //
         // TODO Figure out if we can safely handle each bucket in parallel.
-        let mut reports_per_task: HashMap<Id, Vec<Report>> = HashMap::new();
+        let mut reports_per_task: HashMap<Id, (PartialBatchSelector, Vec<Report>)> = HashMap::new();
         for report_store_id_hex in res.into_iter() {
             let reports_from_durable: Vec<String> = self
                 .durable()
@@ -466,15 +475,16 @@ impl<'a, D> DapLeader<'a, BearerToken> for DaphneWorkerConfig<D> {
                         DapError::fatal("response from ReportStore is not valid hex")
                     })?)?;
 
-                if let Some(reports) = reports_per_task.get_mut(&report.task_id) {
+                if let Some((_par_batch_sel, reports)) = reports_per_task.get_mut(&report.task_id) {
                     reports.push(report);
                 } else {
-                    reports_per_task.insert(report.task_id.clone(), vec![report]);
+                    let par_batch_sel = PartialBatchSelector::TimeInterval; // TODO(issue #100)
+                    reports_per_task.insert(report.task_id.clone(), (par_batch_sel, vec![report]));
                 }
             }
         }
 
-        for (task_id, reports) in reports_per_task.iter() {
+        for (task_id, (_par_batch_sel, reports)) in reports_per_task.iter() {
             console_debug!(
                 "got {} reports for task {}",
                 reports.len(),
@@ -659,6 +669,7 @@ impl<'a, D> DapHelper<'a, BearerToken> for DaphneWorkerConfig<D> {
     async fn mark_aggregated(
         &self,
         task_id: &Id,
+        _part_batch_sel: &PartialBatchSelector,
         report_shares: &[ReportShare],
     ) -> std::result::Result<HashMap<Nonce, TransitionFailure>, DapError> {
         let task_config = self.try_get_task_config_for(task_id)?;


Test patch:
diff --git a/daphne/src/messages_test.rs b/daphne/src/messages_test.rs
index cacf3c6..27489dd 100644
--- a/daphne/src/messages_test.rs
+++ b/daphne/src/messages_test.rs
@@ -2,8 +2,8 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 use crate::messages::{
-    AggregateContinueReq, AggregateInitializeReq, AggregateResp, BatchParameter, Extension,
-    HpkeAeadId, HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Nonce, Report,
+    AggregateContinueReq, AggregateInitializeReq, AggregateResp, Extension, HpkeAeadId,
+    HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Nonce, PartialBatchSelector, Report,
     ReportMetadata, ReportShare, Transition, TransitionVar,
 };
 use prio::codec::{Decode, Encode};
@@ -47,7 +47,7 @@ fn read_agg_init_req() {
         task_id: Id([23; 32]),
         agg_job_id: Id([1; 32]),
         agg_param: b"this is an aggregation parameter".to_vec(),
-        batch_param: BatchParameter::FixedSize {
+        part_batch_sel: PartialBatchSelector::FixedSize {
             batch_id: Id([0; 32]),
         },
         report_shares: vec![
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 525c5df..ce499f9 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -10,12 +10,12 @@ use crate::{
     hpke::{HpkeDecrypter, HpkeReceiverConfig},
     messages::{
         AggregateContinueReq, AggregateInitializeReq, AggregateResp, AggregateShareReq,
-        AggregateShareResp, BatchParameter, BatchSelector, CollectReq, CollectResp, HpkeCiphertext,
-        HpkeKemId, Id, Interval, Query, Report, ReportShare, Time, Transition, TransitionFailure,
-        TransitionVar,
+        AggregateShareResp, BatchSelector, CollectReq, CollectResp, HpkeCiphertext, HpkeKemId, Id,
+        Interval, PartialBatchSelector, Query, Report, ReportShare, Time, Transition,
+        TransitionFailure, TransitionVar,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    testing::{AggStore, MockAggregateInfo, MockAggregator, ReportStore},
+    testing::{AggStore, DapBatchBucketOwned, MockAggregator, MockAggregatorReportSelector},
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateShare, DapCollectJob, DapGlobalConfig, DapLeaderTransition,
     DapMeasurement, DapQueryConfig, DapRequest, DapTaskConfig, DapVersion, Prio3Config, VdafConfig,
@@ -25,8 +25,7 @@ use matchit::Router;
 use prio::codec::{Decode, Encode};
 use rand::{thread_rng, Rng};
 use std::{
-    collections::{HashMap, HashSet, VecDeque},
-    ops::DerefMut,
+    collections::HashMap,
     sync::{Arc, Mutex},
     time::SystemTime,
     vec,
@@ -37,7 +36,8 @@ macro_rules! get_reports {
     ($leader:expr, $selector:expr) => {{
         let reports_per_task = $leader.get_reports($selector).await.unwrap();
         assert_eq!(reports_per_task.len(), 1);
-        reports_per_task.into_iter().next().unwrap()
+        let (task_id, (part_batch_sel, reports)) = reports_per_task.into_iter().next().unwrap();
+        (task_id, part_batch_sel, reports)
     }};
 }
 
@@ -45,12 +45,8 @@ struct Test {
     now: Time,
     leader: MockAggregator,
     helper: MockAggregator,
-    #[allow(dead_code)] // TODO(issue #100) Remove
-    collector_hpke_receiver_config: HpkeReceiverConfig,
-    #[allow(dead_code)] // TODO(issue #100) Remove
     collector_token: BearerToken,
     time_interval_task_id: Id,
-    #[allow(dead_code)] // TODO(issue #100) Remove
     fixed_size_task_id: Id,
 }
 
@@ -156,7 +152,6 @@ impl Test {
             now,
             leader,
             helper,
-            collector_hpke_receiver_config,
             collector_token,
             time_interval_task_id,
             fixed_size_task_id,
@@ -164,8 +159,7 @@ impl Test {
     }
 
     fn gen_test_upload_req(&self, report: Report) -> DapRequest<BearerToken> {
-        let task_id = &self.time_interval_task_id;
-        let task_config = self.leader.tasks.get(task_id).unwrap();
+        let task_config = self.leader.tasks.get(&report.task_id).unwrap();
         let version = task_config.version.clone();
 
         DapRequest {
@@ -184,8 +178,8 @@ impl Test {
     ) -> DapRequest<BearerToken> {
         let mut rng = thread_rng();
         let task_config = self.leader.tasks.get(task_id).unwrap();
-        let batch_param = match task_config.query {
-            DapQueryConfig::TimeInterval { .. } => BatchParameter::TimeInterval,
+        let part_batch_sel = match task_config.query {
+            DapQueryConfig::TimeInterval { .. } => PartialBatchSelector::TimeInterval,
             _ => panic!("TODO(issue #100)"),
         };
 
@@ -197,7 +191,7 @@ impl Test {
                 task_id: task_id.clone(),
                 agg_job_id: Id(rng.gen()),
                 agg_param: Vec::default(),
-                batch_param,
+                part_batch_sel,
                 report_shares,
             },
             task_config.helper_url.join("aggregate").unwrap(),
@@ -241,7 +235,7 @@ impl Test {
             MEDIA_TYPE_AGG_SHARE_REQ,
             AggregateShareReq {
                 task_id: task_id.clone(),
-                batch_selector: BatchSelector::default(),
+                batch_sel: BatchSelector::default(),
                 agg_param: Vec::default(),
                 report_count,
                 checksum,
@@ -282,11 +276,8 @@ impl Test {
         let task_config = wrapped.as_ref().unwrap();
 
         // Leader: Store received report to ReportStore.
-        let selector = &MockAggregateInfo {
-            task_id: task_id.clone(),
-            agg_rate: 1,
-        };
-        let (task_id, reports) = get_reports!(self.leader, selector);
+        let report_sel = MockAggregatorReportSelector(task_id.clone());
+        let (task_id, part_batch_sel, reports) = get_reports!(self.leader, &report_sel);
 
         // Leader: Consume report share.
         let mut rng = thread_rng();
@@ -298,6 +289,7 @@ impl Test {
                 &task_config.vdaf_verify_key,
                 &task_id,
                 &agg_job_id,
+                &part_batch_sel,
                 reports,
             )
             .await?;
@@ -344,7 +336,9 @@ impl Test {
         let out_shares = task_config
             .vdaf
             .handle_final_agg_resp(leader_uncommitted, agg_resp)?;
-        self.leader.put_out_shares(&task_id, out_shares).await?;
+        self.leader
+            .put_out_shares(&task_id, &part_batch_sel, out_shares)
+            .await?;
 
         Ok(())
     }
@@ -387,7 +381,7 @@ impl Test {
         // Leader->Helper: HTTP POST /aggregate_share
         let agg_share_req = AggregateShareReq {
             task_id: collect_req.task_id.clone(),
-            batch_selector: collect_req.query.clone(),
+            batch_sel: collect_req.query.clone(),
             agg_param: collect_req.agg_param.clone(),
             report_count: leader_agg_share.report_count,
             checksum: leader_agg_share.checksum,
@@ -408,6 +402,7 @@ impl Test {
 
         // Leader: Complete the collect job.
         let collect_resp = CollectResp {
+            part_batch_sel: collect_req.query.clone().into(),
             report_count: leader_agg_share.report_count,
             encrypted_agg_shares: vec![leader_enc_agg_share, agg_share_resp.encrypted_agg_share],
         };
@@ -415,7 +410,7 @@ impl Test {
             .finish_collect_job(task_id, collect_id, &collect_resp)
             .await?;
         self.leader
-            .mark_collected(task_id, &agg_share_req.batch_selector)
+            .mark_collected(task_id, &agg_share_req.batch_sel)
             .await?;
 
         // Collector: Poll the collect job.
@@ -468,7 +463,7 @@ impl Test {
 
 // Test that the Helper properly handles the batch parameter in the AggregateInitializeReq.
 #[tokio::test]
-async fn http_post_aggregate_invalid_batch_selector() {
+async fn http_post_aggregate_invalid_batch_sel() {
     let mut rng = thread_rng();
     let t = Test::new();
     let task_id = &t.time_interval_task_id;
@@ -484,7 +479,7 @@ async fn http_post_aggregate_invalid_batch_selector() {
                 task_id: task_id.clone(),
                 agg_job_id: Id(rng.gen()),
                 agg_param: Vec::default(),
-                batch_param: BatchParameter::FixedSize {
+                part_batch_sel: PartialBatchSelector::FixedSize {
                     batch_id: Id(rng.gen()),
                 },
                 report_shares: Vec::default(),
@@ -606,21 +601,20 @@ async fn http_post_aggregate_share_unauthorized_request() {
 
 // Test that the Helper handles the batch selector sent from the Leader properly.
 #[tokio::test]
-async fn http_post_aggregate_share_invalid_batch_selector() {
+async fn http_post_aggregate_share_invalid_batch_sel() {
     let mut rng = thread_rng();
     let t = Test::new();
-    let task_id = &t.time_interval_task_id;
-    let task_config = t.leader.tasks.get(task_id).unwrap();
 
     // Helper expects "time_interval" query, but Leader sent "fixed_size".
+    let task_config = t.leader.tasks.get(&t.time_interval_task_id).unwrap();
     let req = t
         .leader_authorized_req(
-            task_id,
+            &t.time_interval_task_id,
             task_config.version,
             MEDIA_TYPE_AGG_SHARE_REQ,
             AggregateShareReq {
-                task_id: task_id.clone(),
-                batch_selector: BatchSelector::FixedSize {
+                task_id: t.time_interval_task_id.clone(),
+                batch_sel: BatchSelector::FixedSize {
                     batch_id: Id(rng.gen()),
                 },
                 agg_param: Vec::default(),
@@ -634,6 +628,30 @@ async fn http_post_aggregate_share_invalid_batch_selector() {
         t.helper.http_post_aggregate_share(&req).await.unwrap_err(),
         DapAbort::QueryMismatch
     );
+
+    // Leader sends aggregate share request for unrecognized batch ID.
+    let task_config = t.leader.tasks.get(&t.fixed_size_task_id).unwrap();
+    let req = t
+        .leader_authorized_req(
+            &t.fixed_size_task_id,
+            task_config.version,
+            MEDIA_TYPE_AGG_SHARE_REQ,
+            AggregateShareReq {
+                task_id: t.fixed_size_task_id.clone(),
+                batch_sel: BatchSelector::FixedSize {
+                    batch_id: Id(rng.gen()), // Unrecognized batch ID
+                },
+                agg_param: Vec::default(),
+                report_count: 0,
+                checksum: [0; 32],
+            },
+            task_config.helper_url.join("aggregate_share").unwrap(),
+        )
+        .await;
+    assert_matches!(
+        t.helper.http_post_aggregate_share(&req).await.unwrap_err(),
+        DapAbort::BatchInvalid
+    );
 }
 
 #[tokio::test]
@@ -738,19 +756,16 @@ async fn http_post_aggregate_failure_report_replayed() {
     }];
     let req = t.gen_test_agg_init_req(task_id, report_shares).await;
 
-    // The following scope is to ensure that the report_store lock is released before http_post_aggregate is called.
+    // 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 report_store_mutex_guard = t.helper.report_store.lock().expect("lock() failed");
-        let report_store = report_store_mutex_guard.deref_mut();
-        let mut processed = HashSet::new();
-        processed.insert(report.metadata.nonce.clone());
-        report_store.insert(
-            task_id.clone(),
-            ReportStore {
-                pending: VecDeque::new(),
-                processed,
-            },
-        );
+        let mut guard = t
+            .helper
+            .report_store
+            .lock()
+            .expect("report_store: failed to lock");
+        let report_store = guard.entry(task_id.clone()).or_default();
+        report_store.processed.insert(report.metadata.nonce.clone());
     }
 
     // Get AggregateResp and then extract the transition data from inside.
@@ -792,11 +807,8 @@ async fn http_post_aggregate_failure_batch_collected() {
         let agg_store = guard.entry(task_id.clone()).or_default();
 
         agg_store.insert(
-            BatchSelector::TimeInterval {
-                batch_interval: Interval {
-                    start: task_config.truncate_time(t.now),
-                    duration: task_config.time_precision,
-                },
+            DapBatchBucketOwned::TimeInterval {
+                batch_window: task_config.truncate_time(t.now),
             },
             AggStore {
                 agg_share: DapAggregateShare::default(),
@@ -861,11 +873,18 @@ async fn http_post_aggregate_fail_send_cont_req() {
 async fn http_post_upload_fail_send_invalid_report() {
     let t = Test::new();
     let task_id = &t.time_interval_task_id;
+    let task_config = t.leader.tasks.get(task_id).unwrap();
 
     // Construct a report payload with an invalid task ID.
     let mut report_empty_task_id = t.gen_test_report(task_id).await;
     report_empty_task_id.task_id = Id([0; 32]);
-    let req = t.gen_test_upload_req(report_empty_task_id);
+    let req = DapRequest {
+        version: task_config.version,
+        media_type: Some(MEDIA_TYPE_REPORT),
+        payload: report_empty_task_id.get_encoded(),
+        url: task_config.leader_url.join("upload").unwrap(),
+        sender_auth: None,
+    };
 
     // Expect failure due to invalid task ID in report.
     assert_matches!(
@@ -926,18 +945,15 @@ async fn get_reports_empty_response() {
     // Get one report. This should return with the report that was uploaded earlier.
     // We also check that the task ID associated to the report is the same one we
     // requested.
-    let selector = &MockAggregateInfo {
-        task_id: task_id.clone(),
-        agg_rate: 1,
-    };
-    let (returned_task_id, reports) = get_reports!(t.leader, selector);
+    let report_sel = MockAggregatorReportSelector(task_id.clone());
+    let (returned_task_id, _part_batch_sel, reports) = get_reports!(t.leader, &report_sel);
     assert_eq!(reports.len(), 1);
     assert_eq!(&returned_task_id, task_id);
 
     // Try to get another report. This should not return an error, but simply
     // an empty vector, as we drained the ReportStore above. The task ID
     // associated to the report should be the same one we requested.
-    let (returned_task_id, reports) = get_reports!(t.leader, selector);
+    let (returned_task_id, _part_batch_sel, reports) = get_reports!(t.leader, &report_sel);
     assert_eq!(reports.len(), 0);
     assert_eq!(&returned_task_id, task_id);
 }
@@ -979,6 +995,7 @@ async fn poll_collect_job_test_results() {
     let resp = t.leader.get_pending_collect_jobs().await.unwrap();
     let (collect_id, _collect_req) = &resp[0];
     let collect_resp = CollectResp {
+        part_batch_sel: PartialBatchSelector::TimeInterval,
         report_count: 0,
         encrypted_agg_shares: Vec::default(),
     };
@@ -1206,16 +1223,15 @@ async fn http_post_collect_success() {
 async fn http_post_collect_invalid_query() {
     let mut rng = thread_rng();
     let t = Test::new();
-    let task_id = &t.time_interval_task_id;
-    let task_config = t.leader.tasks.get(task_id).unwrap();
 
     // Leader expects "time_interval" query, but Collector sent "fixed_size".
+    let task_config = t.leader.tasks.get(&t.time_interval_task_id).unwrap();
     let req = t
         .collector_authorized_req(
             task_config.version,
             MEDIA_TYPE_COLLECT_REQ,
             CollectReq {
-                task_id: task_id.clone(),
+                task_id: t.time_interval_task_id.clone(),
                 query: Query::FixedSize {
                     batch_id: Id(rng.gen()),
                 },
@@ -1228,6 +1244,27 @@ async fn http_post_collect_invalid_query() {
         t.leader.http_post_collect(&req).await.unwrap_err(),
         DapAbort::QueryMismatch
     );
+
+    // Collector indicates unrecognized batch ID.
+    let task_config = t.leader.tasks.get(&t.fixed_size_task_id).unwrap();
+    let req = t
+        .collector_authorized_req(
+            task_config.version,
+            MEDIA_TYPE_COLLECT_REQ,
+            CollectReq {
+                task_id: t.fixed_size_task_id.clone(),
+                query: Query::FixedSize {
+                    batch_id: Id(rng.gen()), // Unrecognized batch ID
+                },
+                agg_param: Vec::default(),
+            },
+            task_config.leader_url.join("collect").unwrap(),
+        )
+        .await;
+    assert_matches!(
+        t.leader.http_post_collect(&req).await.unwrap_err(),
+        DapAbort::BatchInvalid
+    );
 }
 
 // Test HTTP POST requests with a wrong DAP version.
@@ -1262,9 +1299,8 @@ async fn successful_http_post_upload() {
         .expect("upload failed unexpectedly");
 }
 
-// Test the end-to-end protocol.
 #[tokio::test]
-async fn e2e() {
+async fn e2e_time_interval() {
     let t = Test::new();
     let task_id = &t.time_interval_task_id;
 
@@ -1286,3 +1322,24 @@ async fn e2e() {
         .query_for_current_batch_window(t.now);
     t.run_col_job(task_id, &query).await.unwrap();
 }
+
+#[tokio::test]
+async fn e2e_fixed_size() {
+    let t = Test::new();
+    let task_id = &t.fixed_size_task_id;
+
+    let report = t.gen_test_report(task_id).await;
+    let req = t.gen_test_upload_req(report);
+
+    // Client: Send upload request to Leader.
+    t.leader.http_post_upload(&req).await.unwrap();
+
+    // Leader: Run aggregation job.
+    t.run_agg_job(task_id).await.unwrap();
+
+    // Collector: Create collection job and poll result.
+    let query = Query::FixedSize {
+        batch_id: t.leader.current_batch(task_id).unwrap(),
+    };
+    t.run_col_job(task_id, &query).await.unwrap();
+}
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 0dd7201..a35c282 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -7,17 +7,19 @@ use crate::{
     auth::{BearerToken, BearerTokenProvider},
     hpke::{HpkeDecrypter, HpkeReceiverConfig},
     messages::{
-        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, HpkeConfig, Id, Interval, Nonce,
-        Report, ReportMetadata, ReportShare, Time, TransitionFailure,
+        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, HpkeConfig, Id, Nonce,
+        PartialBatchSelector, Report, ReportMetadata, ReportShare, Time, TransitionFailure,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
-    DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
+    DapAbort, DapAggregateShare, DapBatchBucket, DapCollectJob, DapError, DapGlobalConfig,
+    DapHelperState, DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
 };
+use assert_matches::assert_matches;
 use async_trait::async_trait;
 use rand::{thread_rng, Rng};
 use serde::{Deserialize, Serialize};
 use std::{
+    borrow::Borrow,
     collections::{HashMap, HashSet, VecDeque},
     hash::Hash,
     ops::DerefMut,
@@ -26,11 +28,39 @@ use std::{
 };
 use url::Url;
 
-pub(crate) struct MockAggregateInfo {
-    pub(crate) task_id: Id,
-    pub(crate) agg_rate: u64,
+#[derive(Eq, Hash, PartialEq)]
+pub(crate) enum DapBatchBucketOwned {
+    FixedSize { batch_id: Id },
+    TimeInterval { batch_window: Time },
 }
 
+impl From<DapBatchBucketOwned> for PartialBatchSelector {
+    fn from(bucket: DapBatchBucketOwned) -> Self {
+        match bucket {
+            DapBatchBucketOwned::FixedSize { batch_id } => Self::FixedSize { batch_id },
+            DapBatchBucketOwned::TimeInterval { .. } => Self::TimeInterval,
+        }
+    }
+}
+
+impl<'a> DapBatchBucket<'a> {
+    // TODO(cjpatton) Figure out how to use `ToOwned` properly. The lifetime parameter causes
+    // confusion for the compiler for implementing `Borrow`. The goal is to avoid cloning the
+    // bucket each time we need to check if it exists in the set.
+    pub(crate) fn to_owned_bucket(&self) -> DapBatchBucketOwned {
+        match self {
+            Self::FixedSize { batch_id } => DapBatchBucketOwned::FixedSize {
+                batch_id: (*batch_id).clone(),
+            },
+            Self::TimeInterval { batch_window } => DapBatchBucketOwned::TimeInterval {
+                batch_window: *batch_window,
+            },
+        }
+    }
+}
+
+pub(crate) struct MockAggregatorReportSelector(pub(crate) Id);
+
 #[allow(dead_code)]
 pub(crate) struct MockAggregator {
     pub(crate) now: Time,
@@ -42,7 +72,7 @@ pub(crate) struct MockAggregator {
     pub(crate) report_store: Arc<Mutex<HashMap<Id, ReportStore>>>,
     pub(crate) leader_state_store: Arc<Mutex<HashMap<Id, LeaderState>>>,
     pub(crate) helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapHelperState>>>,
-    pub(crate) agg_store: Arc<Mutex<HashMap<Id, HashMap<BatchSelector, AggStore>>>>,
+    pub(crate) agg_store: Arc<Mutex<HashMap<Id, HashMap<DapBatchBucketOwned, AggStore>>>>,
 }
 
 #[allow(dead_code)]
@@ -53,26 +83,14 @@ impl MockAggregator {
     async fn check_report_early_fail(
         &self,
         task_id: &Id,
+        bucket: &DapBatchBucketOwned,
         metadata: &ReportMetadata,
     ) -> Option<TransitionFailure> {
-        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
-        if matches!(task_config.query, DapQueryConfig::FixedSize { .. }) {
-            panic!("TODO(issue #100)");
-        }
-
-        let batch_selector = BatchSelector::TimeInterval {
-            batch_interval: Interval {
-                start: task_config.truncate_time(metadata.time),
-                duration: task_config.time_precision,
-            },
-        };
-
         // Check AggStateStore to see whether the report is part of a batch that has already
         // been collected.
         let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
         let agg_store = guard.entry(task_id.clone()).or_default();
-        if matches!(agg_store.get(&batch_selector), Some(inner_agg_store) if inner_agg_store.collected)
-        {
+        if matches!(agg_store.get(bucket), Some(inner_agg_store) if inner_agg_store.collected) {
             return Some(TransitionFailure::BatchCollected);
         }
 
@@ -81,12 +99,8 @@ impl MockAggregator {
             .report_store
             .lock()
             .expect("report_store: failed to lock");
-        if guard
-            .entry(task_id.clone())
-            .or_default()
-            .processed
-            .contains(&metadata.nonce)
-        {
+        let report_store = guard.entry(task_id.clone()).or_default();
+        if report_store.processed.contains(&metadata.nonce) {
             return Some(TransitionFailure::ReportReplayed);
         }
 
@@ -98,6 +112,76 @@ impl MockAggregator {
             .iter()
             .find(|&hpke_receiver_config| hpke_config_id == hpke_receiver_config.config.id)
     }
+
+    /// Assign the report to a bucket.
+    ///
+    /// TODO(cjpatton) Figure out if we can avoid returning and owned thing here.
+    fn assign_report_to_bucket(&self, report: &Report) -> Option<DapBatchBucketOwned> {
+        let mut rng = thread_rng();
+        let task_config = self
+            .tasks
+            .get(&report.task_id)
+            .expect("tasks: unrecognized task");
+
+        match task_config.query {
+            // For fixed-size queries, the bucket corresponds to a single batch.
+            DapQueryConfig::FixedSize {
+                min_batch_size,
+                max_batch_size: _,
+            } => {
+                let mut guard = self
+                    .leader_state_store
+                    .lock()
+                    .expect("leader_state_store: failed to lock");
+                let leader_state_store = guard.entry(report.task_id.clone()).or_default();
+
+                // Assign the report to the first unsaturated batch.
+                for (batch_id, report_count) in leader_state_store.batch_queue.iter_mut() {
+                    if *report_count < min_batch_size {
+                        *report_count += 1;
+                        return Some(DapBatchBucketOwned::FixedSize {
+                            batch_id: batch_id.clone(),
+                        });
+                    }
+                }
+
+                // No unsaturated batch exists, so create a new batch.
+                let batch_id = Id(rng.gen());
+                leader_state_store
+                    .batch_queue
+                    .push_back((batch_id.clone(), 1));
+                Some(DapBatchBucketOwned::FixedSize { batch_id })
+            }
+
+            // For time-interval queries, the bucket is the batch window computed by truncating the
+            // report timestamp.
+            DapQueryConfig::TimeInterval { .. } => Some(DapBatchBucketOwned::TimeInterval {
+                batch_window: task_config.truncate_time(report.metadata.time),
+            }),
+        }
+    }
+
+    /// Return the ID of the batch currently being filled with reports. Panics unless the task is
+    /// configured for fixed-size queries.
+    pub(crate) fn current_batch(&self, task_id: &Id) -> Option<Id> {
+        // Calling current_batch() is only well-defined for fixed-size tasks.
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
+        assert_matches!(task_config.query, DapQueryConfig::FixedSize { .. });
+
+        let guard = self
+            .leader_state_store
+            .lock()
+            .expect("leader_state_store: failed to lock");
+        let leader_state_store = guard
+            .get(task_id)
+            .expect("leader_state_store: unrecognized task");
+
+        leader_state_store
+            .batch_queue
+            .front()
+            .cloned() // TODO(cjpatton) Avoid clone by returning MutexGuard
+            .map(|(batch_id, _report_count)| batch_id)
+    }
 }
 
 #[async_trait(?Send)]
@@ -199,7 +283,7 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
         Ok(self.tasks.get(task_id))
     }
 
-    fn get_current_time(&self) -> u64 {
+    fn get_current_time(&self) -> Time {
         SystemTime::now()
             .duration_since(SystemTime::UNIX_EPOCH)
             .unwrap()
@@ -209,56 +293,58 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
     async fn is_batch_overlapping(
         &self,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
     ) -> Result<bool, DapError> {
         let guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
         let agg_store = if let Some(agg_store) = guard.get(task_id) {
             agg_store
         } else {
             return Ok(false);
         };
 
-        let batch_interval = batch_selector.unwrap_interval();
-        for (inner_selector, inner_agg_store) in agg_store.iter() {
-            let inner_interval = inner_selector.unwrap_interval();
-            if batch_interval.start <= inner_interval.start
-                && batch_interval.end() > inner_interval.start
-                && inner_agg_store.collected
-            {
-                return Ok(true);
+        for bucket in task_config.batch_span_for_sel(batch_sel)? {
+            if let Some(inner_agg_store) = agg_store.get(&bucket.to_owned_bucket()) {
+                if inner_agg_store.collected {
+                    return Ok(true);
+                }
             }
         }
 
         Ok(false)
     }
 
+    async fn batch_exists(&self, task_id: &Id, batch_id: &Id) -> Result<bool, DapError> {
+        let guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        if let Some(agg_store) = guard.get(task_id) {
+            Ok(agg_store
+                .get(&DapBatchBucketOwned::FixedSize {
+                    batch_id: batch_id.clone(),
+                })
+                .is_some())
+        } else {
+            Ok(false)
+        }
+    }
+
     async fn put_out_shares(
         &self,
         task_id: &Id,
+        part_batch_sel: &PartialBatchSelector,
         out_shares: Vec<DapOutputShare>,
     ) -> Result<(), DapError> {
         let task_config = self
             .get_task_config_for(task_id)
             .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
-        if matches!(task_config.query, DapQueryConfig::FixedSize { .. }) {
-            panic!("TODO(issue #100)");
-        }
-
-        let agg_shares =
-            DapAggregateShare::batches_from_out_shares(out_shares, task_config.time_precision)?;
 
         let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
         let agg_store = guard.entry(task_id.clone()).or_default();
-        for (window, agg_share_delta) in agg_shares.into_iter() {
-            let batch_selector = BatchSelector::TimeInterval {
-                batch_interval: Interval {
-                    start: window,
-                    duration: task_config.time_precision,
-                },
-            };
-
-            let inner_agg_store = agg_store.entry(batch_selector).or_default();
+        for (bucket, agg_share_delta) 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)?;
         }
 
@@ -268,19 +354,16 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
     async fn get_agg_share(
         &self,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
     ) -> Result<DapAggregateShare, DapError> {
-        let guard = self.agg_store.lock().expect("agg_store: failed to lock");
-        let agg_store = guard.get(task_id).expect("agg_store: unrecognized task");
+        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        let agg_store = guard.entry(task_id.clone()).or_default();
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
 
         // Fetch aggregate shares.
         let mut agg_share = DapAggregateShare::default();
-        let batch_interval = batch_selector.unwrap_interval();
-        for (inner_selector, inner_agg_store) in agg_store.iter() {
-            let inner_interval = inner_selector.unwrap_interval();
-            if batch_interval.start <= inner_interval.start
-                && batch_interval.end() > inner_interval.start
-            {
+        for bucket in task_config.batch_span_for_sel(batch_sel)? {
+            if let Some(inner_agg_store) = agg_store.get(&bucket.to_owned_bucket()) {
                 if inner_agg_store.collected {
                     return Err(DapError::Abort(DapAbort::BatchOverlap));
                 } else {
@@ -295,17 +378,14 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
     async fn mark_collected(
         &self,
         task_id: &Id,
-        batch_selector: &BatchSelector,
+        batch_sel: &BatchSelector,
     ) -> Result<(), DapError> {
         let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
         let agg_store = guard.entry(task_id.clone()).or_default();
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
 
-        let batch_interval = batch_selector.unwrap_interval();
-        for (inner_selector, inner_agg_store) in agg_store.iter_mut() {
-            let inner_interval = inner_selector.unwrap_interval();
-            if batch_interval.start <= inner_interval.start
-                && batch_interval.end() > inner_interval.start
-            {
+        for bucket in task_config.batch_span_for_sel(batch_sel)? {
+            if let Some(inner_agg_store) = agg_store.get_mut(&bucket.to_owned_bucket()) {
                 inner_agg_store.collected = true;
             }
         }
@@ -319,19 +399,24 @@ impl<'a> DapHelper<'a, BearerToken> for MockAggregator {
     async fn mark_aggregated(
         &self,
         task_id: &Id,
+        part_batch_sel: &PartialBatchSelector,
         report_shares: &[ReportShare],
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError> {
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
         let mut early_fails = HashMap::new();
         for report_share in report_shares.iter() {
+            let bucket =
+                task_config.batch_bucket_for_meta(&report_share.metadata, part_batch_sel)?;
+
             // Check whether Report has been collected or replayed.
             if let Some(transition_failure) = self
-                .check_report_early_fail(task_id, &report_share.metadata)
+                .check_report_early_fail(task_id, &bucket.to_owned_bucket(), &report_share.metadata)
                 .await
             {
                 early_fails.insert(report_share.metadata.nonce.clone(), transition_failure);
             };
 
-            // Mark Report processed.
+            // Mark report processed.
             let mut guard = self
                 .report_store
                 .lock()
@@ -407,21 +492,16 @@ impl<'a> DapHelper<'a, BearerToken> for MockAggregator {
 
 #[async_trait(?Send)]
 impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
-    type ReportSelector = MockAggregateInfo;
+    type ReportSelector = MockAggregatorReportSelector;
 
     async fn put_report(&self, report: &Report) -> Result<(), DapError> {
-        let task_id = &report.task_id;
-        let task_config = self
-            .get_task_config_for(task_id)
-            .await?
-            .ok_or_else(|| DapError::fatal("task not found"))?;
-        if matches!(task_config.query, DapQueryConfig::FixedSize { .. }) {
-            panic!("TODO(issue #100)");
-        }
+        let bucket = self
+            .assign_report_to_bucket(report)
+            .expect("could not determine batch for report");
 
         // Check whether Report has been collected or replayed.
         if let Some(transition_failure) = self
-            .check_report_early_fail(task_id, &report.metadata)
+            .check_report_early_fail(&report.task_id, bucket.borrow(), &report.metadata)
             .await
         {
             return Err(DapError::Transition(transition_failure));
@@ -432,49 +512,60 @@ impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
             .report_store
             .lock()
             .expect("report_store: failed to lock");
-        let report_store = guard.entry(task_id.clone()).or_default();
-        report_store.pending.push_back(report.clone());
+        let queue = guard
+            .get_mut(&report.task_id)
+            .expect("report_store: unrecognized task")
+            .pending
+            .entry(bucket)
+            .or_default();
+        queue.push_back(report.clone());
         Ok(())
     }
 
     async fn get_reports(
         &self,
-        selector: &MockAggregateInfo,
-    ) -> Result<HashMap<Id, Vec<Report>>, DapError> {
-        // Lock report_store.
-        let agg_rate = selector
-            .agg_rate
-            .try_into()
-            .expect("agg_rate is larger than usize");
-        let mut reports = Vec::with_capacity(agg_rate);
-        let mut report_store_mutex_guard = self
+        report_sel: &MockAggregatorReportSelector,
+    ) -> Result<HashMap<Id, (PartialBatchSelector, Vec<Report>)>, DapError> {
+        let mut guard = self
             .report_store
             .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let report_store = report_store_mutex_guard.deref_mut();
-
-        // Fetch reports.
-        for (inner_task_id, store) in report_store.iter_mut() {
-            if &selector.task_id == inner_task_id {
-                let num_reports_remaining = agg_rate - reports.len();
-                let num_reports_drained = std::cmp::min(num_reports_remaining, store.pending.len());
-                let mut reports_drained: Vec<Report> =
-                    store.pending.drain(..num_reports_drained).collect();
-                if reports_drained.len() + reports.len() > agg_rate {
-                    return Err(DapError::fatal(
-                        "number of reports received from report_store exceeds the number requested",
-                    ));
-                }
-
-                reports.append(&mut reports_drained);
+            .expect("report_store: failed to lock");
+        let task_id = &report_sel.0;
+        let task_config = self.tasks.get(task_id).expect("tasks: unrecognized task");
+        let report_store = guard.entry(task_id.clone()).or_default();
 
-                if reports.len() == agg_rate {
-                    break;
+        // For the task indicated by the report selector, choose a single report to aggregate.
+        match task_config.query {
+            DapQueryConfig::TimeInterval { .. } => {
+                // Aggregate reports in any order.
+                let mut reports = Vec::new();
+                for (_bucket, queue) in report_store.pending.iter_mut() {
+                    if !queue.is_empty() {
+                        reports.append(&mut queue.drain(..1).collect());
+                        break;
+                    }
                 }
+                return Ok(HashMap::from([(
+                    task_id.clone(),
+                    (PartialBatchSelector::TimeInterval, reports),
+                )]));
+            }
+            DapQueryConfig::FixedSize { .. } => {
+                // Drain the batch that is being filled.
+                let bucket = if let Some(batch_id) = self.current_batch(task_id) {
+                    DapBatchBucketOwned::FixedSize { batch_id }
+                } else {
+                    return Ok(HashMap::default());
+                };
+
+                let queue = report_store
+                    .pending
+                    .get_mut(&bucket)
+                    .expect("report_store: unknown bucket");
+                let reports = queue.drain(..1).collect();
+                return Ok(HashMap::from([(task_id.clone(), (bucket.into(), reports))]));
             }
         }
-
-        Ok(HashMap::from([(selector.task_id.clone(), reports)]))
     }
 
     // Called after receiving a CollectReq from Collector.
@@ -505,7 +596,7 @@ impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
         // Store Collect ID and CollectReq into LeaderState.
         let leader_state = leader_state_store
             .entry(collect_req.task_id.clone())
-            .or_insert_with(LeaderState::new);
+            .or_default();
         leader_state.collect_ids.push_back(collect_id.clone());
         let collect_job_state = CollectJobState::Pending(collect_req.clone());
         leader_state
@@ -562,7 +653,6 @@ impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
         Ok(res)
     }
 
-    // Called after finishing aggregation job to put resuts into LeaderState.
     async fn finish_collect_job(
         &self,
         task_id: &Id,
@@ -583,6 +673,13 @@ impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
             .get_mut(collect_id)
             .ok_or_else(|| DapError::fatal("collect job not found for collect_id"))?;
 
+        // Remove the batch from the batch queue.
+        if let PartialBatchSelector::FixedSize { ref batch_id } = collect_resp.part_batch_sel {
+            leader_state
+                .batch_queue
+                .retain(|(id, _report_count)| id != batch_id);
+        }
+
         match collect_job {
             CollectJobState::Pending(_) => {
                 // Mark collect job as Processed.
@@ -619,7 +716,7 @@ pub(crate) struct HelperStateInfo {
 /// Stores the reports received from Clients.
 #[derive(Default)]
 pub(crate) struct ReportStore {
-    pub(crate) pending: VecDeque<Report>,
+    pub(crate) pending: HashMap<DapBatchBucketOwned, VecDeque<Report>>,
     pub(crate) processed: HashSet<Nonce>,
 }
 
@@ -632,18 +729,11 @@ pub(crate) enum CollectJobState {
 /// LeaderState keeps track of the following:
 /// * Collect IDs in their order of arrival.
 /// * The state of the collect job associated to the Collect ID.
+#[derive(Default)]
 pub(crate) struct LeaderState {
     collect_ids: VecDeque<Id>,
     collect_jobs: HashMap<Id, CollectJobState>,
-}
-
-impl LeaderState {
-    pub(crate) fn new() -> LeaderState {
-        Self {
-            collect_ids: VecDeque::default(),
-            collect_jobs: HashMap::default(),
-        }
-    }
+    batch_queue: VecDeque<(Id, u64)>, // Batch ID, batch size
 }
 
 /// AggStore keeps track of the following:
diff --git a/daphne/src/vdaf/mod_test.rs b/daphne/src/vdaf/mod_test.rs
index 0d8e0cd..9a395c4 100644
--- a/daphne/src/vdaf/mod_test.rs
+++ b/daphne/src/vdaf/mod_test.rs
@@ -5,8 +5,8 @@ use crate::{
     hpke::HpkeReceiverConfig,
     messages::{
         AggregateContinueReq, AggregateInitializeReq, AggregateResp, BatchSelector, HpkeAeadId,
-        HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Interval, Nonce, Report, Transition,
-        TransitionFailure, TransitionVar,
+        HpkeCiphertext, HpkeConfig, HpkeKdfId, HpkeKemId, Id, Interval, Nonce,
+        PartialBatchSelector, Report, Transition, TransitionFailure, TransitionVar,
     },
     DapAbort, DapAggregateResult, DapAggregateShare, DapError, DapHelperState, DapHelperTransition,
     DapLeaderState, DapLeaderTransition, DapLeaderUncommitted, DapMeasurement, DapOutputShare,
@@ -628,6 +628,7 @@ impl<'a> Test<'a> {
                 &self.vdaf_verify_key,
                 &self.task_id,
                 &self.agg_job_id,
+                &PartialBatchSelector::TimeInterval,
                 reports,
             )
             .await
@@ -795,18 +796,12 @@ impl<'a> Test<'a> {
         let report_count = u64::try_from(leader_out_shares.len()).unwrap();
 
         // Leader: Aggregation
-        let leader_agg_shares =
-            DapAggregateShare::batches_from_out_shares(leader_out_shares, 3600).unwrap();
-        assert_eq!(leader_agg_shares.len(), 1);
-        let (_batch_window, leader_agg_share) = leader_agg_shares.into_iter().next().unwrap();
+        let leader_agg_share = DapAggregateShare::try_from_out_shares(leader_out_shares).unwrap();
         let leader_encrypted_agg_share =
             self.produce_leader_encrypted_agg_share(&batch_selector, &leader_agg_share);
 
         // Helper: Aggregation
-        let helper_agg_shares =
-            DapAggregateShare::batches_from_out_shares(helper_out_shares, 3600).unwrap();
-        assert_eq!(helper_agg_shares.len(), 1);
-        let (_batch_window, helper_agg_share) = helper_agg_shares.into_iter().next().unwrap();
+        let helper_agg_share = DapAggregateShare::try_from_out_shares(helper_out_shares).unwrap();
         let helper_encrypted_agg_share =
             self.produce_helper_encrypted_agg_share(&batch_selector, &helper_agg_share);