Repository: cloudflare/daphne
Pull request #132 (2022-10-03 16:39:33)

Prepare `MockAggregator` for fixed-size query tasks
Partially addresses #100.
Based on #131 (merge that first).

Extend `MockAggregator`'s data stores so that data can be bucketed by batch ID and not just by batch window. Along the way, simplify some of the mutex code.


Patch:
diff --git a/daphne/src/lib.rs b/daphne/src/lib.rs
index 7814eb2..1cb3415 100644
--- a/daphne/src/lib.rs
+++ b/daphne/src/lib.rs
@@ -404,10 +404,10 @@ pub struct DapTaskConfig {
     pub(crate) collector_hpke_config: HpkeConfig,
 }
 
-#[cfg(test)]
 impl DapTaskConfig {
     /// Convert at timestamp `now` into an [`Interval`] that contains it. The timestamp is the
     /// numbre of seconds since the beginning of UNIX time.
+    #[cfg(test)]
     pub fn query_for_current_batch_window(&self, now: u64) -> crate::messages::Query {
         let start = now - (now % self.time_precision);
         crate::messages::Query::TimeInterval {
@@ -417,6 +417,10 @@ impl DapTaskConfig {
             },
         }
     }
+
+    pub(crate) fn truncate_time(&self, time: Time) -> Time {
+        time - (time % self.time_precision)
+    }
 }
 
 impl AsRef<DapTaskConfig> for DapTaskConfig {
diff --git a/daphne/src/messages.rs b/daphne/src/messages.rs
index b0b59d9..531ecdf 100644
--- a/daphne/src/messages.rs
+++ b/daphne/src/messages.rs
@@ -430,7 +430,7 @@ impl Decode for AggregateResp {
 }
 
 /// A batch interval.
-#[derive(Clone, Debug, Default, Deserialize, PartialEq, Eq, Serialize)]
+#[derive(Clone, Debug, Default, Deserialize, Eq, Hash, PartialEq, Serialize)]
 #[allow(missing_docs)]
 pub struct Interval {
     pub start: Time,
@@ -461,7 +461,7 @@ impl Decode for Interval {
 }
 
 /// A query issued by the Collector in a collect request.
-#[derive(Clone, Debug, Deserialize, PartialEq, Eq, Serialize)]
+#[derive(Clone, Debug, Deserialize, Hash, PartialEq, Eq, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum Query {
     TimeInterval { batch_interval: Interval },


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 92c3fc4..525c5df 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -15,7 +15,7 @@ use crate::{
         TransitionVar,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
-    testing::{AggStoreState, BucketInfo, MockAggregateInfo, MockAggregator, ReportStore},
+    testing::{AggStore, MockAggregateInfo, MockAggregator, ReportStore},
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateShare, DapCollectJob, DapGlobalConfig, DapLeaderTransition,
     DapMeasurement, DapQueryConfig, DapRequest, DapTaskConfig, DapVersion, Prio3Config, VdafConfig,
@@ -781,16 +781,28 @@ async fn http_post_aggregate_failure_batch_collected() {
     }];
     let req = t.gen_test_agg_init_req(task_id, report_shares).await;
 
-    // This is to ensure that the lock is released before we call http_post_aggregate.
+    // 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 agg_store_mutex_guard = t.helper.agg_store.lock().expect("lock() failed");
-        let agg_store = agg_store_mutex_guard.deref_mut();
-        let bucket_info = BucketInfo::new(task_config, task_id, report.metadata.time);
-        let agg_store_state = AggStoreState {
-            agg_share: DapAggregateShare::default(),
-            collected: true,
-        };
-        agg_store.insert(bucket_info, agg_store_state);
+        let mut guard = t
+            .helper
+            .agg_store
+            .lock()
+            .expect("agg_store: failed to lock");
+        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,
+                },
+            },
+            AggStore {
+                agg_share: DapAggregateShare::default(),
+                collected: true,
+            },
+        );
     }
 
     // Get AggregateResp and then extract the transition data from inside.
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 16f7819..0dd7201 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -7,12 +7,12 @@ use crate::{
     auth::{BearerToken, BearerTokenProvider},
     hpke::{HpkeDecrypter, HpkeReceiverConfig},
     messages::{
-        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, HpkeConfig, Id, Nonce, Report,
-        ReportShare, Time, TransitionFailure,
+        BatchSelector, CollectReq, CollectResp, HpkeCiphertext, HpkeConfig, Id, Interval, Nonce,
+        Report, ReportMetadata, ReportShare, Time, TransitionFailure,
     },
     roles::{DapAggregator, DapAuthorizedSender, DapHelper, DapLeader},
     DapAbort, DapAggregateShare, DapCollectJob, DapError, DapGlobalConfig, DapHelperState,
-    DapOutputShare, DapRequest, DapResponse, DapTaskConfig,
+    DapOutputShare, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig,
 };
 use async_trait::async_trait;
 use rand::{thread_rng, Rng};
@@ -20,7 +20,7 @@ use serde::{Deserialize, Serialize};
 use std::{
     collections::{HashMap, HashSet, VecDeque},
     hash::Hash,
-    ops::{Deref, DerefMut},
+    ops::DerefMut,
     sync::{Arc, Mutex},
     time::SystemTime,
 };
@@ -42,7 +42,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<BucketInfo, AggStoreState>>>,
+    pub(crate) agg_store: Arc<Mutex<HashMap<Id, HashMap<BatchSelector, AggStore>>>>,
 }
 
 #[allow(dead_code)]
@@ -50,26 +50,47 @@ impl MockAggregator {
     /// Conducts checks on a received report to see whether:
     /// 1) the report falls into a batch that has been already collected, or
     /// 2) the report has been submitted by the client in the past.
-    fn check_report(
+    async fn check_report_early_fail(
         &self,
-        bucket_info: &BucketInfo,
-        nonce: &Nonce,
-        report_store: &ReportStore,
-        agg_store: &HashMap<BucketInfo, AggStoreState>,
-    ) -> Result<(), TransitionFailure> {
+        task_id: &Id,
+        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.
-        if matches!(agg_store.get(bucket_info), Some(agg_store_state) if agg_store_state.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)
         {
-            return Err(TransitionFailure::BatchCollected);
+            return Some(TransitionFailure::BatchCollected);
         }
 
         // Check whether the same report has been submitted in the past.
-        if report_store.processed.contains(nonce) {
-            return Err(TransitionFailure::ReportReplayed);
+        let mut guard = self
+            .report_store
+            .lock()
+            .expect("report_store: failed to lock");
+        if guard
+            .entry(task_id.clone())
+            .or_default()
+            .processed
+            .contains(&metadata.nonce)
+        {
+            return Some(TransitionFailure::ReportReplayed);
         }
 
-        Ok(())
+        None
     }
 
     fn get_hpke_receiver_config_for(&self, hpke_config_id: u8) -> Option<&HpkeReceiverConfig> {
@@ -190,21 +211,24 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
         task_id: &Id,
         batch_selector: &BatchSelector,
     ) -> Result<bool, DapError> {
-        let mut agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref_mut();
+        let guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        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_bucket_info, agg_store_state) in agg_store.iter() {
-            if task_id == &inner_bucket_info.task_id
-                && batch_interval.start <= inner_bucket_info.window
-                && batch_interval.end() > inner_bucket_info.window
-                && agg_store_state.collected
+        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);
             }
         }
+
         Ok(false)
     }
 
@@ -217,33 +241,25 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
             .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 agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref_mut();
-        let mut bucket_info = BucketInfo {
-            task_id: task_id.clone(),
-            window: 0,
-        };
+        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() {
-            bucket_info.window = window;
-
-            if let Some(agg_store_state) = agg_store.get_mut(&bucket_info) {
-                agg_store_state.agg_share.merge(agg_share_delta)?;
-            } else {
-                agg_store.insert(
-                    bucket_info.clone(),
-                    AggStoreState {
-                        agg_share: agg_share_delta,
-                        collected: false,
-                    },
-                );
-            }
+            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();
+            inner_agg_store.agg_share.merge(agg_share_delta)?;
         }
 
         Ok(())
@@ -254,25 +270,21 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
         task_id: &Id,
         batch_selector: &BatchSelector,
     ) -> Result<DapAggregateShare, DapError> {
-        // Lock agg_store.
-        let mut agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref_mut();
+        let guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        let agg_store = guard.get(task_id).expect("agg_store: unrecognized task");
 
         // Fetch aggregate shares.
         let mut agg_share = DapAggregateShare::default();
         let batch_interval = batch_selector.unwrap_interval();
-        for (inner_bucket_info, agg_store_state) in agg_store.iter() {
-            if task_id == &inner_bucket_info.task_id
-                && batch_interval.start <= inner_bucket_info.window
-                && batch_interval.end() > inner_bucket_info.window
+        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
             {
-                if agg_store_state.collected {
+                if inner_agg_store.collected {
                     return Err(DapError::Abort(DapAbort::BatchOverlap));
                 } else {
-                    agg_share.merge(agg_store_state.agg_share.clone())?;
+                    agg_share.merge(inner_agg_store.agg_share.clone())?;
                 }
             }
         }
@@ -285,19 +297,16 @@ impl<'a> DapAggregator<'a, BearerToken> for MockAggregator {
         task_id: &Id,
         batch_selector: &BatchSelector,
     ) -> Result<(), DapError> {
-        // Mark aggregate shares as collected.
-        let mut agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref_mut();
+        let mut guard = self.agg_store.lock().expect("agg_store: failed to lock");
+        let agg_store = guard.entry(task_id.clone()).or_default();
+
         let batch_interval = batch_selector.unwrap_interval();
-        for (inner_bucket_info, agg_store_state) in agg_store.iter_mut() {
-            if task_id == &inner_bucket_info.task_id
-                && batch_interval.start <= inner_bucket_info.window
-                && batch_interval.end() > inner_bucket_info.window
+        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
             {
-                agg_store_state.collected = true;
+                inner_agg_store.collected = true;
             }
         }
 
@@ -312,43 +321,22 @@ impl<'a> DapHelper<'a, BearerToken> for MockAggregator {
         task_id: &Id,
         report_shares: &[ReportShare],
     ) -> Result<HashMap<Nonce, TransitionFailure>, DapError> {
-        let task_config = self
-            .get_task_config_for(task_id)
-            .await?
-            .ok_or_else(|| DapError::fatal("task not found"))?;
-        let mut early_fails: HashMap<Nonce, TransitionFailure> = HashMap::new();
-
-        // Lock AggStateStore.
-        let agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref();
-
-        // Lock ReportStore.
-        let mut report_store_mutex_guard = self
-            .report_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let report_store = report_store_mutex_guard.deref_mut();
-        let report_store = report_store
-            .entry(task_id.clone())
-            .or_insert_with(ReportStore::new);
-
+        let mut early_fails = HashMap::new();
         for report_share in report_shares.iter() {
-            let bucket_info = BucketInfo::new(task_config, task_id, report_share.metadata.time);
-
             // Check whether Report has been collected or replayed.
-            if let Err(transition_failure) = self.check_report(
-                &bucket_info,
-                &report_share.metadata.nonce,
-                report_store,
-                agg_store,
-            ) {
+            if let Some(transition_failure) = self
+                .check_report_early_fail(task_id, &report_share.metadata)
+                .await
+            {
                 early_fails.insert(report_share.metadata.nonce.clone(), transition_failure);
             };
 
             // Mark Report processed.
+            let mut guard = self
+                .report_store
+                .lock()
+                .expect("report_store: failed to lock");
+            let report_store = guard.entry(task_id.clone()).or_default();
             report_store
                 .processed
                 .insert(report_share.metadata.nonce.clone());
@@ -427,36 +415,24 @@ impl<'a> DapLeader<'a, BearerToken> for MockAggregator {
             .get_task_config_for(task_id)
             .await?
             .ok_or_else(|| DapError::fatal("task not found"))?;
-        let bucket_info = BucketInfo::new(task_config, task_id, report.metadata.time);
-
-        // Lock AggStateStore.
-        let agg_store_mutex_guard = self
-            .agg_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let agg_store = agg_store_mutex_guard.deref();
-
-        // Lock ReportStore.
-        let mut report_store_mutex_guard = self
-            .report_store
-            .lock()
-            .map_err(|e| DapError::Fatal(e.to_string()))?;
-        let report_store = report_store_mutex_guard.deref_mut();
-        let report_store = report_store
-            .entry(task_id.clone())
-            .or_insert_with(ReportStore::new);
+        if matches!(task_config.query, DapQueryConfig::FixedSize { .. }) {
+            panic!("TODO(issue #100)");
+        }
 
         // Check whether Report has been collected or replayed.
-        if let Err(transition_failure) = self.check_report(
-            &bucket_info,
-            &report.metadata.nonce,
-            report_store,
-            agg_store,
-        ) {
+        if let Some(transition_failure) = self
+            .check_report_early_fail(task_id, &report.metadata)
+            .await
+        {
             return Err(DapError::Transition(transition_failure));
         };
 
         // Store Report for future processing.
+        let mut guard = self
+            .report_store
+            .lock()
+            .expect("report_store: failed to lock");
+        let report_store = guard.entry(task_id.clone()).or_default();
         report_store.pending.push_back(report.clone());
         Ok(())
     }
@@ -640,39 +616,13 @@ pub(crate) struct HelperStateInfo {
     agg_job_id: Id,
 }
 
-/// Information associated to a certain report for a given task and nonce to decide which bucket it would be put into.
-#[derive(Clone, Eq, Hash, PartialEq, Deserialize, Serialize)]
-pub(crate) struct BucketInfo {
-    task_id: Id,
-    window: Time,
-}
-
-impl BucketInfo {
-    pub(crate) fn new(task_config: &DapTaskConfig, task_id: &Id, time: Time) -> Self {
-        let window = time - (time % task_config.time_precision);
-
-        Self {
-            task_id: task_id.clone(),
-            window,
-        }
-    }
-}
-
 /// Stores the reports received from Clients.
+#[derive(Default)]
 pub(crate) struct ReportStore {
     pub(crate) pending: VecDeque<Report>,
     pub(crate) processed: HashSet<Nonce>,
 }
 
-impl ReportStore {
-    pub(crate) fn new() -> Self {
-        let pending: VecDeque<Report> = VecDeque::new();
-        let processed: HashSet<Nonce> = HashSet::new();
-
-        Self { pending, processed }
-    }
-}
-
 /// Stores the state of the collect job.
 pub(crate) enum CollectJobState {
     Pending(CollectReq),
@@ -696,10 +646,11 @@ impl LeaderState {
     }
 }
 
-/// AggStoreState keeps track of the following:
+/// AggStore keeps track of the following:
 /// * Aggregate share
 /// * Whether this aggregate share has been collected
-pub(crate) struct AggStoreState {
+#[derive(Default)]
+pub(crate) struct AggStore {
     pub(crate) agg_share: DapAggregateShare,
     pub(crate) collected: bool,
 }