[0.6] Scrub reports during aggregation job creation, rather than the initial aggregation step. (#2559)

Other than being the earliest point that we can scrub reports, this
avoids (briefly) duplicating report data in both the `client_reports`
and `report_aggregations` tables. I think placing this logic in the
aggregation job creator also makes more sense for when we'll eventually
expand it to handle VDAFs requiring multiple aggregations per report.

Author: branlwyd

aggregator/src/aggregator/aggregation_job_creator.rs

@@ -4,6 +4,7 @@ use fixed::{
     types::extra::{U15, U31},
     FixedI16, FixedI32,
 };
+use futures::future::try_join_all;
 use janus_aggregator_core::{
     datastore::models::{AggregationJob, AggregationJobState},
     datastore::{self, Datastore},
@@ -33,8 +34,15 @@ use prio::{
     },
 };
 use rand::{random, thread_rng, Rng};
-use std::{collections::HashMap, sync::Arc, time::Duration};
-use tokio::time::{self, sleep_until, Instant, MissedTickBehavior};
+use std::{
+    collections::{HashMap, HashSet},
+    sync::Arc,
+    time::Duration,
+};
+use tokio::{
+    time::{self, sleep_until, Instant, MissedTickBehavior},
+    try_join,
+};
 use tracing::{debug, error, info};
 use trillium_tokio::{CloneCounterObserver, Stopper};
 
@@ -528,6 +536,7 @@ impl<C: Clock + 'static> AggregationJobCreator<C> {
 
                     // Generate aggregation jobs & report aggregations based on the reports we read.
                     let mut aggregation_job_writer = AggregationJobWriter::new(Arc::clone(&task));
+                    let mut report_ids_to_scrub = HashSet::new();
                     for agg_job_reports in reports.chunks(this.max_aggregation_job_size) {
                         if agg_job_reports.len() < this.min_aggregation_job_size {
                             if !agg_job_reports.is_empty() {
@@ -585,12 +594,22 @@ impl<C: Clock + 'static> AggregationJobCreator<C> {
                                 ))
                             })
                             .collect::<Result<_, datastore::Error>>()?;
+                        report_ids_to_scrub
+                            .extend(agg_job_reports.iter().map(|report| *report.metadata().id()));
 
                         aggregation_job_writer.put(aggregation_job, report_aggregations)?;
                     }
 
                     // Write the aggregation jobs & report aggregations we created.
-                    aggregation_job_writer.write(tx, vdaf).await?;
+                    try_join!(
+                        aggregation_job_writer.write(tx, vdaf),
+                        try_join_all(
+                            report_ids_to_scrub
+                                .iter()
+                                .map(|report_id| tx.scrub_client_report(task.id(), report_id))
+                        )
+                    )?;
+
                     Ok(!aggregation_job_writer.is_empty())
                 })
             })
@@ -2585,7 +2604,7 @@ mod tests {
         for<'a> A::PrepareState: ParameterizedDecode<(&'a A, usize)>,
         A::PublicShare: PartialEq,
     {
-        try_join!(
+        let (agg_jobs_and_report_ids, batches) = try_join!(
             try_join_all(
                 tx.get_aggregation_jobs_for_task(task_id)
                     .await
@@ -2602,7 +2621,7 @@ mod tests {
                         .await
                         .map(|report_aggs| {
                             // Verify that each report aggregation has the expected state.
-                            let report_ids = report_aggs
+                            let report_ids: Vec<_> = report_aggs
                                 .into_iter()
                                 .map(|ra| {
                                     let want_ra_state =
@@ -2624,6 +2643,28 @@ mod tests {
             ),
             tx.get_batches_for_task(task_id),
         )
-        .unwrap()
+        .unwrap();
+
+        // Verify that all reports we saw a report aggregation for are scrubbed.
+        let all_seen_report_ids: HashSet<_> = agg_jobs_and_report_ids
+            .iter()
+            .flat_map(|(_, report_ids)| report_ids)
+            .collect();
+        for report_id in &all_seen_report_ids {
+            tx.verify_client_report_scrubbed(task_id, report_id).await;
+        }
+
+        // Verify that all reports we did not see a report aggregation for are not scrubbed. (We do
+        // so by reading the report, since reading a report will fail if the report is scrubbed.)
+        for report_id in want_ra_states.keys() {
+            if all_seen_report_ids.contains(report_id) {
+                continue;
+            }
+            tx.get_client_report(vdaf, task_id, report_id)
+                .await
+                .unwrap();
+        }
+
+        (agg_jobs_and_report_ids, batches)
     }
 }

aggregator/src/aggregator/aggregation_job_driver.rs

@@ -23,7 +23,7 @@ use janus_messages::{
     query_type::{FixedSize, TimeInterval},
     AggregationJobContinueReq, AggregationJobInitializeReq, AggregationJobResp,
     PartialBatchSelector, PrepareContinue, PrepareError, PrepareInit, PrepareResp,
-    PrepareStepResult, ReportId, ReportShare, Role,
+    PrepareStepResult, ReportShare, Role,
 };
 use opentelemetry::{
     metrics::{Counter, Histogram, Meter, Unit},
@@ -310,16 +310,6 @@ impl AggregationJobDriver {
         A::PrepareMessage: PartialEq + Eq + Send + Sync,
         A::PublicShare: PartialEq + Send + Sync,
     {
-        // We currently scrub all reports included in an aggregation job as part of completing the
-        // first step. Once we support VDAFs which accept use an aggregation parameter &
-        // permit/require multiple aggregations per report, we will need a more complicated strategy
-        // where we only scrub a report as part of the _final_ aggregation over the report.
-        let report_ids_to_scrub = report_aggregations
-            .iter()
-            .map(|ra| ra.report_id())
-            .copied()
-            .collect();
-
         // Only process non-failed report aggregations.
         let report_aggregations: Vec<_> = report_aggregations
             .into_iter()
@@ -479,7 +469,6 @@ impl AggregationJobDriver {
             aggregation_job,
             &stepped_aggregations,
             report_aggregations_to_write,
-            report_ids_to_scrub,
             resp.prepare_resps(),
         )
         .await
@@ -577,7 +566,6 @@ impl AggregationJobDriver {
             aggregation_job,
             &stepped_aggregations,
             report_aggregations_to_write,
-            Vec::new(), /* reports are only scrubbed on the initial step */
             resp.prepare_resps(),
         )
         .await
@@ -598,7 +586,6 @@ impl AggregationJobDriver {
         aggregation_job: AggregationJob<SEED_SIZE, Q, A>,
         stepped_aggregations: &[SteppedAggregation<SEED_SIZE, A>],
         mut report_aggregations_to_write: Vec<ReportAggregation<SEED_SIZE, A>>,
-        report_ids_to_scrub: Vec<ReportId>,
         helper_prep_resps: &[PrepareResp],
     ) -> Result<(), Error>
     where
@@ -759,27 +746,19 @@ impl AggregationJobDriver {
             report_aggregations_to_write,
         )?;
         let aggregation_job_writer = Arc::new(aggregation_job_writer);
-
-        let report_ids_to_scrub = Arc::new(report_ids_to_scrub);
         let accumulator = Arc::new(accumulator);
+
         datastore
             .run_tx("step_aggregation_job_2", |tx| {
-                let task_id = *task.id();
                 let vdaf = Arc::clone(&vdaf);
                 let aggregation_job_writer = Arc::clone(&aggregation_job_writer);
                 let accumulator = Arc::clone(&accumulator);
-                let report_ids_to_scrub = Arc::clone(&report_ids_to_scrub);
                 let lease = Arc::clone(&lease);
 
                 Box::pin(async move {
-                    let (unwritable_ra_report_ids, unwritable_ba_report_ids, _, _) = try_join!(
+                    let (unwritable_ra_report_ids, unwritable_ba_report_ids, _) = try_join!(
                         aggregation_job_writer.write(tx, Arc::clone(&vdaf)),
                         accumulator.flush_to_datastore(tx, &vdaf),
-                        try_join_all(
-                            report_ids_to_scrub
-                                .iter()
-                                .map(|report_id| tx.scrub_client_report(&task_id, report_id))
-                        ),
                         tx.release_aggregation_job(&lease),
                     )?;
 

aggregator/src/aggregator/batch_creator.rs

@@ -7,13 +7,13 @@ use janus_aggregator_core::datastore::{
 };
 use janus_core::time::{Clock, DurationExt, TimeExt};
 use janus_messages::{
-    query_type::FixedSize, AggregationJobStep, BatchId, Duration, Interval, TaskId, Time,
+    query_type::FixedSize, AggregationJobStep, BatchId, Duration, Interval, ReportId, TaskId, Time,
 };
 use prio::{codec::Encode, vdaf::Aggregator};
 use rand::random;
 use std::{
     cmp::{max, min, Ordering},
-    collections::{binary_heap::PeekMut, hash_map, BinaryHeap, HashMap, VecDeque},
+    collections::{binary_heap::PeekMut, hash_map, BinaryHeap, HashMap, HashSet, VecDeque},
     ops::RangeInclusive,
     sync::Arc,
 };
@@ -32,8 +32,9 @@ where
 {
     properties: Properties,
     aggregation_job_writer: &'a mut AggregationJobWriter<SEED_SIZE, FixedSize, A>,
-    map: HashMap<Option<Time>, Bucket<SEED_SIZE, A>>,
+    buckets: HashMap<Option<Time>, Bucket<SEED_SIZE, A>>,
     new_batches: Vec<(BatchId, Option<Time>)>,
+    report_ids_to_scrub: HashSet<ReportId>,
 }
 
 /// Common properties used by [`BatchCreator`]. This is broken out into a separate structure to make
@@ -72,8 +73,9 @@ where
                 task_batch_time_window_size,
             },
             aggregation_job_writer,
-            map: HashMap::new(),
+            buckets: HashMap::new(),
             new_batches: Vec::new(),
+            report_ids_to_scrub: HashSet::new(),
         }
     }
 
@@ -95,15 +97,15 @@ where
                     .to_batch_interval_start(&batch_time_window_size)
             })
             .transpose()?;
-        let mut map_entry = self.map.entry(time_bucket_start_opt);
+        let mut map_entry = self.buckets.entry(time_bucket_start_opt);
         let bucket = match &mut map_entry {
             hash_map::Entry::Occupied(occupied) => occupied.get_mut(),
             hash_map::Entry::Vacant(_) => {
                 // Lazily find existing unfilled batches.
                 let outstanding_batches = tx
                     .get_outstanding_batches(&self.properties.task_id, &time_bucket_start_opt)
                     .await?;
-                self.map
+                self.buckets
                     .entry(time_bucket_start_opt)
                     .or_insert_with(|| Bucket::new(outstanding_batches))
             }
@@ -115,6 +117,7 @@ where
         Self::process_batches(
             &self.properties,
             self.aggregation_job_writer,
+            &mut self.report_ids_to_scrub,
             &mut self.new_batches,
             &time_bucket_start_opt,
             bucket,
@@ -136,6 +139,7 @@ where
     fn process_batches(
         properties: &Properties,
         aggregation_job_writer: &mut AggregationJobWriter<SEED_SIZE, FixedSize, A>,
+        report_ids_to_scrub: &mut HashSet<ReportId>,
         new_batches: &mut Vec<(BatchId, Option<Time>)>,
         time_bucket_start: &Option<Time>,
         bucket: &mut Bucket<SEED_SIZE, A>,
@@ -182,6 +186,7 @@ where
                             desired_aggregation_job_size,
                             &mut bucket.unaggregated_reports,
                             aggregation_job_writer,
+                            report_ids_to_scrub,
                         )?;
                         largest_outstanding_batch.add_reports(desired_aggregation_job_size);
                     } else {
@@ -209,6 +214,7 @@ where
                             desired_aggregation_job_size,
                             &mut bucket.unaggregated_reports,
                             aggregation_job_writer,
+                            report_ids_to_scrub,
                         )?;
                         largest_outstanding_batch.add_reports(desired_aggregation_job_size);
                     } else {
@@ -249,6 +255,7 @@ where
                     desired_aggregation_job_size,
                     &mut bucket.unaggregated_reports,
                     aggregation_job_writer,
+                    report_ids_to_scrub,
                 )?;
 
                 // Loop to the top of this method to create more aggregation jobs in this newly
@@ -268,6 +275,7 @@ where
         aggregation_job_size: usize,
         unaggregated_reports: &mut VecDeque<LeaderStoredReport<SEED_SIZE, A>>,
         aggregation_job_writer: &mut AggregationJobWriter<SEED_SIZE, FixedSize, A>,
+        report_ids_to_scrub: &mut HashSet<ReportId>,
     ) -> Result<(), Error> {
         let aggregation_job_id = random();
         debug!(
@@ -280,7 +288,7 @@ where
         let mut min_client_timestamp = None;
         let mut max_client_timestamp = None;
 
-        let report_aggregations = (0u64..)
+        let report_aggregations: Vec<_> = (0u64..)
             .zip(unaggregated_reports.drain(..aggregation_job_size))
             .map(|(ord, report)| {
                 let client_timestamp = *report.metadata().time();
@@ -294,6 +302,7 @@ where
                 report.as_start_leader_report_aggregation(aggregation_job_id, ord)
             })
             .collect();
+        report_ids_to_scrub.extend(report_aggregations.iter().map(|ra| *ra.report_id()));
 
         let min_client_timestamp = min_client_timestamp.unwrap(); // unwrap safety: aggregation_job_size > 0
         let max_client_timestamp = max_client_timestamp.unwrap(); // unwrap safety: aggregation_job_size > 0
@@ -329,10 +338,11 @@ where
         // be smaller than max_aggregation_job_size. We will only create jobs smaller than
         // min_aggregation_job_size if the remaining headroom in a batch requires it, otherwise
         // remaining reports will be added to unaggregated_report_ids, to be marked as unaggregated.
-        for (time_bucket_start, mut bucket) in self.map.into_iter() {
+        for (time_bucket_start, mut bucket) in self.buckets.into_iter() {
             Self::process_batches(
                 &self.properties,
                 self.aggregation_job_writer,
+                &mut self.report_ids_to_scrub,
                 &mut self.new_batches,
                 &time_bucket_start,
                 &mut bucket,
@@ -348,6 +358,11 @@ where
 
         try_join!(
             self.aggregation_job_writer.write(tx, vdaf),
+            try_join_all(
+                self.report_ids_to_scrub
+                    .iter()
+                    .map(|report_id| tx.scrub_client_report(&self.properties.task_id, report_id))
+            ),
             try_join_all(
                 self.new_batches
                     .iter()

aggregator_core/src/datastore.rs

@@ -1555,6 +1555,34 @@ impl<C: Clock> Transaction<'_, C> {
         )
     }
 
+    #[cfg(feature = "test-util")]
+    #[cfg_attr(docsrs, doc(cfg(feature = "test-util")))]
+    pub async fn verify_client_report_scrubbed(&self, task_id: &TaskId, report_id: &ReportId) {
+        let row = self
+            .query_one(
+                "SELECT
+                    client_reports.extensions,
+                    client_reports.public_share,
+                    client_reports.leader_input_share,
+                    client_reports.helper_encrypted_input_share
+                FROM client_reports
+                JOIN tasks ON tasks.id = client_reports.task_id
+                WHERE tasks.task_id = $1
+                AND client_reports.report_id = $2",
+                &[task_id.as_ref(), report_id.as_ref()],
+            )
+            .await
+            .unwrap();
+
+        assert_eq!(row.get::<_, Option<Vec<u8>>>("extensions"), None);
+        assert_eq!(row.get::<_, Option<Vec<u8>>>("public_share"), None);
+        assert_eq!(row.get::<_, Option<Vec<u8>>>("leader_input_share"), None);
+        assert_eq!(
+            row.get::<_, Option<Vec<u8>>>("helper_encrypted_input_share"),
+            None
+        );
+    }
+
     /// put_report_share stores a report share, given its associated task ID.
     ///
     /// This method is intended for use by aggregators acting in the Helper role; notably, it does

aggregator_core/src/datastore/tests.rs

@@ -660,29 +660,7 @@ async fn roundtrip_report(ephemeral_datastore: EphemeralDatastore) {
         Box::pin(async move {
             tx.scrub_client_report(&task_id, &report_id).await.unwrap();
 
-            let row = tx
-                .query_one(
-                    "SELECT
-                        client_reports.extensions,
-                        client_reports.public_share,
-                        client_reports.leader_input_share,
-                        client_reports.helper_encrypted_input_share
-                    FROM client_reports
-                    JOIN tasks ON tasks.id = client_reports.task_id
-                    WHERE tasks.task_id = $1
-                    AND client_reports.report_id = $2",
-                    &[&task_id.as_ref(), &report_id.as_ref()],
-                )
-                .await
-                .unwrap();
-
-            assert_eq!(row.get::<_, Option<Vec<u8>>>("extensions"), None);
-            assert_eq!(row.get::<_, Option<Vec<u8>>>("public_share"), None);
-            assert_eq!(row.get::<_, Option<Vec<u8>>>("leader_input_share"), None);
-            assert_eq!(
-                row.get::<_, Option<Vec<u8>>>("helper_encrypted_input_share"),
-                None
-            );
+            tx.verify_client_report_scrubbed(&task_id, &report_id).await;
 
             assert_matches!(
                 tx.get_client_report::<0, dummy_vdaf::Vdaf>(