Repository: cloudflare/daphne
Pull request #58 (2022-07-22 19:58:02)

Move `task_id` out of `DapLeader::get_reports()`
Partially addresses #25.

Moves `task_id` out of the function parameters. Existing implementors of
`DapLeader` can put this field intp the report selector. This allows for
deployments in which work is scheduled not per task but for the entire
deployment. In particular, this unblocks progress on scaling up
Daphne-Worker.

cjpatton: Reviewer note: There should be no functional changes, *except* in Daphne-Worker:
```
/internal/process/task/:task_id
```
has been changed to
```
/internal/process
```
The task ID is now carried in the body of the request.

Patch:
diff --git a/daphne/src/roles.rs b/daphne/src/roles.rs
index 5d3faad..b65dc17 100644
--- a/daphne/src/roles.rs
+++ b/daphne/src/roles.rs
@@ -147,13 +147,12 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     /// Store a report for use later on.
     async fn put_report(&self, report: &Report) -> Result<(), DapError>;
 
-    /// Fetch a sequence of reports to aggregate. The reports returned are removed from persistent
-    /// storage.
+    /// Fetch a sequence of reports to aggregate, grouped by task ID. The reports returned are
+    /// removed from persistent storage.
     async fn get_reports(
         &self,
-        task_id: &Id,
         selector: &Self::ReportSelector,
-    ) -> Result<Vec<Report>, DapError>;
+    ) -> Result<HashMap<Id, Vec<Report>>, DapError>;
 
     /// Create a collect job.
     //
@@ -239,7 +238,142 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
         Ok(self.init_collect_job(&collect_req).await?)
     }
 
-    /// Run an aggregation job if there are pending reports, then process all pending collect jobs.
+    /// Run the aggregation sub-protocol for the given set of reports. Return the number of reports
+    /// that were aggregated successfully.
+    async fn run_agg_job(
+        &self,
+        task_id: &Id,
+        task_config: &DapTaskConfig,
+        reports: Vec<Report>,
+    ) -> Result<u64, DapAbort> {
+        let mut rng = thread_rng();
+
+        // Prepare AggregateInitializeReq.
+        let agg_job_id = Id(rng.gen());
+        let transition = task_config.vdaf.produce_agg_init_req(
+            self,
+            &task_config.vdaf_verify_key,
+            task_id,
+            &agg_job_id,
+            reports,
+        )?;
+        let (state, agg_init_req) = match transition {
+            DapLeaderTransition::Continue(state, agg_init_req) => (state, agg_init_req),
+            DapLeaderTransition::Skip => return Ok(0),
+            DapLeaderTransition::Uncommitted(..) => {
+                return Err(DapError::fatal("unexpected state transition (uncommitted)").into())
+            }
+        };
+
+        // Send AggregateInitializeReq and receive AggregateResp.
+        let resp = leader_post!(
+            self,
+            task_id,
+            task_config,
+            "/aggregate",
+            MEDIA_TYPE_AGG_INIT_REQ,
+            agg_init_req.get_encoded()
+        );
+        let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
+
+        // Prepare AggreagteContinueReq.
+        let transition = task_config
+            .vdaf
+            .handle_agg_resp(task_id, &agg_job_id, state, agg_resp)?;
+        let (uncommited, agg_cont_req) = match transition {
+            DapLeaderTransition::Uncommitted(uncommited, agg_cont_req) => {
+                (uncommited, agg_cont_req)
+            }
+            DapLeaderTransition::Skip => return Ok(0),
+            DapLeaderTransition::Continue(..) => {
+                return Err(DapError::fatal("unexpected state transition (continue)").into())
+            }
+        };
+
+        // Send AggregateContinueReq and receive AggregateResp.
+        let resp = leader_post!(
+            self,
+            task_id,
+            task_config,
+            "/aggregate",
+            MEDIA_TYPE_AGG_CONT_REQ,
+            agg_cont_req.get_encoded()
+        );
+        let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
+
+        // Commit the output shares.
+        let out_shares = task_config
+            .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?;
+        Ok(out_shares_count)
+    }
+
+    /// Handle a pending collect request. If the results are reasdy, then cokmpute the aggregate
+    /// results and store them to be retrieved by the Collector later. Returns the number of
+    /// reports in the batch.
+    async fn run_collect_job(
+        &self,
+        collect_id: &Id,
+        task_id: &Id,
+        task_config: &DapTaskConfig,
+        collect_req: CollectReq,
+    ) -> Result<u64, DapAbort> {
+        let leader_agg_share = self
+            .get_agg_share(&collect_req.task_id, &collect_req.batch_interval)
+            .await?;
+
+        // Check that the minimum batch size is met.
+        if leader_agg_share.report_count < task_config.min_batch_size {
+            return Ok(0);
+        }
+
+        // Prepare the Leader's aggregate share.
+        let leader_enc_agg_share = task_config.vdaf.produce_leader_encrypted_agg_share(
+            &task_config.collector_hpke_config,
+            &collect_req.task_id,
+            &collect_req.batch_interval,
+            &leader_agg_share,
+        )?;
+
+        // Prepare AggregateShareReq.
+        let agg_share_req = AggregateShareReq {
+            task_id: collect_req.task_id.clone(),
+            batch_interval: collect_req.batch_interval.clone(),
+            agg_param: collect_req.agg_param.clone(),
+            report_count: leader_agg_share.report_count,
+            checksum: leader_agg_share.checksum,
+        };
+
+        // Send AggregateShareReq and receive AggregateShareResp.
+        let resp = leader_post!(
+            self,
+            task_id,
+            task_config,
+            "/aggregate_share",
+            MEDIA_TYPE_AGG_SHARE_REQ,
+            agg_share_req.get_encoded()
+        );
+        let agg_share_resp = AggregateShareResp::get_decoded(&resp.payload)?;
+
+        // Complete the collect job.
+        let collect_resp = CollectResp {
+            encrypted_agg_shares: vec![leader_enc_agg_share, agg_share_resp.encrypted_agg_share],
+        };
+        self.finish_collect_job(task_id, collect_id, &collect_resp)
+            .await?;
+
+        // Mark reports as collected.
+        self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_interval)
+            .await?;
+
+        Ok(agg_share_req.report_count)
+    }
+
+    /// Fetch a set of reports grouped by task. For each task, first aggregate the reports, then
+    /// process the collect job queue. It is not safe to run multiple instances of this function in
+    /// parallel.
     ///
     /// This method is geared primarily towards testing. It also demonstrates how to properly
     /// synchronize collect and aggregation jobs. If used in a large DAP deployment, it is likely
@@ -247,140 +381,33 @@ pub trait DapLeader<S>: DapAuthorizedSender<S> + DapAggregator<S> {
     /// jobs in parallel.
     async fn process(
         &self,
-        task_id: &Id,
         selector: &Self::ReportSelector,
     ) -> Result<DapLeaderProcessTelemetry, DapAbort> {
-        let mut rng = thread_rng();
         let mut telem = DapLeaderProcessTelemetry::default();
 
-        let task_config = self
-            .get_task_config_for(task_id)
-            .ok_or(DapAbort::UnrecognizedTask)?;
-
-        // Pick the set of candidate reports.
-        let reports = self.get_reports(task_id, selector).await?;
-        telem.reports_processed += reports.len() as u64;
-        if !reports.is_empty() {
-            // Prepare AggregateInitializeReq.
-            let agg_job_id = Id(rng.gen());
-            let transition = task_config.vdaf.produce_agg_init_req(
-                self,
-                &task_config.vdaf_verify_key,
-                task_id,
-                &agg_job_id,
-                reports,
-            )?;
-            let (state, agg_init_req) = match transition {
-                DapLeaderTransition::Continue(state, agg_init_req) => (state, agg_init_req),
-                DapLeaderTransition::Skip => return Ok(telem),
-                DapLeaderTransition::Uncommitted(..) => {
-                    return Err(DapError::fatal("unexpected state transition (uncommitted)").into())
-                }
-            };
-
-            // Send AggregateInitializeReq and receive AggregateResp.
-            let resp = leader_post!(
-                self,
-                task_id,
-                task_config,
-                "/aggregate",
-                MEDIA_TYPE_AGG_INIT_REQ,
-                agg_init_req.get_encoded()
-            );
-            let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
-
-            // Prepare AggreagteContinueReq.
-            let transition =
-                task_config
-                    .vdaf
-                    .handle_agg_resp(task_id, &agg_job_id, state, agg_resp)?;
-            let (uncommited, agg_cont_req) = match transition {
-                DapLeaderTransition::Uncommitted(uncommited, agg_cont_req) => {
-                    (uncommited, agg_cont_req)
-                }
-                DapLeaderTransition::Skip => return Ok(telem),
-                DapLeaderTransition::Continue(..) => {
-                    return Err(DapError::fatal("unexpected state transition (continue)").into())
-                }
-            };
-
-            // Send AggregateContinueReq and receive AggregateResp.
-            let resp = leader_post!(
-                self,
-                task_id,
-                task_config,
-                "/aggregate",
-                MEDIA_TYPE_AGG_CONT_REQ,
-                agg_cont_req.get_encoded()
-            );
-            let agg_resp = AggregateResp::get_decoded(&resp.payload)?;
-
-            // Commit the output shares.
-            let out_shares = task_config
-                .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?;
-            telem.reports_aggregated += out_shares_count;
-        }
+        // TODO Handle tasks in parallel.
+        for (task_id, reports) in self.get_reports(selector).await?.into_iter() {
+            let task_config = self
+                .get_task_config_for(&task_id)
+                .ok_or(DapAbort::UnrecognizedTask)?;
 
-        // Process pending collect jobs. We wait until all aggregation jobs are finished before
-        // proceeding to this step. This is to prevent a race condition involving an aggregate
-        // share computed during a collect job and any output shares computed during an aggregation
-        // job.
-        let pending = self.get_pending_collect_jobs(task_id).await?;
-        for (collect_id, collect_req) in pending.into_iter() {
-            let leader_agg_share = self
-                .get_agg_share(&collect_req.task_id, &collect_req.batch_interval)
-                .await?;
-
-            // Check that the minimum batch size is met.
-            if leader_agg_share.report_count < task_config.min_batch_size {
-                continue;
+            // Aggregate the reports.
+            telem.reports_processed += reports.len() as u64;
+            if !reports.is_empty() {
+                telem.reports_aggregated +=
+                    self.run_agg_job(&task_id, task_config, reports).await?;
             }
 
-            // Prepare the Leader's aggregate share.
-            let leader_enc_agg_share = task_config.vdaf.produce_leader_encrypted_agg_share(
-                &task_config.collector_hpke_config,
-                &collect_req.task_id,
-                &collect_req.batch_interval,
-                &leader_agg_share,
-            )?;
-
-            // Prepare AggregateShareReq.
-            let agg_share_req = AggregateShareReq {
-                task_id: collect_req.task_id.clone(),
-                batch_interval: collect_req.batch_interval.clone(),
-                agg_param: collect_req.agg_param.clone(),
-                report_count: leader_agg_share.report_count,
-                checksum: leader_agg_share.checksum,
-            };
-
-            // Send AggregateShareReq and receive AggregateShareResp.
-            let resp = leader_post!(
-                self,
-                task_id,
-                task_config,
-                "/aggregate_share",
-                MEDIA_TYPE_AGG_SHARE_REQ,
-                agg_share_req.get_encoded()
-            );
-            let agg_share_resp = AggregateShareResp::get_decoded(&resp.payload)?;
-
-            // Complete the collect job.
-            let collect_resp = CollectResp {
-                encrypted_agg_shares: vec![
-                    leader_enc_agg_share,
-                    agg_share_resp.encrypted_agg_share,
-                ],
-            };
-            self.finish_collect_job(task_id, &collect_id, &collect_resp)
-                .await?;
-
-            // Mark reports as collected.
-            self.mark_collected(&agg_share_req.task_id, &agg_share_req.batch_interval)
-                .await?;
-            telem.reports_collected += agg_share_req.report_count;
+            // Process pending collect jobs. We wait until all aggregation jobs are finished before
+            // proceeding to this step. This is to prevent a race condition involving an aggregate
+            // share computed during a collect job and any output shares computed during an aggregation
+            // job.
+            let pending = self.get_pending_collect_jobs(&task_id).await?;
+            for (collect_id, collect_req) in pending.into_iter() {
+                telem.reports_collected += self
+                    .run_collect_job(&collect_id, &task_id, task_config, collect_req)
+                    .await?;
+            }
         }
 
         Ok(telem)
diff --git a/daphne_worker/src/dap.rs b/daphne_worker/src/dap.rs
index 3cc513f..418c568 100644
--- a/daphne_worker/src/dap.rs
+++ b/daphne_worker/src/dap.rs
@@ -247,13 +247,12 @@ impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
 
     async fn get_reports(
         &self,
-        task_id: &Id,
-        selector: &Self::ReportSelector,
-    ) -> std::result::Result<Vec<Report>, DapError> {
+        selector: &InternalAggregateInfo,
+    ) -> std::result::Result<HashMap<Id, Vec<Report>>, DapError> {
         let durable_namespace = self.durable_object("DAP_REPORT_STORE")?;
 
         let task_config = self
-            .get_task_config_for(task_id)
+            .get_task_config_for(&selector.task_id)
             .ok_or_else(|| DapError::fatal(INT_ERR_UNRECOGNIZED_TASK))?;
 
         let batch_interval = if let Some(ref interval) = selector.batch_info {
@@ -274,7 +273,7 @@ impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
         // the Helper could preprocess its input shares in the same way. See
         // https://github.com/abetterinternet/ppm-specification/pull/174.)
         let mut reports = Vec::with_capacity(selector.agg_rate.try_into().unwrap());
-        for durable_name in self.iter_report_store_names(task_id, &batch_interval)? {
+        for durable_name in self.iter_report_store_names(&selector.task_id, &batch_interval)? {
             let num_reports_remaining = selector.agg_rate - reports.len() as u64;
             let stub = durable_namespace.id_from_name(&durable_name)?.get_stub()?;
             // TODO Don't block on DO request (issue multiple requests simultaneously).
@@ -307,7 +306,7 @@ impl<D> DapLeader<BearerToken> for DaphneWorkerConfig<D> {
             }
         }
 
-        Ok(reports)
+        Ok(HashMap::from([(selector.task_id.clone(), reports)]))
     }
 
     async fn init_collect_job(
diff --git a/daphne_worker/src/lib.rs b/daphne_worker/src/lib.rs
index 5a90c57..299ab7c 100644
--- a/daphne_worker/src/lib.rs
+++ b/daphne_worker/src/lib.rs
@@ -20,6 +20,8 @@ use worker::*;
 /// Parameters used by the Leader to select a set of reports for aggregation.
 #[derive(Debug, Deserialize, Serialize)]
 pub struct InternalAggregateInfo {
+    pub task_id: Id,
+
     #[serde(skip_serializing_if = "Option::is_none")]
     pub batch_info: Option<Interval>,
 
@@ -177,18 +179,14 @@ pub async fn run_daphne_worker(req: Request, env: Env) -> Result<Response> {
                         }
                     },
                 )
-                .post_async(
-                    "/internal/process/task/:task_id",
-                    |mut req, ctx| async move {
-                        let task_id = parse_id!(ctx.param("task_id"));
-                        let config = DaphneWorkerConfig::from_worker_context(ctx)?;
-                        let agg_info: InternalAggregateInfo = req.json().await?;
-                        match config.process(&task_id, &agg_info).await {
-                            Ok(telem) => Response::from_json(&telem),
-                            Err(e) => abort(e),
-                        }
-                    },
-                )
+                .post_async("/internal/process", |mut req, ctx| async move {
+                    let config = DaphneWorkerConfig::from_worker_context(ctx)?;
+                    let agg_info: InternalAggregateInfo = req.json().await?;
+                    match config.process(&agg_info).await {
+                        Ok(telem) => Response::from_json(&telem),
+                        Err(e) => abort(e),
+                    }
+                })
         }
 
         "helper" => router


Test patch:
diff --git a/daphne/src/roles_test.rs b/daphne/src/roles_test.rs
index 8e39abd..2bee4f3 100644
--- a/daphne/src/roles_test.rs
+++ b/daphne/src/roles_test.rs
@@ -25,6 +25,16 @@ use prio::codec::{Decode, Encode};
 use rand::{thread_rng, Rng};
 use std::{ops::DerefMut, vec};
 
+// MockAggregator's implementation of DapLeader::get_report() always returns reports for a single
+// task. This macro is used to conveniently unwrap the task ID and reports for testing purposes.
+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()
+    }};
+}
+
 impl MockAggregator {
     fn gen_test_upload_req(&self, report: Report) -> DapRequest<BearerToken> {
         let task_id = self.nominal_task_id();
@@ -437,10 +447,11 @@ async fn get_reports_fail_invalid_batch_interval() {
         duration: 0,
     };
     let selector = &MockAggregateInfo {
+        task_id: task_id.clone(),
         batch_info: Some(empty_interval),
         agg_rate: 1,
     };
-    let err = leader.get_reports(task_id, selector).await.unwrap_err();
+    let err = leader.get_reports(selector).await.unwrap_err();
 
     // Fails because a 0 second duration interval is not permitted.
     assert_matches!(err, DapError::Fatal(e) =>
@@ -466,10 +477,11 @@ async fn get_reports_empty_response() {
     // Get report.
     let now = report.nonce.time - 200000000;
     let selector = &MockAggregateInfo {
+        task_id: task_id.clone(),
         batch_info: Some(task_config.current_batch_window(now)),
         agg_rate: 1,
     };
-    let reports = leader.get_reports(task_id, selector).await.unwrap();
+    let (task_id, reports) = get_reports!(leader, selector);
 
     // We get an empty response due to no reports existing within requested batch_window.
     assert_eq!(reports.len(), 0);
@@ -477,10 +489,11 @@ async fn get_reports_empty_response() {
     // Attempt to get reports from the future.
     let now = report.nonce.time + 200000000;
     let selector = &MockAggregateInfo {
+        task_id: task_id.clone(),
         batch_info: Some(task_config.current_batch_window(now)),
         agg_rate: 1,
     };
-    let reports = leader.get_reports(task_id, selector).await.unwrap();
+    let (_task_id, reports) = get_reports!(leader, selector);
 
     // We get an empty response due to no reports existing within requested batch_window.
     assert_eq!(reports.len(), 0);
@@ -520,11 +533,11 @@ async fn e2e() {
     // Leader: Store received report to ReportStore.
     let now = 1637361337;
     let selector = &MockAggregateInfo {
+        task_id: task_id.clone(),
         batch_info: Some(task_config.current_batch_window(now)),
         agg_rate: 1,
     };
-    let res = leader.get_reports(task_id, selector).await;
-    let reports = res.unwrap();
+    let (task_id, reports) = get_reports!(leader, selector);
     assert_eq!(report, reports[0]);
 
     // Leader: Consume report share.
@@ -535,7 +548,7 @@ async fn e2e() {
         .produce_agg_init_req(
             &leader,
             &task_config.vdaf_verify_key,
-            task_id,
+            &task_id,
             &agg_job_id,
             reports,
         )
@@ -556,7 +569,7 @@ async fn e2e() {
     // Leader: Produce Leader output share and prepare aggregate continue request for Helper.
     let transition = task_config
         .vdaf
-        .handle_agg_resp(task_id, &agg_job_id, leader_state, agg_resp)
+        .handle_agg_resp(&task_id, &agg_job_id, leader_state, agg_resp)
         .unwrap();
     assert_matches!(transition, DapLeaderTransition::Uncommitted(..));
     let (leader_uncommitted, agg_cont_req) = transition.unwrap_uncommitted();
@@ -576,7 +589,7 @@ async fn e2e() {
         .vdaf
         .handle_final_agg_resp(leader_uncommitted, agg_resp)
         .unwrap();
-    leader.put_out_shares(task_id, out_shares).await.unwrap();
+    leader.put_out_shares(&task_id, out_shares).await.unwrap();
 
     // TODO(nakatsuka-y) Call http_post_collect to post a collect job.
     //                   Assume for now that the leader has already picked a collect job.
diff --git a/daphne/src/testing.rs b/daphne/src/testing.rs
index 50d6042..d43c389 100644
--- a/daphne/src/testing.rs
+++ b/daphne/src/testing.rs
@@ -65,6 +65,7 @@ pub const LEADER_BEARER_TOKEN: &str = "ivA1e7LpnySDNn1AulaZggFLQ1n7jZ8GWOUO7GY4h
 pub const COLLECTOR_BEARER_TOKEN: &str = "syfRfvcvNFF5MJk4Y-B7xjRIqD_iNzhaaEB9mYqO9hk=";
 
 pub(crate) struct MockAggregateInfo {
+    pub(crate) task_id: Id,
     pub(crate) batch_info: Option<Interval>,
     pub(crate) agg_rate: u64,
 }
@@ -409,11 +410,10 @@ impl DapLeader<BearerToken> for MockAggregator {
 
     async fn get_reports(
         &self,
-        task_id: &Id,
-        selector: &Self::ReportSelector,
-    ) -> Result<Vec<Report>, DapError> {
+        selector: &MockAggregateInfo,
+    ) -> Result<HashMap<Id, Vec<Report>>, DapError> {
         let task_config = self
-            .get_task_config_for(task_id)
+            .get_task_config_for(&selector.task_id)
             .ok_or_else(|| DapError::fatal("no task found"))?;
 
         let now = SystemTime::now()
@@ -445,7 +445,7 @@ impl DapLeader<BearerToken> for MockAggregator {
 
         // Fetch reports.
         for (inner_bucket_info, store) in report_store.iter_mut() {
-            if task_id == &inner_bucket_info.task_id
+            if selector.task_id == inner_bucket_info.task_id
                 && batch_interval.start <= inner_bucket_info.window
                 && batch_interval.end() > inner_bucket_info.window
             {
@@ -467,7 +467,7 @@ impl DapLeader<BearerToken> for MockAggregator {
             }
         }
 
-        Ok(reports)
+        Ok(HashMap::from([(selector.task_id.clone(), reports)]))
     }
 
     async fn init_collect_job(&self, _collect_req: &CollectReq) -> Result<Url, DapError> {
diff --git a/daphne_worker_test/tests/e2e.rs b/daphne_worker_test/tests/e2e.rs
index 61293ca..aec1363 100644
--- a/daphne_worker_test/tests/e2e.rs
+++ b/daphne_worker_test/tests/e2e.rs
@@ -174,6 +174,7 @@ async fn e2e_internal_leader_process() {
     let hpke_config_list = t.get_hpke_configs(&client).await;
 
     let agg_info = InternalAggregateInfo {
+        task_id: t.task_id.clone(),
         batch_info: t.batch_info(),
         agg_rate: t.min_batch_size,
     };
@@ -225,6 +226,7 @@ async fn e2e_internal_leader_process_current_batch_window() {
     let hpke_config_list = t.get_hpke_configs(&client).await;
 
     let agg_info = InternalAggregateInfo {
+        task_id: t.task_id.clone(),
         batch_info: None,
         agg_rate: 1,
     };
@@ -296,6 +298,7 @@ async fn e2e_leader_collect_ok() {
         .internal_process(
             &client,
             &InternalAggregateInfo {
+                task_id: t.task_id.clone(),
                 batch_info: batch_info.clone(),
                 agg_rate: 100,
             },
@@ -388,6 +391,7 @@ async fn e2e_leader_collect_not_ready_min_batch_size() {
         .internal_process(
             &client,
             &InternalAggregateInfo {
+                task_id: t.task_id.clone(),
                 batch_info,
                 agg_rate: 100,
             },
diff --git a/daphne_worker_test/tests/janus.rs b/daphne_worker_test/tests/janus.rs
index ab80ef2..579ce7a 100644
--- a/daphne_worker_test/tests/janus.rs
+++ b/daphne_worker_test/tests/janus.rs
@@ -60,6 +60,7 @@ async fn janus_client() {
     janus_client.upload(&23).await.unwrap();
 
     let agg_info = InternalAggregateInfo {
+        task_id: t.task_id.clone(),
         batch_info: t.batch_info(),
         agg_rate: 1,
     };
@@ -79,6 +80,7 @@ async fn janus_helper() {
     let client = t.http_client();
     let hpke_config_list = t.get_hpke_configs(&client).await;
     let agg_info = InternalAggregateInfo {
+        task_id: t.task_id.clone(),
         batch_info: t.batch_info(),
         agg_rate: t.min_batch_size,
     };
diff --git a/daphne_worker_test/tests/test_runner.rs b/daphne_worker_test/tests/test_runner.rs
index 4ebf34c..fc03cbf 100644
--- a/daphne_worker_test/tests/test_runner.rs
+++ b/daphne_worker_test/tests/test_runner.rs
@@ -330,14 +330,7 @@ impl TestRunner {
         client: &reqwest::Client,
         agg_info: &InternalAggregateInfo,
     ) -> DapLeaderProcessTelemetry {
-        let url = self
-            .leader_url
-            .join(&format!(
-                "/internal/process/task/{}",
-                self.task_id.to_base64url()
-            ))
-            .unwrap();
-
+        let url = self.leader_url.join("/internal/process").unwrap();
         let resp = client
             .post(url.as_str())
             .body(serde_json::to_string(&agg_info).unwrap())