From 21495e1e9813fa39835d6b0d2311bbedafa5b1a3 Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Tue, 13 Oct 2026 10:21:25 -0400 Subject: [PATCH 1/6] Leader-selected collection job IDs Transfers responsibility for choosing collection job IDs from the collector to the Leader, per incoming DAP changes in [1]. In order to make this work for the Leader-Selected batch mode, we also add an idempotency key to the Leader-Selected batch mode query config (also detailed in [1]). The route/handler for collection job creation: - is now POST instead of PUT - no longer includes a collection job ID - provides colleciton job ID in a Location header - sets HTTP 200 OK status and a retry-after header - checks for replayed creation requests by comparing the query instead of the collection job ID And we add an index on `(task_id, query)` on the collection jobs table to make replayed collection job lookups faster. [1]: https://github.com/ietf-wg-ppm/draft-ietf-ppm-dap/pull/797 --- aggregator/src/aggregator.rs | 194 ++++--- aggregator/src/aggregator/batch_mode.rs | 72 ++- .../src/aggregator/collection_job_tests.rs | 510 +++++++++--------- aggregator/src/aggregator/error.rs | 8 +- .../src/aggregator/garbage_collector.rs | 10 +- aggregator/src/aggregator/http_handlers.rs | 102 ++-- .../http_handlers/tests/aggregate_share.rs | 5 +- .../tests/aggregation_job_init.rs | 7 +- .../http_handlers/tests/collection_job.rs | 125 ++--- .../http_handlers/tests/helper_e2e.rs | 8 +- aggregator/src/aggregator/problem_details.rs | 13 +- aggregator/src/aggregator/taskprov_tests.rs | 17 +- aggregator_core/src/batch_mode.rs | 18 +- aggregator_core/src/datastore.rs | 99 +++- aggregator_core/src/datastore/tests.rs | 177 +++++- aggregator_core/src/task.rs | 21 +- core/src/http.rs | 12 +- core/src/lib.rs | 1 + core/src/messages.rs | 45 ++ db/00000000000001_initial_schema.up.sql | 1 + 20 files changed, 885 insertions(+), 560 deletions(-) create mode 100644 core/src/messages.rs diff --git a/aggregator/src/aggregator.rs b/aggregator/src/aggregator.rs index 598c68a49..62e1e83a0 100644 --- a/aggregator/src/aggregator.rs +++ b/aggregator/src/aggregator.rs @@ -43,6 +43,7 @@ use janus_core::{ auth_tokens::AuthenticationToken, hpke::{self, HpkeApplicationInfo, HpkeCiphersuite, Label}, http::ReqwestAuthenticationToken, + messages::DeriveId, retries::{HttpResponse, retry_http_request_notify}, time::{Clock, DateTimeExt, IntervalExt, TimeExt}, vdaf::{ @@ -74,7 +75,7 @@ use prio::{ flp::gadgets::{Mul, ParallelSum}, vdaf::prio3::{Prio3, Prio3Count, Prio3Histogram, Prio3Sum, Prio3SumVec}, }; -use rand::{RngExt, random, rng}; +use rand::{RngExt, rng}; use reqwest::Client; use tokio::{select, try_join}; use tracing::{Level, debug, error, info, warn}; @@ -636,10 +637,9 @@ impl Aggregator { async fn handle_create_collection_job( &self, task_id: &TaskId, - collection_job_id: &CollectionJobId, req_bytes: &[u8], auth_token: Option, - ) -> Result, Error> { + ) -> Result<(CollectionJobId, Option>), Error> { let task_aggregator = self .task_aggregators .get(task_id) @@ -656,7 +656,7 @@ impl Aggregator { } task_aggregator - .handle_create_collection_job(&self.datastore, collection_job_id, req_bytes) + .handle_create_collection_job(&self.datastore, req_bytes) .await } @@ -669,7 +669,7 @@ impl Aggregator { task_id: &TaskId, collection_job_id: &CollectionJobId, auth_token: Option, - ) -> Result, Error> { + ) -> Result>, Error> { let task_aggregator = self .task_aggregators .get(task_id) @@ -1206,16 +1206,10 @@ impl TaskAggregator { async fn handle_create_collection_job( &self, datastore: &Datastore, - collection_job_id: &CollectionJobId, req_bytes: &[u8], - ) -> Result, Error> { + ) -> Result<(CollectionJobId, Option>), Error> { self.vdaf_ops - .handle_create_collection_job( - datastore, - Arc::clone(&self.task), - collection_job_id, - req_bytes, - ) + .handle_create_collection_job(datastore, Arc::clone(&self.task), req_bytes) .await } @@ -1223,7 +1217,7 @@ impl TaskAggregator { &self, datastore: &Datastore, collection_job_id: &CollectionJobId, - ) -> Result, Error> { + ) -> Result>, Error> { self.vdaf_ops .handle_get_collection_job(datastore, Arc::clone(&self.task), collection_job_id) .await @@ -3136,9 +3130,8 @@ impl VdafOps { &self, datastore: &Datastore, task: Arc, - collection_job_id: &CollectionJobId, collection_req_bytes: &[u8], - ) -> Result, Error> { + ) -> Result<(CollectionJobId, Option>), Error> { match task.batch_mode() { task::BatchMode::TimeInterval => { vdaf_ops_dispatch!(self, (vdaf, VdafType, VERIFY_KEY_LENGTH) => { @@ -3146,7 +3139,7 @@ impl VdafOps { VERIFY_KEY_LENGTH, TimeInterval, VdafType, - >(datastore, task, Arc::clone(vdaf), collection_job_id, collection_req_bytes) + >(datastore, task, Arc::clone(vdaf), collection_req_bytes) .await }) } @@ -3156,7 +3149,7 @@ impl VdafOps { VERIFY_KEY_LENGTH, LeaderSelected, VdafType, - >(datastore, task, Arc::clone(vdaf), collection_job_id, collection_req_bytes) + >(datastore, task, Arc::clone(vdaf), collection_req_bytes) .await }) } @@ -3171,77 +3164,79 @@ impl VdafOps { datastore: &Datastore, task: Arc, vdaf: Arc, - collection_job_id: &CollectionJobId, req_bytes: &[u8], - ) -> Result, Error> { + ) -> Result<(CollectionJobId, Option>), Error> { let req = Arc::new(CollectionJobReq::::get_decoded(req_bytes).map_err(Error::MessageDecode)?); - Self::handle_create_collection_job_inner(datastore, task, vdaf, collection_job_id, req) - .await + Self::handle_create_collection_job_inner(datastore, task, vdaf, req).await } + /// Handle creation of a collection job, returning the job ID and possibly an encoded + /// `CollectionJobResp`. #[tracing::instrument( skip_all, fields(query = ?req.query().query_body()), err(level = Level::DEBUG) )] - async fn handle_create_collection_job_inner< - const SEED_SIZE: usize, - B: CollectableBatchMode, - A: AsyncAggregator, - >( + async fn handle_create_collection_job_inner( datastore: &Datastore, task: Arc, vdaf: Arc, - collection_job_id: &CollectionJobId, req: Arc>, - ) -> Result, Error> { + ) -> Result<(CollectionJobId, Option>), Error> + where + B: CollectableBatchMode, + A: AsyncAggregator, + { validate_collection_job_extensions(task.id(), req.extensions())?; + A::AggregationParam::get_decoded(req.aggregation_parameter()) + .map_err(Error::MessageDecode)?; - let aggregation_param = Arc::new( - A::AggregationParam::get_decoded(req.aggregation_parameter()) - .map_err(Error::MessageDecode)?, - ); - - datastore + let (collection_job_id, response) = datastore .run_tx("collect", move |tx| { - let (task, vdaf, collection_job_id, req, aggregation_param) = ( - Arc::clone(&task), - Arc::clone(&vdaf), - *collection_job_id, - Arc::clone(&req), - Arc::clone(&aggregation_param), - ); + let (task, vdaf, req) = (Arc::clone(&task), Arc::clone(&vdaf), Arc::clone(&req)); Box::pin(async move { - // Check if this collection job already exists, ensuring that all parameters - // match. - if let Some(collection_job) = tx - .get_collection_job::(&vdaf, task.id(), &collection_job_id) - .await? - { - if collection_job.query() == req.query() - && collection_job.aggregation_parameter() == aggregation_param.as_ref() - { + let mut jobs = B::collection_jobs_intersecting_query::( + tx, &vdaf, &task, &req, + ) + .await?; + // If there are any collection jobs intersecting the incoming query, then there + // are two possibilities: + // + // - there is a single intersecting job and it exactly matches the incoming + // query. It's not possible for any others jobs to intersect because the + // existence of the first job precludes them. We have a benign replayed + // request and serve up the job we already have. + // - there are one or more more intersecting jobs and none of them exactly match + // the incoming query. We fail with batch overlap. + // + // In either case, we don't need to examine anything but the first item in the + // iterator. + if let Some(job) = jobs.next() { + if *job.id() == req.derive_id()? { debug!( - collection_job_id = %collection_job_id, + collection_job_id = %job.id(), collect_request = ?req, "collection job already exists" ); - return Ok(()); + debug_assert!(jobs.next().is_none()); + return Ok(( + *job.id(), + Self::collection_job_resp_from_job(&task, &job) + .map_err(|e| datastore::Error::User(e.into()))?, + )); } else { return Err(datastore::Error::User( - Error::ForbiddenMutation { - resource_type: "collection job", - identifier: collection_job_id.to_string(), - } - .into(), + Error::BatchOverlap(*task.id()).into(), )); } } + // We found no collection jobs matching the query. Assign a batch identifier and + // create a new collection job. let collection_identifier = - B::collection_identifier_for_query(tx, &task, req.query()) + B::batch_identifier_for_query(tx, &task, req.query()) .await? .ok_or_else(|| { datastore::Error::User( @@ -3263,16 +3258,8 @@ impl VdafOps { } debug!(collect_request = ?req, "Cache miss, creating new collection job"); - let (_, report_count) = try_join!( - B::validate_query_count::( - tx, - &vdaf, - &task, - &collection_identifier, - &aggregation_param, - ), - B::count_client_reports(tx, &task, &collection_identifier), - )?; + let report_count = + B::count_client_reports(tx, &task, &collection_identifier).await?; // Batch size must be validated while handling CollectReq and hence before // creating a collection job. @@ -3283,23 +3270,20 @@ impl VdafOps { )); } - tx.put_collection_job(&CollectionJob::::new( - *task.id(), - collection_job_id, - random(), - req.query().clone(), - aggregation_param.as_ref().clone(), - collection_identifier, - CollectionJobState::Start, - )) - .await?; + let collection_job_id = tx + .put_collection_job_by_request::( + task.id(), + &req, + collection_identifier, + ) + .await?; - Ok(()) + Ok((collection_job_id, None)) }) }) .await?; - Ok(vec![0; 0]) + Ok((collection_job_id, response)) } /// Handle GET requests to the leader's `tasks/{task-id}/collection_jobs/{collection-job-id}` @@ -3315,7 +3299,7 @@ impl VdafOps { datastore: &Datastore, task: Arc, collection_job_id: &CollectionJobId, - ) -> Result, Error> { + ) -> Result>, Error> { match task.batch_mode() { task::BatchMode::TimeInterval => { vdaf_ops_dispatch!(self, (vdaf, VdafType, VERIFY_KEY_LENGTH) => { @@ -3340,17 +3324,17 @@ impl VdafOps { } } - // return value is an encoded CollectResp - async fn handle_get_collection_job_generic< - const SEED_SIZE: usize, - B: CollectableBatchMode, - A: AsyncAggregator, - >( + /// return value is an encoded `CollectionJobResp` + async fn handle_get_collection_job_generic( datastore: &Datastore, task: Arc, vdaf: Arc, collection_job_id: &CollectionJobId, - ) -> Result, Error> { + ) -> Result>, Error> + where + B: CollectableBatchMode, + A: AsyncAggregator, + { let collection_job = datastore .run_tx("get_collection_job", |tx| { let (task, vdaf, collection_job_id) = @@ -3368,23 +3352,35 @@ impl VdafOps { }) .await?; + Self::collection_job_resp_from_job(&task, &collection_job) + } + + /// return value is an encoded `CollectionJobResp` + fn collection_job_resp_from_job( + task: &AggregatorTask, + collection_job: &CollectionJob, + ) -> Result>, Error> + where + B: CollectableBatchMode, + A: AsyncAggregator, + { match collection_job.state() { CollectionJobState::Start => { debug!( - %collection_job_id, + collection_job_id = %collection_job.id(), task_id = %task.id(), "collection job has not run yet" ); - Ok(vec![0; 0]) + Ok(None) } CollectionJobState::Poll => { debug!( - %collection_job_id, + collection_job_id = %collection_job.id(), task_id = %task.id(), "collection job has not completed yet" ); - Ok(vec![0; 0]) + Ok(None) } CollectionJobState::Finished { @@ -3404,7 +3400,7 @@ impl VdafOps { // has been long enough since the encrypted helper share was cached -- tricky // thing is deciding what "long enough" is. debug!( - %collection_job_id, + collection_job_id = %collection_job.id(), task_id = %task.id(), "Serving cached collection job response" ); @@ -3439,17 +3435,19 @@ impl VdafOps { helper_encrypted_agg_share: encrypted_helper_aggregate_share.clone(), } .get_encoded() + .map(Some) .map_err(Error::MessageEncode) } CollectionJobState::Abandoned => Err(Error::AbandonedCollectionJob( *task.id(), - *collection_job_id, + *collection_job.id(), )), - CollectionJobState::Deleted => { - Err(Error::DeletedCollectionJob(*task.id(), *collection_job_id)) - } + CollectionJobState::Deleted => Err(Error::DeletedCollectionJob( + *task.id(), + *collection_job.id(), + )), } } diff --git a/aggregator/src/aggregator/batch_mode.rs b/aggregator/src/aggregator/batch_mode.rs index c08dcf0b3..710fcb27a 100644 --- a/aggregator/src/aggregator/batch_mode.rs +++ b/aggregator/src/aggregator/batch_mode.rs @@ -2,11 +2,14 @@ use async_trait::async_trait; use janus_aggregator_core::{ AsyncAggregator, batch_mode::{AccumulableBatchMode, CollectableBatchMode as CoreCollectableBatchMode}, - datastore::{self, Transaction, models::LeaderStoredReport}, + datastore::{ + self, Transaction, + models::{CollectionJob, LeaderStoredReport}, + }, task::AggregatorTask, }; use janus_messages::{ - Role, + CollectionJobReq, Role, batch_mode::{BatchMode, LeaderSelected, TimeInterval}, }; @@ -73,6 +76,11 @@ impl UploadableBatchMode for LeaderSelected { pub trait CollectableBatchMode: CoreCollectableBatchMode + AccumulableBatchMode { /// Validates query count for a given batch, per the size checks in /// . + /// + /// This method is exclusively used in the Helper's aggregate share handler. Once that moves to + /// draft-ietf-ppm-dap-20 style, server selected identifier handling, usage of this method will + /// be replaced with calls to a new method aggregate_share_jobs_intersecting_query and this + /// method will be deleted. async fn validate_query_count>( tx: &Transaction<'_>, vdaf: &A, @@ -80,6 +88,24 @@ pub trait CollectableBatchMode: CoreCollectableBatchMode + AccumulableBatchMode batch_identifier: &Self::BatchIdentifier, aggregation_param: &A::AggregationParam, ) -> Result<(), datastore::Error>; + + /// Find any collection jobs that intersect with the request's query. + /// + /// For time interval batch mode, multiple collection jobs could have queries that overlap with + /// the query's time interval. For leader selected batch mode, the query is a batch ID, so only + /// zero or one collection job will ever be returned. + async fn collection_jobs_intersecting_query< + const SEED_SIZE: usize, + A: AsyncAggregator, + >( + tx: &Transaction<'_>, + vdaf: &A, + task: &AggregatorTask, + request: &CollectionJobReq, + ) -> Result< + impl Iterator> + Send + Sync, + datastore::Error, + >; } #[async_trait] @@ -132,7 +158,7 @@ impl CollectableBatchMode for TimeInterval { // Check that all intersecting collect intervals are equal to this collect interval. if found_overlapping_nonequal_interval { return Err(datastore::Error::User( - Error::BatchOverlap(*task.id(), *collect_interval).into(), + Error::BatchOverlap(*task.id()).into(), )); } @@ -151,6 +177,28 @@ impl CollectableBatchMode for TimeInterval { } Ok(()) } + + async fn collection_jobs_intersecting_query< + const SEED_SIZE: usize, + A: AsyncAggregator, + >( + tx: &Transaction<'_>, + vdaf: &A, + task: &AggregatorTask, + request: &CollectionJobReq, + ) -> Result< + impl Iterator> + Send + Sync, + datastore::Error, + > { + Ok(tx + .get_collection_jobs_intersecting_interval::( + vdaf, + task.id(), + request.query().query_body(), + ) + .await? + .into_iter()) + } } #[async_trait] @@ -196,4 +244,22 @@ impl CollectableBatchMode for LeaderSelected { } Ok(()) } + + async fn collection_jobs_intersecting_query< + const SEED_SIZE: usize, + A: AsyncAggregator, + >( + tx: &Transaction<'_>, + vdaf: &A, + task: &AggregatorTask, + request: &CollectionJobReq, + ) -> Result< + impl Iterator> + Send + Sync, + datastore::Error, + > { + Ok(tx + .get_collection_job_by_query(vdaf, task.id(), request.query()) + .await? + .into_iter()) + } } diff --git a/aggregator/src/aggregator/collection_job_tests.rs b/aggregator/src/aggregator/collection_job_tests.rs index 2e18a8b46..6c27913c7 100644 --- a/aggregator/src/aggregator/collection_job_tests.rs +++ b/aggregator/src/aggregator/collection_job_tests.rs @@ -1,4 +1,4 @@ -use std::{collections::HashSet, sync::Arc}; +use std::sync::Arc; use assert_matches::assert_matches; use axum::{Router, body::Body}; @@ -8,7 +8,8 @@ use janus_aggregator_core::{ Datastore, models::{ AggregationJob, AggregationJobState, BatchAggregation, BatchAggregationState, - CollectionJobState, LeaderStoredReport, ReportAggregation, ReportAggregationState, + CollectionJob, CollectionJobState, LeaderStoredReport, ReportAggregation, + ReportAggregationState, }, test_util::{EphemeralDatastore, ephemeral_datastore}, }, @@ -21,13 +22,14 @@ use janus_aggregator_core::{ use janus_core::{ auth_tokens::{AuthenticationToken, test_util::WithAuthenticationToken}, hpke::{self, HpkeApplicationInfo, Label}, + retries::HeaderMapExt, test_util::{install_test_trace_subscriber, runtime::TestRuntime}, time::{Clock, DateTimeExt, MockClock}, vdaf::VdafInstance, }; use janus_messages::{ AggregateShareAad, AggregationJobStep, BatchId, CollectionJobId, CollectionJobReq, - CollectionJobResp, Interval, MediaType, Query, Role, Time, TimePrecision, + CollectionJobResp, Interval, MediaType, Query, Role, TaskId, Time, TimePrecision, batch_mode::{ BatchMode as BatchModeTrait, LeaderSelected, LeaderSelectedQueryConfig, TimeInterval, }, @@ -56,20 +58,14 @@ pub(crate) struct CollectionJobTestCase { } impl CollectionJobTestCase { - pub(super) async fn put_collection_job_with_auth_token( + pub(super) async fn post_collection_job_with_auth_token( &self, - collection_job_id: &CollectionJobId, request: &CollectionJobReq, auth_token: Option<&AuthenticationToken>, ) -> http::Response { let mut builder = Request::builder() - .method("PUT") - .uri( - self.task - .collection_job_uri(collection_job_id) - .unwrap() - .path(), - ) + .method("POST") + .uri(self.task.collection_job_creation_uri().unwrap().path()) .header( http::header::CONTENT_TYPE, CollectionJobReq::::MEDIA_TYPE, @@ -88,17 +84,12 @@ impl CollectionJobTestCase { .unwrap() } - pub(super) async fn put_collection_job( + pub(super) async fn post_collection_job( &self, - collection_job_id: &CollectionJobId, request: &CollectionJobReq, ) -> http::Response { - self.put_collection_job_with_auth_token( - collection_job_id, - request, - Some(self.task.collector_auth_token()), - ) - .await + self.post_collection_job_with_auth_token(request, Some(self.task.collector_auth_token())) + .await } pub(super) async fn get_collection_job_with_auth_token( @@ -263,6 +254,19 @@ impl CollectionJobTestCase { } } +/// Parse task ID and collection job ID out of response location header. +/// +/// Panics if header is missing from response or value is malformed. +pub(super) fn ids_from_collection_job_response( + response: &http::Response, +) -> (TaskId, CollectionJobId) { + response + .headers() + .ids_from_collection_job_location() + .unwrap() + .unwrap() +} + pub(crate) async fn setup_collection_job_test_case( role: Role, batch_mode: BatchMode, @@ -336,35 +340,37 @@ async fn setup_leader_selected_current_batch_collection_job_test_case() #[tokio::test] async fn collection_job_success_leader_selected() { // This test drives two current batch collection jobs to completion, verifying that distinct - // batch IDs are collected each time. Then, we attempt to collect another current batch, which - // must fail as there are no more outstanding batches. + // collection job IDs are minted and distinct batch IDs are collected each time. Then, we + // attempt to collect another current batch, which must fail as there are no more outstanding + // batches. let (test_case, batch_id_1, batch_id_2, spanned_interval) = setup_leader_selected_current_batch_collection_job_test_case().await; + let first_idempotency_key: LeaderSelectedQueryConfig = random(); let mut saw_batch_id_1 = false; let mut saw_batch_id_2 = false; + let mut collection_job_ids = vec![]; + let mut collection_job_resps = vec![]; let vdaf = dummy::Vdaf::new(1); let leader_aggregate_share = dummy::AggregateShare(0); let helper_aggregate_share = dummy::AggregateShare(1); let aggregation_param = dummy::AggregationParam::default(); - let request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), - aggregation_param.get_encoded().unwrap(), - ); - - for _ in 0..2 { - let collection_job_id: CollectionJobId = random(); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); + for idempotency_key in [first_idempotency_key.clone(), random()] { + let request = CollectionJobReq::new( + Query::new_leader_selected(idempotency_key), + aggregation_param.get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (task_id, collection_job_id) = ids_from_collection_job_response(&response); + collection_job_ids.push(collection_job_id); + assert_eq!(test_case.task.id(), &task_id); assert!(take_response_body(&mut response).await.is_empty()); let mut response = test_case.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); assert!(take_response_body(&mut response).await.is_empty()); // Update the collection job with the aggregate shares. collection job should now be @@ -375,6 +381,7 @@ async fn collection_job_success_leader_selected() { let task = test_case.task.clone(); let vdaf = vdaf.clone(); let helper_aggregate_share_bytes = helper_aggregate_share.get_encoded().unwrap(); + let request = request.clone(); Box::pin(async move { let collection_job = tx .get_collection_job::<0, LeaderSelected, dummy::Vdaf>( @@ -395,18 +402,9 @@ async fn collection_job_success_leader_selected() { &Role::Collector, ), &helper_aggregate_share_bytes, - &AggregateShareAad::new( - *task.id(), - task.task_configuration(), - CollectionJobReq::new( - Query::new_leader_selected( - LeaderSelectedQueryConfig::test_default(), - ), - aggregation_param.get_encoded().unwrap(), - ), - ) - .get_encoded() - .unwrap(), + &AggregateShareAad::new(*task.id(), task.task_configuration(), request) + .get_encoded() + .unwrap(), ) .unwrap(); @@ -448,7 +446,7 @@ async fn collection_job_success_leader_selected() { leader_encrypted_aggregate_share, helper_encrypted_aggregate_share, ) = assert_matches!( - collect_resp, + collect_resp.clone(), CollectionJobResp { report_count, interval, @@ -473,12 +471,7 @@ async fn collection_job_success_leader_selected() { &AggregateShareAad::new( *test_case.task.id(), test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), - aggregation_param.get_encoded().unwrap(), - ), + request.clone(), ) .get_encoded() .unwrap(), @@ -496,12 +489,7 @@ async fn collection_job_success_leader_selected() { &AggregateShareAad::new( *test_case.task.id(), test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), - aggregation_param.get_encoded().unwrap(), - ), + request.clone(), ) .get_encoded() .unwrap(), @@ -511,16 +499,22 @@ async fn collection_job_success_leader_selected() { helper_aggregate_share, dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() ); + + collection_job_resps.push(collect_resp); } assert!(saw_batch_id_1 && saw_batch_id_2); + // Collection job IDs should be distinct + assert_eq!(collection_job_ids.len(), 2); + assert_ne!(collection_job_ids[0], collection_job_ids[1]); // We have run the two ready batches to completion. Further attempts to collect current batch // ought to fail. - let collection_job_id: CollectionJobId = random(); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) + .post_collection_job(&CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + )) .await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( @@ -532,207 +526,171 @@ async fn collection_job_success_leader_selected() { "taskid": format!("{}", test_case.task.id()), }), ); + + // Re-send the very first collection job creation request. We should get the representation of + // the completed collection job without needing a further GET. + let request = CollectionJobReq::new( + Query::new_leader_selected(first_idempotency_key), + aggregation_param.get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id, collection_job_ids[0]); + assert_eq!( + response.headers().get("content-type").unwrap(), + CollectionJobResp::MEDIA_TYPE, + ); + decode_response_body::(&mut response).await; } #[tokio::test] -async fn collection_job_put_idempotence_time_interval() { +async fn collection_job_post_idempotence_time_interval() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + // Seed DB with reports, aggregations, etc. corresponding to the two time intervals used later + // in the test test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, &TimePrecision::from_seconds(1), )) .await; + test_case + .setup_time_interval_batch(Time::from_seconds_since_epoch( + 3, + &TimePrecision::from_seconds(1), + )) + .await; - let collection_job_id = random(); - let request = CollectionJobReq::new( - Query::new_time_interval( - Interval::minimal(Time::from_seconds_since_epoch( - 0, - test_case.task.time_precision(), - )) - .unwrap(), - ), - dummy::AggregationParam::default().get_encoded().unwrap(), - ); - + // We should get the same collection job ID back both times + let mut collection_job_id = None; for _ in 0..2 { let response = test_case - .put_collection_job(&collection_job_id, &request) + .post_collection_job(&CollectionJobReq::new( + Query::new_time_interval( + Interval::minimal(Time::from_seconds_since_epoch( + 0, + test_case.task.time_precision(), + )) + .unwrap(), + ), + dummy::AggregationParam::default().get_encoded().unwrap(), + )) .await; - assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.status(), StatusCode::OK); + let (_, created_collection_job_id) = ids_from_collection_job_response(&response); + if let Some(collection_job_id) = collection_job_id { + assert_eq!(created_collection_job_id, collection_job_id); + } else { + collection_job_id = Some(created_collection_job_id); + } } - // There should only be a single collection job despite two successful PUTs - test_case - .datastore - .run_unnamed_tx(|tx| { - let task_id = *test_case.task.id(); - let vdaf = dummy::Vdaf::new(1); - Box::pin(async move { - let collection_jobs = tx - .get_collection_jobs_for_task::<0, TimeInterval, dummy::Vdaf>(&vdaf, &task_id) - .await - .unwrap(); - assert_eq!(collection_jobs.len(), 1); - assert_eq!(collection_jobs[0].id(), &collection_job_id); - - Ok(()) - }) - }) - .await - .unwrap(); -} - -#[tokio::test] -async fn collection_job_put_idempotence_time_interval_varied_collection_id() { - // This test sends repeated, identical collection requests with differing collection job IDs and - // validates that they are accepted. They should be accepted because the check for repeated - // collection should allow a batch to be repeatedly collected as long as it uses the same - // aggregation parameter each time. - - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; - test_case - .setup_time_interval_batch(Time::from_seconds_since_epoch( - 0, - test_case.task.time_precision(), + // Send another request, with non-overlapping time interval. We should create a distinct + // collection job. + let response = test_case + .post_collection_job(&CollectionJobReq::new( + Query::new_time_interval( + Interval::minimal(Time::from_seconds_since_epoch( + 3 * test_case.task.time_precision().as_seconds(), + test_case.task.time_precision(), + )) + .unwrap(), + ), + dummy::AggregationParam::default().get_encoded().unwrap(), )) .await; + assert_eq!(response.status(), StatusCode::OK); + let (_, next_collection_job_id) = ids_from_collection_job_response(&response); + assert_ne!(collection_job_id.unwrap(), next_collection_job_id); - let collection_job_ids = HashSet::from(random::<[CollectionJobId; 2]>()); - let request = CollectionJobReq::new( - Query::new_time_interval( - Interval::minimal(Time::from_seconds_since_epoch( - 0, - test_case.task.time_precision(), - )) - .unwrap(), - ), - dummy::AggregationParam::default().get_encoded().unwrap(), - ); - - for collection_job_id in &collection_job_ids { - let response = test_case - .put_collection_job(collection_job_id, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); - } + // There should only be two collection jobs despite three successful POSTs + let mut expect_collection_job_ids = vec![collection_job_id.unwrap(), next_collection_job_id]; + expect_collection_job_ids.sort(); - test_case + let got_collection_job_ids = test_case .datastore .run_unnamed_tx(|tx| { let task_id = *test_case.task.id(); - let collection_job_ids = collection_job_ids.clone(); - + let vdaf = dummy::Vdaf::new(1); Box::pin(async move { - let collection_jobs = tx - .get_collection_jobs_for_task::<0, TimeInterval, dummy::Vdaf>( - &dummy::Vdaf::new(1), - &task_id, - ) + let mut collection_job_ids: Vec<_> = tx + .get_collection_jobs_for_task::<0, TimeInterval, dummy::Vdaf>(&vdaf, &task_id) .await - .unwrap(); - - assert_eq!(collection_jobs.len(), 2); - assert_eq!( - collection_jobs - .into_iter() - .map(|job| *job.id()) - .collect::>(), - collection_job_ids - ); + .unwrap() + .iter() + .map(CollectionJob::id) + .cloned() + .collect(); + collection_job_ids.sort(); - Ok(()) + Ok(collection_job_ids) }) }) .await .unwrap(); -} - -#[tokio::test] -async fn collection_job_put_idempotence_time_interval_mutate_time_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; - test_case - .setup_time_interval_batch(Time::from_time_precision_units(0)) - .await; - - let collection_job_id = random(); - let request = CollectionJobReq::new( - Query::new_time_interval(Interval::minimal(Time::from_time_precision_units(0)).unwrap()), - dummy::AggregationParam::default().get_encoded().unwrap(), - ); - - let response = test_case - .put_collection_job(&collection_job_id, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); - let mutated_request = CollectionJobReq::new( - Query::new_time_interval( - Interval::minimal(Time::from_seconds_since_epoch( - test_case.task.time_precision().as_seconds(), - test_case.task.time_precision(), - )) - .unwrap(), - ), - dummy::AggregationParam::default().get_encoded().unwrap(), - ); - - let response = test_case - .put_collection_job(&collection_job_id, &mutated_request) - .await; - assert_eq!(response.status(), StatusCode::CONFLICT); + assert_eq!(got_collection_job_ids, expect_collection_job_ids); } #[tokio::test] -async fn collection_job_put_idempotence_time_interval_mutate_aggregation_param() { +async fn collection_job_post_idempotence_time_interval_mutate_aggregation_param() { + // Create a collection job for some time interval, then try to create another, varying the + // aggregation parameter. It should fail due to batch overlap. let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_time_precision_units(0)) .await; - let collection_job_id = random(); let request = CollectionJobReq::new( Query::new_time_interval(Interval::minimal(Time::from_time_precision_units(0)).unwrap()), dummy::AggregationParam(0).get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); let mutated_request = CollectionJobReq::new( Query::new_time_interval(Interval::minimal(Time::from_time_precision_units(0)).unwrap()), dummy::AggregationParam(1).get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &mutated_request) - .await; - assert_eq!(response.status(), StatusCode::CONFLICT); + let mut response = test_case.post_collection_job(&mutated_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); } #[tokio::test] -async fn collection_job_put_idempotence_leader_selected() { +async fn collection_job_post_idempotence_leader_selected() { let (test_case, batch_id_1, batch_id_2, _) = setup_leader_selected_current_batch_collection_job_test_case().await; - let collection_job_id = random(); - let request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), - dummy::AggregationParam(0).get_encoded().unwrap(), - ); let mut seen_batch_id = None; + let mut collection_job_id = None; + let idempotency_key: LeaderSelectedQueryConfig = random(); for _ in 0..2 { let response = test_case - .put_collection_job(&collection_job_id, &request) + .post_collection_job(&CollectionJobReq::new( + Query::new_leader_selected(idempotency_key.clone()), + dummy::AggregationParam(0).get_encoded().unwrap(), + )) .await; - - assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.status(), StatusCode::OK); + let (_, created_collection_job_id) = ids_from_collection_job_response(&response); + if let Some(collection_job_id) = collection_job_id { + assert_eq!(created_collection_job_id, collection_job_id); + } else { + collection_job_id = Some(created_collection_job_id); + } // Make sure that there is only ever a single collection job, and that it uses the same // batch ID after each PUT @@ -741,15 +699,15 @@ async fn collection_job_put_idempotence_leader_selected() { .run_unnamed_tx(|tx| { let task_id = *test_case.task.id(); Box::pin(async move { - let vdaf = dummy::Vdaf::new(1); let collection_jobs = tx .get_collection_jobs_for_task::<0, LeaderSelected, dummy::Vdaf>( - &vdaf, &task_id, + &dummy::Vdaf::new(1), + &task_id, ) .await .unwrap(); assert_eq!(collection_jobs.len(), 1); - assert_eq!(collection_jobs[0].id(), &collection_job_id); + assert_eq!(collection_jobs[0].id(), &collection_job_id.unwrap()); assert!( collection_jobs[0].batch_identifier().eq(&batch_id_1) || collection_jobs[0].batch_identifier().eq(&batch_id_2) @@ -765,83 +723,122 @@ async fn collection_job_put_idempotence_leader_selected() { Some(seen_batch_id) => assert_eq!(seen_batch_id, batch_id), } } + + // POST another collection job with a different idempotency key. We should create a distinct + // collection job. + let response = test_case + .post_collection_job(&CollectionJobReq::new( + Query::new_leader_selected(random()), + dummy::AggregationParam(0).get_encoded().unwrap(), + )) + .await; + assert_eq!(response.status(), StatusCode::OK); + let (_, next_collection_job_id) = ids_from_collection_job_response(&response); + assert_ne!(next_collection_job_id, collection_job_id.unwrap()); + + // There should only be two collection jobs despite three successful POSTs + let mut expect_collection_job_ids = vec![collection_job_id.unwrap(), next_collection_job_id]; + expect_collection_job_ids.sort(); + + let got_collection_job_ids = test_case + .datastore + .run_unnamed_tx(|tx| { + let task_id = *test_case.task.id(); + Box::pin(async move { + let mut collection_job_ids: Vec<_> = tx + .get_collection_jobs_for_task::<0, LeaderSelected, dummy::Vdaf>( + &dummy::Vdaf::new(1), + &task_id, + ) + .await + .unwrap() + .iter() + .map(CollectionJob::id) + .cloned() + .collect(); + collection_job_ids.sort(); + + Ok(collection_job_ids) + }) + }) + .await + .unwrap(); + + assert_eq!(got_collection_job_ids, expect_collection_job_ids); } #[tokio::test] -async fn collection_job_put_idempotence_leader_selected_mutate_aggregation_param() { +async fn collection_job_post_idempotence_leader_selected_mutate_aggregation_param() { + // Create a collection job for some idempotency key, then try to create another, varying the + // aggregation parameter. It should fail because the batch is marked as aggregated. let (test_case, _, _, _) = setup_leader_selected_current_batch_collection_job_test_case().await; - let collection_job_id = random(); + let idempotency_key: LeaderSelectedQueryConfig = random(); + let request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(idempotency_key.clone()), dummy::AggregationParam(0).get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &request) - .await; - - assert_eq!(response.status(), StatusCode::CREATED); + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); let mutated_request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(idempotency_key), dummy::AggregationParam(1).get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &mutated_request) - .await; - assert_eq!(response.status(), StatusCode::CONFLICT); + let mut response = test_case.post_collection_job(&mutated_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); } #[tokio::test] -async fn collection_job_put_idempotence_leader_selected_no_extra_reports() { +async fn collection_job_post_idempotence_leader_selected_no_extra_reports() { let (test_case, _batch_id_1, _batch_id_2, _) = setup_leader_selected_current_batch_collection_job_test_case().await; - let collection_job_id_1 = random(); - let collection_job_id_2 = random(); - let request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + let request_1 = CollectionJobReq::new( + Query::new_leader_selected(random()), dummy::AggregationParam(0).get_encoded().unwrap(), ); - // Create the first collection job. - let response = test_case - .put_collection_job(&collection_job_id_1, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); - - // Fetch the first collection job, to advance the current batch. + // Create and then fetch the first collection job. + let response = test_case.post_collection_job(&request_1).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id_1) = ids_from_collection_job_response(&response); let response = test_case.get_collection_job(&collection_job_id_1).await; assert_eq!(response.status(), StatusCode::OK); - // Create the second collection job. - let response = test_case - .put_collection_job(&collection_job_id_2, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); + let request_2 = CollectionJobReq::new( + Query::new_leader_selected(random()), + dummy::AggregationParam(0).get_encoded().unwrap(), + ); - // Fetch the second collection job, to advance the current batch. There are now no outstanding - // batches left. + // Create and then fetch the second collection job. There are now no outstanding batches left. + let response = test_case.post_collection_job(&request_2).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id_2) = ids_from_collection_job_response(&response); let response = test_case.get_collection_job(&collection_job_id_2).await; assert_eq!(response.status(), StatusCode::OK); // Re-send the collection job creation requests to confirm they are still idempotent. - let response = test_case - .put_collection_job(&collection_job_id_1, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); - let response = test_case - .put_collection_job(&collection_job_id_2, &request) - .await; - assert_eq!(response.status(), StatusCode::CREATED); + let response = test_case.post_collection_job(&request_1).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id_1_again) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id_1, collection_job_id_1_again); + let response = test_case.post_collection_job(&request_2).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id_2_again) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id_2, collection_job_id_2_again); } #[tokio::test] @@ -855,17 +852,12 @@ async fn collection_job_batch_mode_misaligned() { .await; // LeaderSelected != TimeInterval - let collection_job_id = random(); let request = CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), dummy::AggregationParam(0).get_encoded().unwrap(), ); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) - .await; + let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( take_problem_details(&mut response).await, diff --git a/aggregator/src/aggregator/error.rs b/aggregator/src/aggregator/error.rs index 740d0ace0..490aa209a 100644 --- a/aggregator/src/aggregator/error.rs +++ b/aggregator/src/aggregator/error.rs @@ -12,7 +12,7 @@ use janus_aggregator_core::{ use janus_core::http::HttpErrorResponse; use janus_messages::{ AggregateShareId, AggregationJobId, AggregationJobStep, CollectionJobId, HpkeConfigId, - Interval, ReportError, ReportId, ReportIdChecksum, Role, TaskId, Time, + ReportError, ReportId, ReportIdChecksum, Role, TaskId, Time, }; use opentelemetry::{KeyValue, metrics::Counter}; use prio::{topology::ping_pong::PingPongError, vdaf::VdafError}; @@ -115,8 +115,8 @@ pub enum Error { BatchMismatch(Box), /// A collect or aggregate share request was rejected because the batch overlaps with a /// previously collected one. - #[error("task {0}: queried batch {1} overlaps with previously collected batch(es)")] - BatchOverlap(TaskId, Interval), + #[error("task {0}: queried batch overlaps with previously collected batch(es)")] + BatchOverlap(TaskId), /// HPKE failure. #[error("HPKE error: {0}")] Hpke(#[from] janus_core::hpke::Error), @@ -339,7 +339,7 @@ impl Error { Error::InvalidBatchSize(_, _) => "invalid_batch_size", Error::Url(_) => "url", Error::BatchMismatch { .. } => "batch_mismatch", - Error::BatchOverlap(_, _) => "batch_overlap", + Error::BatchOverlap(_) => "batch_overlap", Error::Hpke(_) => "hpke", Error::TaskParameters(_) => "task_parameters", Error::HttpClient(_) => "http_client", diff --git a/aggregator/src/aggregator/garbage_collector.rs b/aggregator/src/aggregator/garbage_collector.rs index df1d9b378..673e64b48 100644 --- a/aggregator/src/aggregator/garbage_collector.rs +++ b/aggregator/src/aggregator/garbage_collector.rs @@ -196,7 +196,7 @@ mod tests { use janus_messages::{ AggregationJobStep, CollectionJobReq, Duration, HpkeCiphertext, HpkeConfigId, Interval, Query, ReportIdChecksum, ReportMetadata, ReportShare, Role, TimePrecision, - batch_mode::{LeaderSelected, LeaderSelectedQueryConfig, TimeInterval}, + batch_mode::{LeaderSelected, TimeInterval}, }; use prio::{codec::Encode, vdaf::dummy}; use rand::random; @@ -677,9 +677,7 @@ mod tests { *task.id(), random(), random(), - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), dummy::AggregationParam(0), batch_id, CollectionJobState::Start, @@ -895,9 +893,7 @@ mod tests { 0, ReportIdChecksum::default(), CollectionJobReq::new( - Query::new_leader_selected( - LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), dummy::AggregationParam(0).get_encoded().unwrap(), ), ), diff --git a/aggregator/src/aggregator/http_handlers.rs b/aggregator/src/aggregator/http_handlers.rs index 585660eac..695525627 100644 --- a/aggregator/src/aggregator/http_handlers.rs +++ b/aggregator/src/aggregator/http_handlers.rs @@ -6,7 +6,7 @@ use axum::{ extract::{FromRequestParts, Path, State as AxumState}, middleware, response::{IntoResponse, Response}, - routing::{post, put}, + routing::{get, post, put}, }; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use bytes::Bytes; @@ -24,7 +24,7 @@ use janus_aggregator_core::{ use janus_core::{ Runtime, auth_tokens::{AuthenticationToken, DAP_AUTH_HEADER}, - http::{check_content_type, extract_bearer_token}, + http::{check_content_type, collection_job_location, extract_bearer_token}, taskprov::TASKPROV_HEADER, }; use janus_messages::{ @@ -157,7 +157,7 @@ impl Error { .with_task_id(task_id) .into_response() } - Error::BatchOverlap(task_id, _) => { + Error::BatchOverlap(task_id) => { ProblemDocument::new_dap(DapProblemType::BatchOverlap) .with_task_id(task_id) .into_response() @@ -318,7 +318,7 @@ impl IntoResponse for EncodedBody { /// An empty body with retry-after and location headers. #[derive(Clone)] struct EmptyBody { - location: String, + location: (HeaderName, String), } impl EmptyBody { @@ -330,7 +330,17 @@ impl EmptyBody { step: u16, ) -> Self { Self { - location: format!("/tasks/{task_id}/aggregation_jobs/{aggregation_job_id}?step={step}"), + location: ( + http::header::LOCATION, + format!("/tasks/{task_id}/aggregation_jobs/{aggregation_job_id}?step={step}"), + ), + } + } + + /// Construct an `EmptyBody` for a collection job. + fn for_collection_job(task_id: &TaskId, collection_job_id: &CollectionJobId) -> Self { + Self { + location: collection_job_location(task_id, collection_job_id), } } } @@ -342,7 +352,7 @@ impl IntoResponse for EmptyBody { [ // TODO(#3921): make Retry-After configurable instead of hardcoding. (http::header::RETRY_AFTER, "2".to_string()), - (http::header::LOCATION, self.location), + self.location, ], ) .into_response() @@ -351,8 +361,6 @@ impl IntoResponse for EmptyBody { pub(crate) static AGGREGATION_JOB_ROUTE: &str = "/tasks/{task_id}/aggregation_jobs/{aggregation_job_id}"; -pub(crate) static COLLECTION_JOB_ROUTE: &str = - "/tasks/{task_id}/collection_jobs/{collection_job_id}"; pub(crate) static AGGREGATE_SHARES_ROUTE: &str = "/tasks/{task_id}/aggregate_shares/{aggregate_share_id}"; @@ -585,10 +593,12 @@ impl<'a> AggregatorHandlerBuilder<'a> { ) .merge(aggregation_job_routes) .route( - COLLECTION_JOB_ROUTE, - put(collection_jobs_put) - .get(collection_jobs_get) - .delete(collection_jobs_delete), + "/tasks/{task_id}/collection_jobs", + post(collection_jobs_post), + ) + .route( + "/tasks/{task_id}/collection_jobs/{collection_job_id}", + get(collection_jobs_get).delete(collection_jobs_delete), ) .route( AGGREGATE_SHARES_ROUTE, @@ -874,29 +884,38 @@ async fn aggregation_jobs_delete( Ok(StatusCode::NO_CONTENT) } -/// Axum handler for the "/tasks/.../collection_jobs/..." PUT endpoint. -async fn collection_jobs_put( +/// Axum handler for the "/tasks/.../collection_jobs" POST endpoint. +async fn collection_jobs_post( headers: HeaderMap, - path: CollectionJobPath, + Path(task_id): Path, AxumState(aggregator): AxumState>, body: Bytes, ) -> Result { validate_content_type_headers::>(&headers)?; - let auth_token = parse_auth_token_from_headers(&path.task_id, &headers)?; - let response_bytes = aggregator - .handle_create_collection_job(&path.task_id, &path.collection_job_id, &body, auth_token) + let task_id: TaskId = task_id + .parse() + .map_err(|_| Error::BadRequest("invalid TaskId".into()))?; + + let auth_token = parse_auth_token_from_headers(&task_id, &headers)?; + let (collection_job_id, collection_job_resp) = aggregator + .handle_create_collection_job(&task_id, &body, auth_token) .await?; - Ok(( - StatusCode::CREATED, - [( - CONTENT_TYPE, - HeaderValue::from_static(CollectionJobResp::MEDIA_TYPE), - )], - response_bytes, - ) - .into_response()) + let resp = match collection_job_resp { + Some(response_bytes) => ( + StatusCode::OK, + [ + (CONTENT_TYPE, CollectionJobResp::MEDIA_TYPE.to_owned()), + collection_job_location(&task_id, &collection_job_id), + ], + response_bytes, + ) + .into_response(), + None => EmptyBody::for_collection_job(&task_id, &collection_job_id).into_response(), + }; + + Ok(resp) } /// Axum handler for the "/tasks/.../collection_jobs/..." GET endpoint. @@ -906,19 +925,24 @@ async fn collection_jobs_get( AxumState(aggregator): AxumState>, ) -> Result { let auth_token = parse_auth_token_from_headers(&path.task_id, &headers)?; - let response_bytes = aggregator + let resp = match aggregator .handle_get_collection_job(&path.task_id, &path.collection_job_id, auth_token) - .await?; - - Ok(( - StatusCode::OK, - [( - CONTENT_TYPE, - HeaderValue::from_static(CollectionJobResp::MEDIA_TYPE), - )], - response_bytes, - ) - .into_response()) + .await? + { + Some(response_bytes) => ( + StatusCode::OK, + [( + CONTENT_TYPE, + HeaderValue::from_static(CollectionJobResp::MEDIA_TYPE), + )], + response_bytes, + ) + .into_response(), + None => { + EmptyBody::for_collection_job(&path.task_id, &path.collection_job_id).into_response() + } + }; + Ok(resp) } /// Axum handler for the "/tasks/.../collection_jobs/..." DELETE endpoint. diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs b/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs index b37cca6b9..5a2bcbcb3 100644 --- a/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs +++ b/aggregator/src/aggregator/http_handlers/tests/aggregate_share.rs @@ -31,6 +31,7 @@ use prio::{ codec::{Decode, Encode}, vdaf::dummy, }; +use rand::random; use serde_json::json; use tower::ServiceExt; @@ -1057,9 +1058,7 @@ async fn aggregate_share_request_get_poll_after_put_leader_selected() { let request = AggregateShareReq::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param.get_encoded().unwrap(), ), BatchSelector::new_leader_selected(batch_id), diff --git a/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs b/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs index 347e33131..9351608e8 100644 --- a/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs +++ b/aggregator/src/aggregator/http_handlers/tests/aggregation_job_init.rs @@ -1526,12 +1526,7 @@ async fn aggregate_init_partially_replayed_aggregation_init() { } let request = AggregateShareReq::new( - CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), - agg_param.clone(), - ), + CollectionJobReq::new(Query::new_leader_selected(random()), agg_param.clone()), BatchSelector::new_leader_selected(batch_id), 5, ReportIdChecksum::from_report_ids(&report_ids), diff --git a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs index 4bb7427fa..9c1f9fde9 100644 --- a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs +++ b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs @@ -25,14 +25,14 @@ use serde_json::json; use tower::ServiceExt; use crate::aggregator::{ - collection_job_tests::setup_collection_job_test_case, + collection_job_tests::{ids_from_collection_job_response, setup_collection_job_test_case}, http_handlers::test_util::{ HttpHandlerTest, decode_response_body, take_problem_details, take_response_body, }, }; #[tokio::test] -async fn collection_job_put_request_to_helper() { +async fn collection_job_post_request_to_helper() { let test_case = setup_collection_job_test_case(Role::Helper, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -41,7 +41,6 @@ async fn collection_job_put_request_to_helper() { )) .await; - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval( Interval::minimal(Time::from_seconds_since_epoch( @@ -54,7 +53,7 @@ async fn collection_job_put_request_to_helper() { ); let mut response = test_case - .put_collection_job_with_auth_token(&collection_job_id, &request, Some(&random())) + .post_collection_job_with_auth_token(&request, Some(&random())) .await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); @@ -70,7 +69,7 @@ async fn collection_job_put_request_to_helper() { } #[tokio::test] -async fn collection_job_put_request_invalid_batch_interval() { +async fn collection_job_post_request_invalid_batch_interval() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -79,7 +78,6 @@ async fn collection_job_put_request_invalid_batch_interval() { )) .await; - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval( Interval::new( @@ -95,9 +93,7 @@ async fn collection_job_put_request_invalid_batch_interval() { dummy::AggregationParam::default().get_encoded().unwrap(), ); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) - .await; + let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( @@ -112,7 +108,7 @@ async fn collection_job_put_request_invalid_batch_interval() { } #[tokio::test] -async fn collection_job_put_request_invalid_aggregation_parameter() { +async fn collection_job_post_request_invalid_aggregation_parameter() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -121,7 +117,6 @@ async fn collection_job_put_request_invalid_aggregation_parameter() { )) .await; - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval( Interval::minimal(Time::from_seconds_since_epoch( @@ -135,9 +130,7 @@ async fn collection_job_put_request_invalid_aggregation_parameter() { Vec::from([0u8, 0u8]), ); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) - .await; + let mut response = test_case.post_collection_job(&request).await; // Collect request will be rejected because the aggregation parameter can't be decoded assert_eq!(response.status(), StatusCode::BAD_REQUEST); @@ -152,7 +145,7 @@ async fn collection_job_put_request_invalid_aggregation_parameter() { } #[tokio::test] -async fn collection_job_put_request_unsupported_extension() { +async fn collection_job_post_request_unsupported_extension() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; // No collection job extension types are defined, so any extension is unsupported. @@ -171,7 +164,7 @@ async fn collection_job_put_request_unsupported_extension() { Vec::new(), )])); - let mut response = test_case.put_collection_job(&random(), &request).await; + let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( @@ -187,7 +180,7 @@ async fn collection_job_put_request_unsupported_extension() { } #[tokio::test] -async fn collection_job_put_request_invalid_extension_order() { +async fn collection_job_post_request_invalid_extension_order() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; // Extension types 0x0002 then 0x0001 are not in strictly increasing order. @@ -206,7 +199,7 @@ async fn collection_job_put_request_invalid_extension_order() { CollectionJobExtension::new(CollectionJobExtensionType::Unknown(0x0001), Vec::new()), ])); - let mut response = test_case.put_collection_job(&random(), &request).await; + let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( @@ -222,7 +215,7 @@ async fn collection_job_put_request_invalid_extension_order() { } #[tokio::test] -async fn collection_job_put_request_invalid_batch_size() { +async fn collection_job_post_request_invalid_batch_size() { let HttpHandlerTest { router, ephemeral_datastore: _ephemeral_datastore, @@ -241,7 +234,6 @@ async fn collection_job_put_request_invalid_batch_size() { let leader_task = task.leader_view().unwrap(); datastore.put_aggregator_task(&leader_task).await.unwrap(); - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval(Interval::minimal(Time::from_time_precision_units(0)).unwrap()), dummy::AggregationParam::default().get_encoded().unwrap(), @@ -251,8 +243,8 @@ async fn collection_job_put_request_invalid_batch_size() { .clone() .oneshot( Request::builder() - .method("PUT") - .uri(task.collection_job_uri(&collection_job_id).unwrap().path()) + .method("POST") + .uri(task.collection_job_creation_uri().unwrap().path()) .with_authentication_token(task.collector_auth_token()) .header( http::header::CONTENT_TYPE, @@ -278,7 +270,7 @@ async fn collection_job_put_request_invalid_batch_size() { } #[tokio::test] -async fn collection_job_put_request_unauthenticated() { +async fn collection_job_post_request_unauthenticated() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -292,7 +284,6 @@ async fn collection_job_put_request_unauthenticated() { test_case.task.time_precision(), )) .unwrap(); - let collection_job_id: CollectionJobId = random(); let req = CollectionJobReq::new( Query::new_time_interval(batch_interval), dummy::AggregationParam::default().get_encoded().unwrap(), @@ -300,23 +291,19 @@ async fn collection_job_put_request_unauthenticated() { // Incorrect authentication token. let response = test_case - .put_collection_job_with_auth_token(&collection_job_id, &req, Some(&random())) + .post_collection_job_with_auth_token(&req, Some(&random())) .await; assert_eq!(response.status(), StatusCode::FORBIDDEN); // Aggregator authentication token. let response = test_case - .put_collection_job_with_auth_token( - &collection_job_id, - &req, - Some(test_case.task.aggregator_auth_token()), - ) + .post_collection_job_with_auth_token(&req, Some(test_case.task.aggregator_auth_token())) .await; assert_eq!(response.status(), StatusCode::FORBIDDEN); // Missing authentication token. let response = test_case - .put_collection_job_with_auth_token(&collection_job_id, &req, None) + .post_collection_job_with_auth_token(&req, None) .await; assert_eq!(response.status(), StatusCode::FORBIDDEN); } @@ -337,17 +324,15 @@ async fn collection_job_get_request_unauthenticated_collection_jobs() { )) .unwrap(); - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval(batch_interval), dummy::AggregationParam::default().get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &request) - .await; + let response = test_case.post_collection_job(&request).await; - assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id) = ids_from_collection_job_response(&response); // Incorrect authentication token. let response = test_case @@ -388,15 +373,15 @@ async fn collection_job_success_time_interval() { let leader_aggregate_share = dummy::AggregateShare(0); let helper_aggregate_share = dummy::AggregateShare(1); - let collection_job_id: CollectionJobId = random(); let request = CollectionJobReq::new( Query::new_time_interval(batch_interval), aggregation_param.get_encoded().unwrap(), ); - let mut response = test_case - .put_collection_job(&collection_job_id, &request) - .await; + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (task_id, collection_job_id) = ids_from_collection_job_response(&response); + assert_eq!(&task_id, test_case.task.id()); let got_collection_job = test_case .datastore @@ -426,11 +411,11 @@ async fn collection_job_success_time_interval() { ); assert_eq!(want_collection_job, got_collection_job); - assert_eq!(response.status(), StatusCode::CREATED); assert!(take_response_body(&mut response).await.is_empty()); let mut response = test_case.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); assert!(take_response_body(&mut response).await.is_empty()); // Update the collection job with the aggregate shares and some aggregation jobs. collection @@ -501,7 +486,7 @@ async fn collection_job_success_time_interval() { leader_encrypted_aggregate_share, helper_encrypted_aggregate_share, ) = assert_matches!( - collect_resp, + collect_resp.clone(), CollectionJobResp{ report_count, interval, @@ -555,6 +540,18 @@ async fn collection_job_success_time_interval() { helper_aggregate_share, dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() ); + + // Re-send the very first collection job creation request. We should get the representation of + // the completed collection job without needing a further GET. + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, new_collection_job_id) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id, new_collection_job_id); + assert_eq!( + response.headers().get("content-type").unwrap(), + CollectionJobResp::MEDIA_TYPE, + ); + decode_response_body::(&mut response).await; } #[tokio::test] @@ -576,7 +573,7 @@ async fn collection_job_get_request_no_such_collection_job() { } #[tokio::test] -async fn collection_job_put_request_batch_queried_multiple_times() { +async fn collection_job_post_request_batch_queried_multiple_times() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -591,9 +588,8 @@ async fn collection_job_put_request_batch_queried_multiple_times() { dummy::AggregationParam(0).get_encoded().unwrap(), ); - let response = test_case.put_collection_job(&random(), &request).await; - - assert_eq!(response.status(), StatusCode::CREATED); + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); // This request will not be allowed due to the query count already being consumed. let invalid_request = CollectionJobReq::new( @@ -601,24 +597,21 @@ async fn collection_job_put_request_batch_queried_multiple_times() { dummy::AggregationParam(1).get_encoded().unwrap(), ); - let mut response = test_case - .put_collection_job(&random(), &invalid_request) - .await; + let mut response = test_case.post_collection_job(&invalid_request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( take_problem_details(&mut response).await, json!({ "status": StatusCode::BAD_REQUEST.as_u16(), - "type": "urn:ietf:params:ppm:dap:error:invalidMessage", - "title": "The message type for a response was incorrect or the payload was malformed.", - "detail": "batch has already been collected with another aggregation parameter", + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", "taskid": format!("{}", test_case.task.id()), }) ); } #[tokio::test] -async fn collection_job_put_request_batch_overlap() { +async fn collection_job_post_request_batch_overlap() { let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( @@ -639,9 +632,9 @@ async fn collection_job_put_request_batch_overlap() { dummy::AggregationParam(0).get_encoded().unwrap(), ); - let response = test_case.put_collection_job(&random(), &request).await; + let response = test_case.post_collection_job(&request).await; - assert_eq!(response.status(), StatusCode::CREATED); + assert_eq!(response.status(), StatusCode::OK); // This request will not be allowed due to overlapping with the previous request. let invalid_request = CollectionJobReq::new( @@ -649,9 +642,7 @@ async fn collection_job_put_request_batch_overlap() { dummy::AggregationParam(1).get_encoded().unwrap(), ); - let mut response = test_case - .put_collection_job(&random(), &invalid_request) - .await; + let mut response = test_case.post_collection_job(&invalid_request).await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); assert_eq!( take_problem_details(&mut response).await, @@ -674,8 +665,6 @@ async fn delete_collection_job() { )) .await; - let collection_job_id: CollectionJobId = random(); - // Try to delete a collection job that doesn't exist let response = test_case .router @@ -683,13 +672,7 @@ async fn delete_collection_job() { .oneshot( Request::builder() .method("DELETE") - .uri( - test_case - .task - .collection_job_uri(&collection_job_id) - .unwrap() - .path(), - ) + .uri(test_case.task.collection_job_uri(&random()).unwrap().path()) .with_authentication_token(test_case.task.collector_auth_token()) .body(Body::empty()) .unwrap(), @@ -704,11 +687,9 @@ async fn delete_collection_job() { dummy::AggregationParam::default().get_encoded().unwrap(), ); - let response = test_case - .put_collection_job(&collection_job_id, &request) - .await; - - assert_eq!(response.status(), StatusCode::CREATED); + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (_, collection_job_id) = ids_from_collection_job_response(&response); // Cancel the job let response = test_case diff --git a/aggregator/src/aggregator/http_handlers/tests/helper_e2e.rs b/aggregator/src/aggregator/http_handlers/tests/helper_e2e.rs index 117154c8a..a4cbf25b5 100644 --- a/aggregator/src/aggregator/http_handlers/tests/helper_e2e.rs +++ b/aggregator/src/aggregator/http_handlers/tests/helper_e2e.rs @@ -87,9 +87,7 @@ async fn helper_aggregation_report_share_replay() { .updated_with(other_report_1.report_share().metadata().id()); let agg_share_req_1 = AggregateShareReq::::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), agg_param.get_encoded().unwrap(), ), BatchSelector::new(batch_id_1), @@ -99,9 +97,7 @@ async fn helper_aggregation_report_share_replay() { let checksum_2 = ReportIdChecksum::for_report_id(other_report_2.report_share().metadata().id()); let agg_share_req_2 = AggregateShareReq::::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), agg_param.get_encoded().unwrap(), ), BatchSelector::new(batch_id_2), diff --git a/aggregator/src/aggregator/problem_details.rs b/aggregator/src/aggregator/problem_details.rs index 9674da4cd..44f86e1d4 100644 --- a/aggregator/src/aggregator/problem_details.rs +++ b/aggregator/src/aggregator/problem_details.rs @@ -301,18 +301,7 @@ mod tests { ) .expect_task_id(), TestCase::new( - Box::new(|| { - Error::BatchOverlap( - random(), - Interval::new( - RealClock::default() - .now() - .to_time(&TimePrecision::from_seconds(1)), - Duration::from_seconds(3600, &TimePrecision::from_seconds(1)), - ) - .unwrap(), - ) - }), + Box::new(|| Error::BatchOverlap(random())), Some(DapProblemType::BatchOverlap), ) .expect_task_id(), diff --git a/aggregator/src/aggregator/taskprov_tests.rs b/aggregator/src/aggregator/taskprov_tests.rs index 01ce83059..cfbf83dc4 100644 --- a/aggregator/src/aggregator/taskprov_tests.rs +++ b/aggregator/src/aggregator/taskprov_tests.rs @@ -1131,10 +1131,7 @@ async fn taskprov_aggregate_continue() { 0, ReportIdChecksum::default(), CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default( - ), - ), + Query::new_leader_selected(random()), aggregation_param.get_encoded().unwrap(), ), ), @@ -1274,9 +1271,7 @@ async fn taskprov_aggregate_share() { let request = AggregateShareReq::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param.get_encoded().unwrap(), ), BatchSelector::new_leader_selected(batch_id), @@ -1491,9 +1486,7 @@ async fn end_to_end() { let checksum = ReportIdChecksum::for_report_id(report_share.metadata().id()); let aggregate_share_request = AggregateShareReq::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param.get_encoded().unwrap(), ), BatchSelector::new_leader_selected(batch_id), @@ -1637,9 +1630,7 @@ async fn end_to_end_sumvec_hmac() { let checksum = ReportIdChecksum::for_report_id(report_share.metadata().id()); let aggregate_share_request = AggregateShareReq::new( CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param.get_encoded().unwrap(), ), BatchSelector::new_leader_selected(batch_id), diff --git a/aggregator_core/src/batch_mode.rs b/aggregator_core/src/batch_mode.rs index e2d54c842..a5fb08316 100644 --- a/aggregator_core/src/batch_mode.rs +++ b/aggregator_core/src/batch_mode.rs @@ -146,7 +146,7 @@ pub trait CollectableBatchMode: AccumulableBatchMode { type Iter: Iterator + Send + Sync + Clone; /// Retrieves the batch identifier for a given query. - async fn collection_identifier_for_query( + async fn batch_identifier_for_query( tx: &Transaction<'_>, task: &AggregatorTask, query: &Query, @@ -252,7 +252,7 @@ pub trait CollectableBatchMode: AccumulableBatchMode { impl CollectableBatchMode for TimeInterval { type Iter = TimeIntervalBatchIdentifierIter; - async fn collection_identifier_for_query( + async fn batch_identifier_for_query( _: &Transaction<'_>, _: &AggregatorTask, query: &Query, @@ -328,7 +328,7 @@ impl Iterator for TimeIntervalBatchIdentifierIter { impl CollectableBatchMode for LeaderSelected { type Iter = iter::Once; - async fn collection_identifier_for_query( + async fn batch_identifier_for_query( tx: &Transaction<'_>, task: &AggregatorTask, _: &Query, @@ -341,7 +341,7 @@ impl CollectableBatchMode for LeaderSelected { _batch_id: &Self::BatchIdentifier, _query: &Query, ) -> bool { - // A leader-selected query has an empty body, so any batch ID is consistent with it. + // Leader-selected queries are consistent with any batch ID. true } @@ -367,7 +367,7 @@ impl CollectableBatchMode for LeaderSelected { mod tests { use janus_messages::{ BatchId, CollectionJobReq, Duration, Interval, Query, Time, TimePrecision, - batch_mode::{LeaderSelected, TimeInterval}, + batch_mode::{LeaderSelected, LeaderSelectedQueryConfig, TimeInterval}, }; use prio::{ codec::Encode, @@ -433,18 +433,14 @@ mod tests { let batch_id = BatchId::from([5u8; 32]); let collector_request = CollectionJobReq::::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), encoded_aggregation_parameter, ); let leader_job = CollectionJob::<0, LeaderSelected, Vdaf>::new( random(), random(), random(), - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), aggregation_parameter, batch_id, CollectionJobState::Start, diff --git a/aggregator_core/src/datastore.rs b/aggregator_core/src/datastore.rs index dbfa13d44..424cd81b1 100644 --- a/aggregator_core/src/datastore.rs +++ b/aggregator_core/src/datastore.rs @@ -22,6 +22,7 @@ use futures::future::try_join_all; use janus_core::{ auth_tokens::AuthenticationToken, hpke::{self, HpkePrivateKey}, + messages::DeriveId, time::{Clock, TimeExt}, vdaf::VdafInstance, }; @@ -3024,8 +3025,7 @@ WHERE collection_jobs.task_id = $1 &[ /* task_id */ &task_info.pkey, /* collection_job_id */ &collection_job_id.as_ref(), - /* threshold */ - &task_info.report_expiry_threshold(self.clock.now())?, + /* threshold */ &task_info.report_expiry_threshold(self.clock.now())?, ], ) .await? @@ -3042,6 +3042,60 @@ WHERE collection_jobs.task_id = $1 .transpose() } + /// Returns the collection job matching the query, if any exists. + pub async fn get_collection_job_by_query< + const SEED_SIZE: usize, + B: BatchMode, + A: AsyncAggregator, + >( + &self, + vdaf: &A, + task_id: &TaskId, + query: &Query, + ) -> Result>, Error> { + let task_info = match self.task_info_for(task_id).await? { + Some(task_info) => task_info, + None => return Ok(None), + }; + + let stmt = self + .prepare_cached( + "-- get_collection_job_by_query() +SELECT + collection_job_id, aggregation_param, batch_identifier, state, report_count, + client_timestamp_interval, helper_aggregate_share, leader_aggregate_share, + aggregate_share_id, query +FROM collection_jobs +WHERE collection_jobs.task_id = $1 +AND collection_jobs.query = $2 +AND COALESCE( + LOWER(collection_jobs.batch_interval), + (SELECT MAX(UPPER(client_timestamp_interval)) + FROM batch_aggregations + WHERE batch_aggregations.task_id = collection_jobs.task_id + AND batch_aggregations.batch_identifier = collection_jobs.batch_identifier + AND batch_aggregations.aggregation_param = collection_jobs.aggregation_param), + 0) >= $3", + ) + .await?; + self.query_opt( + &stmt, + &[ + /* task_id */ &task_info.pkey, + /* query */ &query.get_encoded()?, + /* threshold */ &task_info.report_expiry_threshold(self.clock.now())?, + ], + ) + .await? + .map(|row| { + let batch_identifier = B::BatchIdentifier::get_decoded(row.get("batch_identifier"))?; + let collection_job_id = + row.get_bytea_and_convert::("collection_job_id")?; + Self::collection_job_from_row(vdaf, *task_id, batch_identifier, collection_job_id, &row) + }) + .transpose() + } + /// Returns a collection job in state FINISHED with the given parameters, or `None` if no such /// collection job exists. pub async fn get_finished_collection_job< @@ -3399,6 +3453,36 @@ WHERE task_id = $1 )) } + /// Stores a new collection job constructed from a [`CollectionJobReq`]. + /// + /// The collection job ID is deterministically derived from the request and returned. + pub async fn put_collection_job_by_request< + const SEED_SIZE: usize, + B: CollectableBatchMode, + A: AsyncAggregator, + >( + &self, + task_id: &TaskId, + req: &CollectionJobReq, + collection_identifier: B::BatchIdentifier, + ) -> Result { + let collection_job_id = req.derive_id()?; + + self.put_collection_job(&CollectionJob::::new( + *task_id, + collection_job_id, + random(), + req.query().clone(), + A::AggregationParam::get_decoded(req.aggregation_parameter()) + .map_err(|_| Error::InvalidParameter("invalid aggregation parameter"))?, + collection_identifier, + CollectionJobState::Start, + )) + .await?; + + Ok(collection_job_id) + } + /// Stores a new collection job. #[tracing::instrument(skip(self), err(level = Level::DEBUG))] pub async fn put_collection_job< @@ -3409,10 +3493,10 @@ WHERE task_id = $1 &self, collection_job: &CollectionJob, ) -> Result<(), Error> { - let task_info = match self.task_info_for(collection_job.task_id()).await? { - Some(task_info) => task_info, - None => return Err(Error::MutationTargetNotFound), - }; + let task_info = self + .task_info_for(collection_job.task_id()) + .await? + .ok_or_else(|| Error::MutationTargetNotFound)?; let now = self.clock.now(); let batch_interval = B::to_batch_interval(collection_job.batch_identifier()) @@ -6086,6 +6170,9 @@ pub enum Error { /// The transaction was aborted because it retried too many times. #[error("too many retries")] TooManyRetries { source: Option> }, + /// An identifier could not be derived from a request. + #[error("failed to derive identifier: {0}")] + IdentifierDerivation(#[from] janus_core::messages::DeriveIdError), } impl From for Error { diff --git a/aggregator_core/src/datastore/tests.rs b/aggregator_core/src/datastore/tests.rs index c1573d210..104fef2d1 100644 --- a/aggregator_core/src/datastore/tests.rs +++ b/aggregator_core/src/datastore/tests.rs @@ -4036,9 +4036,7 @@ async fn get_collection_jobs_by_batch_id(ephemeral_datastore: EphemeralDatastore *task.id(), random(), random(), - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param, random(), CollectionJobState::Start, @@ -4066,9 +4064,7 @@ async fn get_collection_jobs_by_batch_id(ephemeral_datastore: EphemeralDatastore *task.id(), random(), random(), - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), aggregation_param, random(), CollectionJobState::Start, @@ -4122,6 +4118,167 @@ async fn get_collection_jobs_by_batch_id(ephemeral_datastore: EphemeralDatastore .unwrap(); } +async fn get_collection_jobs_by_query_test( + ephemeral_datastore: EphemeralDatastore, + batch_mode: task::BatchMode, + queries: [Query; 3], + batch_identifiers: [B::BatchIdentifier; 3], +) { + install_test_trace_subscriber(); + + let clock = MockClock::new(START_TIMESTAMP); + let ds = ephemeral_datastore.datastore(clock.clone()).await; + let aggregation_param = dummy::AggregationParam(13); + + ds.run_tx("test-put-collection-job", |tx| { + let queries = queries.clone(); + let batch_identifiers = batch_identifiers.clone(); + Box::pin(async move { + let task = TaskBuilder::new( + batch_mode, + AggregationMode::Synchronous, + VdafInstance::Fake { rounds: 1 }, + ) + .with_report_expiry_age(Some(REPORT_EXPIRY_AGE_DURATION)) + .with_time_precision(TIME_PRECISION) + .build() + .leader_view() + .unwrap(); + tx.put_aggregator_task(&task).await.unwrap(); + + let first_collection_job = CollectionJob::<0, B, dummy::Vdaf>::new( + *task.id(), + random(), + random(), + queries[0].clone(), + aggregation_param, + batch_identifiers[0].clone(), + CollectionJobState::Start, + ); + tx.put_collection_job(&first_collection_job).await.unwrap(); + // We must insert a batch aggregation so that the query can work out a client timestamp + // interval to check report expiry against. It doesn't matter what its + // BatchAggregationState is. + tx.put_batch_aggregation(&BatchAggregation::<0, B, dummy::Vdaf>::new( + *task.id(), + first_collection_job.batch_identifier().clone(), + aggregation_param, + 0, + Interval::new( + START_TIME, + Duration::from_seconds(TIME_PRECISION_SECONDS, &TIME_PRECISION), + ) + .unwrap(), + BatchAggregationState::Scrubbed, + )) + .await + .unwrap(); + + let second_collection_job = CollectionJob::<0, B, dummy::Vdaf>::new( + *task.id(), + random(), + random(), + queries[1].clone(), + aggregation_param, + batch_identifiers[1].clone(), + CollectionJobState::Start, + ); + tx.put_collection_job(&second_collection_job).await.unwrap(); + tx.put_batch_aggregation(&BatchAggregation::<0, B, dummy::Vdaf>::new( + *task.id(), + second_collection_job.batch_identifier().clone(), + aggregation_param, + 0, + Interval::new( + START_TIME, + Duration::from_seconds(TIME_PRECISION_SECONDS, &TIME_PRECISION), + ) + .unwrap(), + BatchAggregationState::Scrubbed, + )) + .await + .unwrap(); + + let check_jobs = + async |want_first_job: bool, want_second_job: bool, query: &Query<_>| { + let job = tx + .get_collection_job_by_query(&dummy::Vdaf::default(), task.id(), query) + .await + .unwrap(); + + if want_first_job { + assert_eq!(job.unwrap().id(), first_collection_job.id()); + } else if want_second_job { + assert_eq!(job.unwrap().id(), second_collection_job.id()); + } else { + assert!(job.is_none()); + } + }; + + check_jobs(true, false, &queries[0]).await; + check_jobs(false, true, &queries[1]).await; + check_jobs(false, false, &queries[2]).await; + + Ok(()) + }) + }) + .await + .unwrap(); +} + +#[rstest_reuse::apply(schema_versions_template)] +#[tokio::test] +async fn get_collection_jobs_by_query_time_interval(ephemeral_datastore: EphemeralDatastore) { + let first_batch_interval = + Interval::minimal(START_TIME.add_duration(&Duration::ONE).unwrap()).unwrap(); + let second_batch_interval = Interval::minimal( + START_TIME + .add_duration(&Duration::from_time_precision_units(2)) + .unwrap(), + ) + .unwrap(); + let third_batch_interval = Interval::minimal( + START_TIME + .add_duration(&Duration::from_time_precision_units(3)) + .unwrap(), + ) + .unwrap(); + + get_collection_jobs_by_query_test( + ephemeral_datastore, + task::BatchMode::TimeInterval, + [ + Query::new_time_interval(first_batch_interval), + Query::new_time_interval(second_batch_interval), + Query::new_time_interval(third_batch_interval), + ], + [ + first_batch_interval, + second_batch_interval, + third_batch_interval, + ], + ) + .await; +} + +#[rstest_reuse::apply(schema_versions_template)] +#[tokio::test] +async fn get_collection_jobs_by_query_leader_selected(ephemeral_datastore: EphemeralDatastore) { + get_collection_jobs_by_query_test( + ephemeral_datastore, + task::BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + [ + Query::new_leader_selected(random()), + Query::new_leader_selected(random()), + Query::new_leader_selected(random()), + ], + [random(), random(), random()], + ) + .await +} + #[rstest_reuse::apply(schema_versions_template)] #[tokio::test] async fn update_collection_jobs(ephemeral_datastore: EphemeralDatastore) { @@ -4313,9 +4470,7 @@ impl TestBatchModeExt for TimeInterval { #[async_trait] impl TestBatchModeExt for LeaderSelected { fn query_for_batch_identifier(_: &Self::BatchIdentifier) -> Query { - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ) + Query::new_leader_selected(random()) } fn batch_identifier_for_client_timestamps(_: &[Time]) -> Self::BatchIdentifier { @@ -6449,9 +6604,7 @@ async fn roundtrip_aggregate_share_job_leader_selected(ephemeral_datastore: Ephe 10, ReportIdChecksum::get_decoded(&[1; 32]).unwrap(), CollectionJobReq::new( - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), dummy::AggregationParam(11).get_encoded().unwrap(), ) .with_extensions(Vec::from([CollectionJobExtension::new( diff --git a/aggregator_core/src/task.rs b/aggregator_core/src/task.rs index d90874d62..3d0322071 100644 --- a/aggregator_core/src/task.rs +++ b/aggregator_core/src/task.rs @@ -1239,22 +1239,29 @@ pub mod test_util { aggregate_share_id: &AggregateShareId, ) -> Result { Ok(self.helper_aggregator_endpoint().join(&format!( - "{}/aggregate_shares/{}", - self.tasks_path(), - aggregate_share_id + "{}/aggregate_shares/{aggregate_share_id}", + self.tasks_path() ))?) } + /// Returns the URI at which the leader resource for collection job creation can be + /// accessed. + pub fn collection_job_creation_uri(&self) -> Result { + Ok(self + .leader_aggregator_endpoint() + .join(&format!("{}/collection_jobs", self.tasks_path()))?) + } + /// Returns the URI at which the leader resource for the specified collection job ID can be /// accessed. pub fn collection_job_uri( &self, collection_job_id: &CollectionJobId, ) -> Result { - Ok(self.leader_aggregator_endpoint().join(&format!( - "{}/collection_jobs/{collection_job_id}", - self.tasks_path() - ))?) + Ok(self + .collection_job_creation_uri()? + .ensure_trailing_slash() + .join(&collection_job_id.to_string())?) } /// Render the leader aggregator's view of this task. diff --git a/core/src/http.rs b/core/src/http.rs index 3acd72238..685e17e7f 100644 --- a/core/src/http.rs +++ b/core/src/http.rs @@ -4,9 +4,9 @@ use std::{ }; use anyhow::{Context, anyhow}; -use http::{StatusCode, header::RETRY_AFTER}; +use http::{HeaderName, StatusCode, header::RETRY_AFTER}; use http_api_problem::{HttpApiProblem, PROBLEM_JSON_MEDIA_TYPE}; -use janus_messages::{MediaType, problem_type::DapProblemType}; +use janus_messages::{CollectionJobId, MediaType, TaskId, problem_type::DapProblemType}; use mime::Mime; use reqwest::{Response, header::CONTENT_TYPE}; use retry_after::RetryAfter; @@ -220,6 +220,14 @@ impl ReqwestAuthenticationToken for reqwest::RequestBuilder { } } +/// Construct an HTTP Location header for the collection job. +pub fn collection_job_location(task_id: &TaskId, job_id: &CollectionJobId) -> (HeaderName, String) { + ( + http::header::LOCATION, + format!("/tasks/{task_id}/collection_jobs/{job_id}"), + ) +} + #[cfg(test)] mod tests { use assert_matches::assert_matches; diff --git a/core/src/lib.rs b/core/src/lib.rs index 66425cfce..f8d705b1c 100644 --- a/core/src/lib.rs +++ b/core/src/lib.rs @@ -10,6 +10,7 @@ pub mod cli; pub mod dp; pub mod hpke; pub mod http; +pub mod messages; pub mod report_id; pub mod retries; pub mod task_config; diff --git a/core/src/messages.rs b/core/src/messages.rs new file mode 100644 index 000000000..6948cb50c --- /dev/null +++ b/core/src/messages.rs @@ -0,0 +1,45 @@ +//! Extensions and utilities for items from `janus_messages`. + +use aws_lc_rs::digest::{SHA256, digest}; +use janus_messages::{CollectionJobId, CollectionJobReq, batch_mode::BatchMode}; +use prio::codec::Encode; + +/// Errors from deriving IDs. +#[derive(Debug, thiserror::Error)] +pub enum DeriveIdError { + /// Error handling a message. + #[error("{0}")] + Messages(#[from] janus_messages::Error), + /// Error encoding or decoding a message. + #[error("{0}")] + Codec(#[from] prio::codec::CodecError), +} + +/// Derive a unique identifier for the receiver. +pub trait DeriveId { + type Identifier; + + fn derive_id(&self) -> Result; +} + +impl DeriveId for CollectionJobReq { + type Identifier = CollectionJobId; + + fn derive_id(&self) -> Result { + derive_id(self, CollectionJobId::LEN) + } +} + +fn derive_id( + req: &Request, + identifier_len: usize, +) -> Result +where + Request: Encode, + Identifier: for<'a> TryFrom<&'a [u8], Error = janus_messages::Error>, +{ + let request_digest = digest(&SHA256, &req.get_encoded()?); + Ok(Identifier::try_from( + &request_digest.as_ref()[..identifier_len], + )?) +} diff --git a/db/00000000000001_initial_schema.up.sql b/db/00000000000001_initial_schema.up.sql index ac2e44ed5..9fc91009e 100644 --- a/db/00000000000001_initial_schema.up.sql +++ b/db/00000000000001_initial_schema.up.sql @@ -447,6 +447,7 @@ CREATE TABLE collection_jobs( CONSTRAINT collection_jobs_unique_id UNIQUE(task_id, collection_job_id), CONSTRAINT fk_task_id FOREIGN KEY(task_id) REFERENCES tasks(id) ON DELETE CASCADE ); +CREATE INDEX collection_jobs_task_id_query ON collection_jobs(task_id, query); CREATE INDEX collection_jobs_task_id_batch_id ON collection_jobs(task_id, batch_identifier); -- TODO(#224): verify that this index is optimal for purposes of acquiring collection jobs. CREATE INDEX collection_jobs_state_and_lease_expiry ON collection_jobs(state, lease_expiry) WHERE state = 'START'; From d0b80eb7d9d421b63f07f774dc5a913f84be27d2 Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Tue, 13 Oct 2026 11:21:47 -0400 Subject: [PATCH 2/6] neuter integration tests temporarily These tests will be re-enabled in #4899. --- integration_tests/tests/integration/janus.rs | 27 +++++++++++++++++++ .../tests/integration/metrics.rs | 5 ++++ .../integration/simulation/reproduction.rs | 3 +++ 3 files changed, 35 insertions(+) diff --git a/integration_tests/tests/integration/janus.rs b/integration_tests/tests/integration/janus.rs index b9a6629f0..8be0d18b9 100644 --- a/integration_tests/tests/integration/janus.rs +++ b/integration_tests/tests/integration/janus.rs @@ -208,6 +208,7 @@ async fn run_in_process_test( /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_sync_count() { run_container_test( "janus_janus_sync_count", @@ -222,6 +223,7 @@ async fn janus_janus_sync_count() { /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_async_count() { // TODO(#3436): allow callers to specify asynchronous aggregation mode (requires updated // interop test design) @@ -237,6 +239,7 @@ async fn janus_janus_async_count() { /// This test exercises Prio3Count with Janus as both the leader and the helper, using synchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_sync_count() { run_in_process_test( "janus_in_process_sync_count", @@ -250,6 +253,7 @@ async fn janus_in_process_sync_count() { /// This test exercises Prio3Count with Janus as both the leader and the helper, using asynchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_async_count() { run_in_process_test( "janus_in_process_async_count", @@ -264,6 +268,7 @@ async fn janus_in_process_async_count() { /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_sync_sum_16() { run_container_test( "janus_janus_sync_sum_16", @@ -280,6 +285,7 @@ async fn janus_janus_sync_sum_16() { /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_async_sum_16() { run_container_test( "janus_janus_async_sum_16", @@ -295,6 +301,7 @@ async fn janus_janus_async_sum_16() { /// This test exercises Prio3Sum with Janus as both the leader and the helper, using synchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_sync_sum_16() { run_in_process_test( "janus_in_process_sync_sum_16", @@ -310,6 +317,7 @@ async fn janus_in_process_sync_sum_16() { /// This test exercises Prio3Sum with Janus as both the leader and the helper, using asynchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_async_sum_16() { run_in_process_test( "janus_in_process_async_sum_16", @@ -326,6 +334,7 @@ async fn janus_in_process_async_sum_16() { /// synchronous aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_sync_histogram_4_buckets() { run_container_test( "janus_janus_sync_histogram_4_buckets", @@ -344,6 +353,7 @@ async fn janus_janus_sync_histogram_4_buckets() { /// asynchronous aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_async_histogram_4_buckets() { run_container_test( "janus_janus_async_histogram_4_buckets", @@ -361,6 +371,7 @@ async fn janus_janus_async_histogram_4_buckets() { /// This test exercises Prio3Histogram with Janus as both the leader and the helper, using /// synchronous aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_sync_histogram_4_buckets() { run_in_process_test( "janus_in_process_sync_histogram_4_buckets", @@ -378,6 +389,7 @@ async fn janus_in_process_sync_histogram_4_buckets() { /// This test exercises Prio3Histogram with Janus as both the leader and the helper, using /// asynchronous aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_async_histogram_4_buckets() { run_in_process_test( "janus_in_process_async_histogram_4_buckets", @@ -396,6 +408,7 @@ async fn janus_in_process_async_histogram_4_buckets() { /// using synchronous aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_leader_selected_sync() { run_container_test( "janus_janus_leader_selected_sync", @@ -412,6 +425,7 @@ async fn janus_janus_leader_selected_sync() { /// using asynchronous aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_leader_selected_async() { run_container_test( "janus_janus_leader_selected_async", @@ -427,6 +441,7 @@ async fn janus_janus_leader_selected_async() { /// This test exercises the leader-selected batch mode with Janus as both the leader and the helper, /// using synchronous aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_leader_selected_sync() { run_in_process_test( "janus_in_process_leader_selected_sync", @@ -442,6 +457,7 @@ async fn janus_in_process_leader_selected_sync() { /// This test exercises the leader-selected batch mode with Janus as both the leader and the helper, /// using asynchronous aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_leader_selected_async() { run_in_process_test( "janus_in_process_leader_selected_async", @@ -458,6 +474,7 @@ async fn janus_in_process_leader_selected_async() { /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_sync_sum_vec() { run_container_test( "janus_janus_sync_sum_vec", @@ -477,6 +494,7 @@ async fn janus_janus_sync_sum_vec() { /// aggregation. #[tokio::test(flavor = "multi_thread")] #[cfg(feature = "testcontainer")] +#[ignore] async fn janus_janus_async_sum_vec() { run_container_test( "janus_janus_async_sum_vec", @@ -495,6 +513,7 @@ async fn janus_janus_async_sum_vec() { /// This test exercises Prio3SumVec with Janus as both the leader and the helper, using synchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_sync_sum_vec() { run_in_process_test( "janus_in_process_sync_sum_vec", @@ -513,6 +532,7 @@ async fn janus_in_process_sync_sum_vec() { /// This test exercises Prio3SumVec with Janus as both the leader and the helper, using asynchronous /// aggregation. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_async_sum_vec() { run_in_process_test( "janus_in_process_async_sum_vec", @@ -531,6 +551,7 @@ async fn janus_in_process_async_sum_vec() { /// This test exercises Prio3SumVecField64MultiproofHmacSha256Aes128 with Janus as both the leader /// and the helper. #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_customized_sum_vec() { install_test_trace_subscriber(); initialize_rustls(); @@ -558,6 +579,7 @@ async fn janus_in_process_customized_sum_vec() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_histogram_dp_noise() { static TEST_NAME: &str = "janus_in_process_histogram_dp_noise"; const HISTOGRAM_LENGTH: usize = 100; @@ -627,6 +649,7 @@ async fn janus_in_process_histogram_dp_noise() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_sumvec_dp_noise() { static TEST_NAME: &str = "janus_in_process_sumvec_dp_noise"; const VECTOR_LENGTH: usize = 50; @@ -700,6 +723,7 @@ async fn janus_in_process_sumvec_dp_noise() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_fake_vdaf_4_round_sync() { install_test_trace_subscriber(); initialize_rustls(); @@ -721,6 +745,7 @@ async fn janus_in_process_fake_vdaf_4_round_sync() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_fake_vdaf_4_round_async() { install_test_trace_subscriber(); initialize_rustls(); @@ -742,6 +767,7 @@ async fn janus_in_process_fake_vdaf_4_round_async() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_fake_vdaf_5_round_sync() { install_test_trace_subscriber(); initialize_rustls(); @@ -763,6 +789,7 @@ async fn janus_in_process_fake_vdaf_5_round_sync() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_in_process_fake_vdaf_5_round_async() { install_test_trace_subscriber(); initialize_rustls(); diff --git a/integration_tests/tests/integration/metrics.rs b/integration_tests/tests/integration/metrics.rs index 920b8673f..3b88db383 100644 --- a/integration_tests/tests/integration/metrics.rs +++ b/integration_tests/tests/integration/metrics.rs @@ -11,11 +11,13 @@ use crate::{ }; #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_metrics_count() { run_metrics_test(VdafInstance::Prio3Count, "Prio3Count", 1).await; } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_metrics_sum() { run_metrics_test( VdafInstance::Prio3Sum { @@ -28,6 +30,7 @@ async fn janus_metrics_sum() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_metrics_sumvec() { run_metrics_test( VdafInstance::Prio3SumVec { @@ -43,6 +46,7 @@ async fn janus_metrics_sumvec() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_metrics_sumvec_custom() { run_metrics_test( VdafInstance::Prio3SumVecField64MultiproofHmacSha256Aes128 { @@ -59,6 +63,7 @@ async fn janus_metrics_sumvec_custom() { } #[tokio::test(flavor = "multi_thread")] +#[ignore] async fn janus_metrics_histogram() { run_metrics_test( VdafInstance::Prio3Histogram { diff --git a/integration_tests/tests/integration/simulation/reproduction.rs b/integration_tests/tests/integration/simulation/reproduction.rs index e28d5aaac..d27e9203c 100644 --- a/integration_tests/tests/integration/simulation/reproduction.rs +++ b/integration_tests/tests/integration/simulation/reproduction.rs @@ -16,6 +16,7 @@ use crate::simulation::{ }; #[test] +#[ignore] fn successful_collection_time_interval() { install_test_trace_subscriber(); initialize_rustls(); @@ -123,6 +124,7 @@ fn successful_collection_time_interval() { } #[test] +#[ignore] fn successful_collection_leader_selected() { install_test_trace_subscriber(); initialize_rustls(); @@ -218,6 +220,7 @@ fn successful_collection_leader_selected() { } #[test] +#[ignore] fn successful_collection_asynchronous() { install_test_trace_subscriber(); initialize_rustls(); From f3951dd7ddb007c1680796623bfe4fb02a59e70a Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Fri, 9 Oct 2026 16:33:38 -0400 Subject: [PATCH 3/6] janus_collector: server selected collection job id Adds support to `janus_collector` and the `collect` bin in `janus_tools` for server selected collection job IDs and leader-selected batch mode idempotency keys, per DAP changes in [1]. [1]: https://github.com/ietf-wg-ppm/draft-ietf-ppm-dap/pull/797 --- collector/Cargo.toml | 2 +- collector/src/lib.rs | 259 ++++++++++-------- .../tests/integration/simulation/run.rs | 2 - messages/src/batch_mode.rs | 6 + tools/src/bin/collect.rs | 158 ++++++----- tools/tests/cmd/collect.trycmd | 6 + 6 files changed, 254 insertions(+), 179 deletions(-) diff --git a/collector/Cargo.toml b/collector/Cargo.toml index 0cdd65ba5..b75d8b7b1 100644 --- a/collector/Cargo.toml +++ b/collector/Cargo.toml @@ -26,7 +26,6 @@ janus_core.workspace = true janus_messages.workspace = true mime.workspace = true prio.workspace = true -rand = { workspace = true } reqwest = { workspace = true, features = ["json"] } retry-after = { workspace = true } serde = { workspace = true } @@ -42,3 +41,4 @@ base64.workspace = true janus_collector = { workspace = true, features = ["test-util"] } janus_core = { workspace = true, features = ["test-util"] } mockito = { workspace = true } +rand = { workspace = true } diff --git a/collector/src/lib.rs b/collector/src/lib.rs index 733588381..fb4c32448 100644 --- a/collector/src/lib.rs +++ b/collector/src/lib.rs @@ -12,7 +12,10 @@ //! use std::{fs::File, str::FromStr}; //! //! use janus_collector::{Collector, ConfiguredVdaf, PrivateCollectorCredential}; -//! use janus_messages::{BatchConfig, Duration, Interval, Query, TaskId, Time, TimePrecision, Url}; +//! use janus_messages::{ +//! BatchConfig, Duration, Interval, Query, TaskId, Time, TimePrecision, Url, +//! batch_mode::LeaderSelectedQueryConfig, +//! }; //! //! # async fn run() { //! # const TIME_PRECISION: u64 = 3600; @@ -62,8 +65,12 @@ //! .await //! .unwrap(); //! -//! // Or if this is a leader-selected task, make a leader-selected query. -//! let query = Query::new_leader_selected(todo!("idempotency key")); +//! // Or if this is a leader-selected task, make a leader-selected query. You must choose an +//! // idempotency key to use as the query configuration. The idempotency key allows collection job +//! // creation to be safely retried, so you may wish to persist their values before sending +//! // requests. +//! let idempotency_key = LeaderSelectedQueryConfig::new(vec![0x00, 0x01, 0x02, 0x03]).unwrap(); +//! let query = Query::new_leader_selected(idempotency_key); //! let aggregation_result = collector.collection(query, &()).collect().await.unwrap(); //! # } //! ``` @@ -81,17 +88,19 @@ pub use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder}; use chrono::{DateTime, Duration, Utc}; pub use credential::PrivateCollectorCredential; use educe::Educe; -pub use janus_core::{auth_tokens::AuthenticationToken, hpke::HpkeKeypair, vdaf::ConfiguredVdaf}; use janus_core::{ + UrlExt, hpke::{self, HpkeApplicationInfo}, http::{HttpErrorResponse, ReqwestAuthenticationToken, check_content_type}, retries::{ - ExponentialWithTotalDelayBuilder, http_request_exponential_backoff, retry_http_request, + ExponentialWithTotalDelayBuilder, HeaderMapExt, HeaderMapExtError, + http_request_exponential_backoff, retry_http_request, }, task_config::build_task_configuration, time::TimeExt, url_for_join, }; +pub use janus_core::{auth_tokens::AuthenticationToken, hpke::HpkeKeypair, vdaf::ConfiguredVdaf}; use janus_messages::{ AggregateShareAad, BatchConfig, CollectionJobExtension, CollectionJobId, CollectionJobReq, CollectionJobResp, Interval, MediaType, Query, Role, TaskConfiguration, TaskId, TimePrecision, @@ -102,10 +111,9 @@ use prio::{ codec::{Decode, Encode, ParameterizedDecode}, vdaf, }; -use rand::random; use reqwest::{ StatusCode, - header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderValue, RETRY_AFTER, ToStrError}, + header::{CONTENT_LENGTH, CONTENT_TYPE, RETRY_AFTER}, }; pub use retry_after; use retry_after::{FromHeaderValueError, RetryAfter}; @@ -124,10 +132,8 @@ pub enum Error { Url(#[from] url::ParseError), #[error("missing Location header in See Other response")] MissingLocationHeader, - #[error("invalid bytes in header")] - InvalidHeader(#[from] ToStrError), - #[error("wrong Content-Type header: {0:?}")] - BadContentType(Option), + #[error("invalid header")] + InvalidHeader(#[from] HeaderMapExtError), #[error("invalid Retry-After header value: {0}")] InvalidRetryAfterHeader(#[from] FromHeaderValueError), #[error("codec error: {0}")] @@ -223,23 +229,16 @@ impl CollectionJob { /// A builder for a collection request, returned by [`Collector::collection`]. /// -/// Configure the optional collection job ID ([`Self::with_id`]) and extensions -/// ([`Self::with_extensions`]), then send the request with [`Self::start`] or [`Self::collect`]. +/// Configure the optional extensions ([`Self::with_extensions`]), then send the request with +/// [`Self::start`] or [`Self::collect`]. pub struct CollectionRequestBuilder<'a, V: vdaf::Collector, B: BatchMode> { collector: &'a Collector, query: Query, aggregation_parameter: &'a V::AggregationParam, - collection_job_id: Option, extensions: Vec, } impl<'a, V: vdaf::Collector, B: BatchMode> CollectionRequestBuilder<'a, V, B> { - /// Use a caller-chosen collection job ID instead of a randomly generated one. - pub fn with_id(mut self, collection_job_id: CollectionJobId) -> Self { - self.collection_job_id = Some(collection_job_id); - self - } - /// Set the collection job extensions. pub fn with_extensions(mut self, extensions: Vec) -> Self { self.extensions = extensions; @@ -250,14 +249,8 @@ impl<'a, V: vdaf::Collector, B: BatchMode> CollectionRequestBuilder<'a, V, B> { /// [`CollectionJob`] that must be polled separately using [`Collector::poll_once`] or /// [`Collector::poll_until_complete`]. pub async fn start(self) -> Result, Error> { - let collection_job_id = self.collection_job_id.unwrap_or_else(random); self.collector - .send_collection_request( - collection_job_id, - self.query, - self.aggregation_parameter, - self.extensions, - ) + .send_collection_request(self.query, self.aggregation_parameter, self.extensions) .await } @@ -649,22 +642,27 @@ impl Collector { &self.task_configuration } + /// Construct a URI for collection job creation. + fn collection_job_creation_uri(&self) -> Result { + Ok( + url_for_join(self.task_configuration.leader_aggregator_endpoint())? + .join(&format!("tasks/{}/collection_jobs", self.task_id))?, + ) + } + /// Construct a URI for a collection. fn collection_job_uri(&self, collection_job_id: CollectionJobId) -> Result { // Joined at request time so the endpoint bytes bound into HPKE AADs stay verbatim // (DAP §4.1). - Ok( - url_for_join(self.task_configuration.leader_aggregator_endpoint())?.join(&format!( - "tasks/{}/collection_jobs/{collection_job_id}", - self.task_id - ))?, - ) + Ok(self + .collection_job_creation_uri()? + .ensure_trailing_slash() + .join(&collection_job_id.to_string())?) } /// Construct a collection request to send to the leader aggregator. /// - /// Returns a [`CollectionRequestBuilder`]; configure the optional collection job ID - /// ([`CollectionRequestBuilder::with_id`]) and extensions + /// Returns a [`CollectionRequestBuilder`]; configure the optional extensions /// ([`CollectionRequestBuilder::with_extensions`]), then send the request with /// [`CollectionRequestBuilder::start`] (returns a [`CollectionJob`] to poll) or /// [`CollectionRequestBuilder::collect`] (sends and waits for the result). @@ -677,7 +675,6 @@ impl Collector { collector: self, query, aggregation_parameter, - collection_job_id: None, extensions: Vec::new(), } } @@ -687,7 +684,6 @@ impl Collector { #[tracing::instrument(skip(self, aggregation_parameter, extensions), err)] async fn send_collection_request( &self, - collection_job_id: CollectionJobId, query: Query, aggregation_parameter: &V::AggregationParam, extensions: Vec, @@ -696,12 +692,12 @@ impl Collector { CollectionJobReq::new(query.clone(), aggregation_parameter.get_encoded()?) .with_extensions(extensions) .get_encoded()?; - let collection_job_url = self.collection_job_uri(collection_job_id)?; + let collection_job_url = self.collection_job_creation_uri()?; let response_res = retry_http_request(self.http_request_retry_parameters.build(), || async { self.http_client - .put(collection_job_url.clone()) + .post(collection_job_url.clone()) .header(CONTENT_TYPE, CollectionJobReq::::MEDIA_TYPE) .body(collect_request.clone()) .authentication_token(&self.authentication) @@ -710,14 +706,20 @@ impl Collector { }) .await; - match response_res { + let collection_job_id = match response_res { // Successful response. Ok(response) => { let status = response.status(); - if status != StatusCode::CREATED { + if status != StatusCode::OK { // Incorrect success status code. return Err(Error::Http(Box::new(status.into()))); } + + let (_, collection_job_id) = response + .headers() + .ids_from_collection_job_location()? + .ok_or_else(|| Error::MissingLocationHeader)?; + collection_job_id } // HTTP-level error. @@ -997,7 +999,7 @@ mod tests { use rand::random; use reqwest::{ StatusCode, - header::{AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE}, + header::{AUTHORIZATION, CONTENT_LENGTH, CONTENT_TYPE, LOCATION}, }; use retry_after::RetryAfter; @@ -1029,15 +1031,20 @@ mod tests { .unwrap() } + /// Matches on the relative path for collection job creation. + fn collection_creation_uri_matcher(task_id: &TaskId) -> Matcher { + Matcher::Regex(format!("^/tasks/{task_id}/collection_jobs")) + } + + /// Matches on the relative path for a collection job resource. fn collection_uri_regex_matcher(task_id: &TaskId) -> Matcher { - // Matches on the relative path for a collection job resource. The Base64 URL-safe encoding - // of a collection ID is always 22 characters. + // The Base64 URL-safe encoding of a collection ID is always 22 characters. Matcher::Regex(format!( "^/tasks/{task_id}/collection_jobs/[A-Za-z0-9-_]{{22}}$" )) } - fn build_collect_response_time( + fn build_collect_response_time_interval( transcript: &VdafTranscript, collector: &Collector, aggregation_parameter: &V::AggregationParam, @@ -1075,9 +1082,10 @@ mod tests { } } - fn build_collect_response_fixed( + fn build_collect_response_leader_selected( transcript: &VdafTranscript, collector: &Collector, + idempotency_key: &LeaderSelectedQueryConfig, aggregation_parameter: &V::AggregationParam, ) -> CollectionJobResp where @@ -1088,7 +1096,7 @@ mod tests { collector.task_id, collector.task_configuration().clone(), CollectionJobReq::new( - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(idempotency_key.clone()), aggregation_parameter.get_encoded().unwrap(), ), ); @@ -1280,6 +1288,10 @@ mod tests { .as_str(), endpoint ); + assert_eq!( + collector.collection_job_creation_uri().unwrap().as_str(), + format!("{joined_base}tasks/{}/collection_jobs", collector.task_id) + ); assert_eq!( collector .collection_job_uri(collection_job_id) @@ -1307,7 +1319,12 @@ mod tests { &random(), &true, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1315,11 +1332,11 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_error = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1329,13 +1346,14 @@ mod tests { .create_async() .await; let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) .match_authentication_token(&collector.authentication) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1350,11 +1368,8 @@ mod tests { let job = job.unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_error = server .mock("GET", collection_job_path.as_str()) .with_status(500) @@ -1412,7 +1427,12 @@ mod tests { &random(), &144, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1420,16 +1440,17 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1440,12 +1461,9 @@ mod tests { .await .unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1485,7 +1503,12 @@ mod tests { &random(), &3, ); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1493,16 +1516,17 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1513,13 +1537,10 @@ mod tests { .await .unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1559,37 +1580,39 @@ mod tests { &random(), &true, ); + let idempotency_key = random(); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); - let collect_resp = build_collect_response_fixed(&transcript, &collector, &()); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collect_resp = + build_collect_response_leader_selected(&transcript, &collector, &idempotency_key, &()); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; let job = collector - .collection( - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), - &(), - ) + .collection(Query::new_leader_selected(idempotency_key), &()) .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); mocked_collect_start_success.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -1647,6 +1670,11 @@ mod tests { .with_collect_poll_backoff(test_http_request_exponential_backoff()) .build() .unwrap(); + let collection_job_id: CollectionJobId = random(); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); let batch_interval = Interval::new( Time::from_seconds_since_epoch(1_000_000, &TEST_TIME_PRECISION), @@ -1654,17 +1682,18 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); - let matcher = collection_uri_regex_matcher(&collector.task_id); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mocked_collect_start_success = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) .match_header(AUTHORIZATION.as_str(), "Bearer AAAAAAAAAAAAAAAA") - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1677,11 +1706,8 @@ mod tests { mocked_collect_start_success.assert_async().await; let job = job.unwrap(); assert_eq!(job.query.batch_interval(), &batch_interval); + assert_eq!(job.collection_job_id, collection_job_id); - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mocked_collect_complete = server .mock("GET", collection_job_path.as_str()) .match_header(AUTHORIZATION.as_str(), "Bearer AAAAAAAAAAAAAAAA") @@ -1715,10 +1741,10 @@ mod tests { let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_server_error = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1746,7 +1772,7 @@ mod tests { mock_server_error.assert_async().await; let mock_server_error_details = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1772,7 +1798,7 @@ mod tests { mock_server_error_details.assert_async().await; let mock_bad_request = server - .mock("PUT", matcher) + .mock("POST", matcher) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, @@ -1809,16 +1835,22 @@ mod tests { initialize_rustls(); let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); + let collection_job_id: CollectionJobId = random(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_collect_start = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -1839,6 +1871,7 @@ mod tests { .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); let error = collector.poll_once(&job).await.unwrap_err(); assert_matches!(error, Error::Http(error_response) => { assert_eq!(error_response.status(), StatusCode::INTERNAL_SERVER_ERROR); @@ -1848,10 +1881,6 @@ mod tests { mock_collect_start.assert_async().await; mock_collection_job_server_error.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); let mock_collection_job_server_error_details = server .mock("GET", collection_job_path.as_str()) .with_status(500) @@ -2057,16 +2086,22 @@ mod tests { initialize_rustls(); let mut server = mockito::Server::new_async().await; let configured_vdaf = ConfiguredVdaf::prio3_count().unwrap(); + let collection_job_id = random(); let collector = setup_collector(&mut server, configured_vdaf); - let matcher = collection_uri_regex_matcher(&collector.task_id); + let collection_job_path = format!( + "/tasks/{}/collection_jobs/{}", + collector.task_id, collection_job_id + ); + let matcher = collection_creation_uri_matcher(&collector.task_id); let mock_collect_start = server - .mock("PUT", matcher.clone()) + .mock("POST", matcher.clone()) .match_header( CONTENT_TYPE.as_str(), CollectionJobReq::::MEDIA_TYPE, ) - .with_status(201) + .with_status(200) + .with_header(LOCATION.as_str(), collection_job_path.as_str()) .expect(1) .create_async() .await; @@ -2080,13 +2115,9 @@ mod tests { .start() .await .unwrap(); + assert_eq!(job.collection_job_id, collection_job_id); mock_collect_start.assert_async().await; - let collection_job_path = format!( - "/tasks/{}/collection_jobs/{}", - collector.task_id, job.collection_job_id - ); - let mock_collect_poll_no_retry_after = server .mock("GET", collection_job_path.as_str()) .with_status(200) @@ -2255,7 +2286,7 @@ mod tests { let collection_job_id = random(); let collection_job = CollectionJob::new( collection_job_id, - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(random()), dummy::AggregationParam(1), ); let matcher = collection_uri_regex_matcher(&collector.task_id); @@ -2293,7 +2324,7 @@ mod tests { let collection_job_id = random(); let collection_job = CollectionJob::new( collection_job_id, - Query::new_leader_selected(LeaderSelectedQueryConfig::test_default()), + Query::new_leader_selected(random()), dummy::AggregationParam(1), ); let matcher = collection_uri_regex_matcher(&collector.task_id); @@ -2338,7 +2369,7 @@ mod tests { ) .unwrap(); let collect_resp = - build_collect_response_time(&transcript, &collector, &(), batch_interval); + build_collect_response_time_interval(&transcript, &collector, &(), batch_interval); let job = CollectionJob { collection_job_id: random(), diff --git a/integration_tests/tests/integration/simulation/run.rs b/integration_tests/tests/integration/simulation/run.rs index e07cb392f..815ae259a 100644 --- a/integration_tests/tests/integration/simulation/run.rs +++ b/integration_tests/tests/integration/simulation/run.rs @@ -442,7 +442,6 @@ impl Simulation { .components .collector .collection(query, &()) - .with_id(*collection_job_id) .start() .await { @@ -462,7 +461,6 @@ impl Simulation { .components .collector .collection(query, &()) - .with_id(*collection_job_id) .start() .await { diff --git a/messages/src/batch_mode.rs b/messages/src/batch_mode.rs index 6832cea72..04edebd6f 100644 --- a/messages/src/batch_mode.rs +++ b/messages/src/batch_mode.rs @@ -215,6 +215,12 @@ impl Decode for LeaderSelectedQueryConfig { } } +impl AsRef<[u8]> for LeaderSelectedQueryConfig { + fn as_ref(&self) -> &[u8] { + self.idempotency_key.as_ref() + } +} + /// DAP protocol message indicating a batch mode. #[derive(Copy, Clone, Debug, PartialEq, Eq, TryFromPrimitive, Serialize, Deserialize)] #[repr(u8)] diff --git a/tools/src/bin/collect.rs b/tools/src/bin/collect.rs index 4c821516a..9a3b52e23 100644 --- a/tools/src/bin/collect.rs +++ b/tools/src/bin/collect.rs @@ -19,13 +19,13 @@ use janus_core::{ }; use janus_messages::{ BatchConfig, CollectionJobId, Duration, HpkeConfig, Interval, Query, TaskConfiguration, TaskId, - Time, batch_mode::BatchMode, + Time, + batch_mode::{BatchMode, LeaderSelectedQueryConfig}, }; use prio::{ codec::Decode, vdaf::{self, Vdaf, VdafError}, }; -use rand::random; use tracing_log::LogTracer; use tracing_subscriber::{EnvFilter, Registry, prelude::*}; @@ -132,6 +132,37 @@ fn private_collector_credential_parser( serde_json::from_str(s) } +#[derive(Debug, Args, PartialEq, Eq)] +struct AuthenticationOptions { + /// Authentication token for the DAP-Auth-Token HTTP header + #[clap( + long, + value_parser = StringValueParser::new().try_map( + AuthenticationToken::new_dap_auth_token_from_string, + ), + env, + hide_env_values = true, + help_heading = "Authorization", + display_order = 0, + conflicts_with = "authorization_bearer_token" + )] + dap_auth_token: Option, + + /// Authentication token for the "Authorization: Bearer ..." HTTP header + #[clap( + long, + value_parser = StringValueParser::new().try_map( + AuthenticationToken::new_bearer_token_from_string, + ), + env, + hide_env_values = true, + help_heading = "Authorization", + display_order = 1, + conflicts_with = "dap_auth_token" + )] + authorization_bearer_token: Option, +} + #[derive(Clone)] struct TaskConfigurationValueParser { inner: NonEmptyStringValueParser, @@ -163,35 +194,35 @@ impl TypedValueParser for TaskConfigurationValueParser { } } -#[derive(Debug, Args, PartialEq, Eq)] -struct AuthenticationOptions { - /// Authentication token for the DAP-Auth-Token HTTP header - #[clap( - long, - value_parser = StringValueParser::new().try_map( - AuthenticationToken::new_dap_auth_token_from_string, - ), - env, - hide_env_values = true, - help_heading = "Authorization", - display_order = 0, - conflicts_with = "authorization_bearer_token" - )] - dap_auth_token: Option, +#[derive(Clone)] +struct IdempotencyKeyValueParser { + inner: StringValueParser, +} - /// Authentication token for the "Authorization: Bearer ..." HTTP header - #[clap( - long, - value_parser = StringValueParser::new().try_map( - AuthenticationToken::new_bearer_token_from_string, - ), - env, - hide_env_values = true, - help_heading = "Authorization", - display_order = 1, - conflicts_with = "dap_auth_token" - )] - authorization_bearer_token: Option, +impl IdempotencyKeyValueParser { + fn new() -> Self { + Self { + inner: StringValueParser::new(), + } + } +} + +impl TypedValueParser for IdempotencyKeyValueParser { + type Value = LeaderSelectedQueryConfig; + + fn parse_ref( + &self, + cmd: &clap::Command, + arg: Option<&clap::Arg>, + value: &std::ffi::OsStr, + ) -> Result { + let input = self.inner.parse_ref(cmd, arg, value)?; + let bytes = URL_SAFE_NO_PAD + .decode(input) + .map_err(|err| clap::Error::raw(ErrorKind::ValueValidation, err))?; + Self::Value::get_decoded(&bytes) + .map_err(|err| clap::Error::raw(ErrorKind::ValueValidation, err)) + } } #[derive(Debug, Args, PartialEq, Eq)] @@ -211,6 +242,13 @@ struct QueryOptions { help_heading = "Collect Request Parameters (Time Interval)" )] batch_interval_duration: Option, + + /// Idempotency key for collection job creation, as base64 encoded bytes. + /// + /// If not set, a random idempotency key is used. + #[clap(long, value_parser = IdempotencyKeyValueParser::new(), + help_heading = "Collect Request Parameters (Leader Selected)")] + idempotency_key: Option, } #[derive(Debug, Args, PartialEq, Eq)] @@ -273,12 +311,7 @@ enum Subcommands { /// Initialize a new collection job /// /// Outputs collection job ID to stdout. - NewJob { - /// Job ID to use for the new collection job. If absent, an ID is randomly generated - /// - /// A valid ID consists of 16 randomly selected bytes, encoded with unpadded base64url. - collection_job_id: Option, - }, + NewJob, /// Poll an existing collection job once /// /// The supplied query options must exactly match the ones used to create the collection job, @@ -438,7 +471,7 @@ async fn main() -> anyhow::Result<()> { /// The batch to collect, resolved against the task's batch configuration. enum ResolvedQuery { TimeInterval(Interval), - LeaderSelected, + LeaderSelected(LeaderSelectedQueryConfig), } /// Resolves the query arguments against the batch configuration in `--task-config`. The batch mode @@ -450,8 +483,9 @@ fn resolve_query(options: &Options) -> Result { batch_config, options.query.batch_interval_start, options.query.batch_interval_duration, + options.query.idempotency_key.clone(), ) { - (BatchConfig::TimeInterval, Some(start), Some(duration)) => { + (BatchConfig::TimeInterval, Some(start), Some(duration), _) => { let time_precision = options.task_config.time_precision(); Ok(ResolvedQuery::TimeInterval( Interval::new( @@ -461,23 +495,25 @@ fn resolve_query(options: &Options) -> Result { .map_err(|err| Error::Anyhow(err.into()))?, )) } - (BatchConfig::TimeInterval, None, None) => Err(clap::Error::raw( + (BatchConfig::TimeInterval, None, None, _) => Err(clap::Error::raw( ErrorKind::MissingRequiredArgument, "the task is a time-interval task, so --batch-interval-start and \ --batch-interval-duration are required\n", ) .into()), - (BatchConfig::TimeInterval, _, _) => { + (BatchConfig::TimeInterval, _, _, _) => { unreachable!("clap requires both batch interval arguments together") } - (BatchConfig::LeaderSelected, None, None) => Ok(ResolvedQuery::LeaderSelected), - (BatchConfig::LeaderSelected, _, _) => Err(clap::Error::raw( + (BatchConfig::LeaderSelected, None, None, Some(idempotency_key)) => { + Ok(ResolvedQuery::LeaderSelected(idempotency_key)) + } + (BatchConfig::LeaderSelected, _, _, _) => Err(clap::Error::raw( ErrorKind::ArgumentConflict, "the task is a leader-selected task, so --batch-interval-start and \ --batch-interval-duration must not be provided\n", ) .into()), - (batch_config, _, _) => Err(clap::Error::raw( + (batch_config, _, _, _) => Err(clap::Error::raw( ErrorKind::ValueValidation, format!( "unsupported batch mode in --task-config: {}\n", @@ -495,10 +531,8 @@ macro_rules! options_query_dispatch { let $query = Query::new_time_interval(batch_interval); $body } - ResolvedQuery::LeaderSelected => { - let $query = Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ); + ResolvedQuery::LeaderSelected(idempotency_key) => { + let $query = Query::new_leader_selected(idempotency_key); $body } } @@ -540,9 +574,8 @@ async fn run(options: Options) -> Result<(), Error> { // Every VDAF this tool dispatches to has a trivial aggregation parameter. let agg_param: &::AggregationParam = &Default::default(); match options.subcommand { - Some(Subcommands::NewJob { collection_job_id }) => { - let collection_job_id = collection_job_id.unwrap_or_else(random); - run_new_job(options, vdaf, http_client, query, agg_param, collection_job_id).await + Some(Subcommands::NewJob) => { + run_new_job(options, vdaf, http_client, query, agg_param).await } Some(Subcommands::PollJob { collection_job_id }) => { run_poll_job(options, vdaf, http_client, query, agg_param, collection_job_id).await @@ -577,14 +610,12 @@ async fn run_new_job( http_client: reqwest::Client, query: Query, agg_param: &V::AggregationParam, - collection_job_id: CollectionJobId, ) -> Result<(), Error> where V::AggregateResult: Debug, { let collection = new_collector(options, vdaf, http_client)? .collection(query, agg_param) - .with_id(collection_job_id) .start() .await .map_err(|err| Error::Anyhow(err.into()))?; @@ -701,7 +732,10 @@ mod tests { initialize_rustls, test_util::install_test_trace_subscriber, }; - use janus_messages::{BatchConfig, TaskConfiguration, TaskId, TimePrecision, VdafConfig}; + use janus_messages::{ + BatchConfig, TaskConfiguration, TaskId, TimePrecision, VdafConfig, + batch_mode::LeaderSelectedQueryConfig, + }; use prio::{codec::Encode, vdaf::prio3::Prio3}; use rand::random; use tempfile::NamedTempFile; @@ -780,6 +814,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); @@ -989,6 +1024,8 @@ mod tests { let encoded_hpke_config = URL_SAFE_NO_PAD.encode(hpke_keypair.config().get_encoded().unwrap()); let encoded_private_key = URL_SAFE_NO_PAD.encode(hpke_keypair.private_key().as_ref()); + let encoded_idempotency_key = + URL_SAFE_NO_PAD.encode(LeaderSelectedQueryConfig::test_default().as_ref()); let auth_token = AuthenticationToken::DapAuth(random()); let task_config_argument = task_config_argument(BatchConfig::LeaderSelected); @@ -1010,6 +1047,7 @@ mod tests { query: QueryOptions { batch_interval_start: None, batch_interval_duration: None, + idempotency_key: Some(LeaderSelectedQueryConfig::test_default()), }, }; let correct_arguments = [ @@ -1019,6 +1057,7 @@ mod tests { &format!("--dap-auth-token={}", auth_token.as_str()), &format!("--hpke-config={encoded_hpke_config}"), &format!("--hpke-private-key={encoded_private_key}"), + &format!("--idempotency-key={encoded_idempotency_key}"), ]; match Options::try_parse_from(correct_arguments) { Ok(got) => assert_eq!(got, expected), @@ -1313,9 +1352,7 @@ mod tests { let task_config_argument = task_config_argument(BatchConfig::TimeInterval); let mut expected = Options { - subcommand: Some(Subcommands::NewJob { - collection_job_id: None, - }), + subcommand: Some(Subcommands::NewJob), task_id, task_config: task_config(BatchConfig::TimeInterval), authentication: AuthenticationOptions { @@ -1331,6 +1368,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); @@ -1352,10 +1390,7 @@ mod tests { Err(e) => panic!("{e}\narguments were {correct_arguments:?}"), } - let collection_job_id = random(); - expected.subcommand = Some(Subcommands::NewJob { - collection_job_id: Some(collection_job_id), - }); + expected.subcommand = Some(Subcommands::NewJob); let correct_arguments = [ "collect", &format!("--task-id={task_id_encoded}"), @@ -1368,8 +1403,6 @@ mod tests { "--batch-interval-duration", "1000", "new-job", - "--", // prevent ID from being interpreted as a flag, in case it starts with a hyphen. - &format!("{collection_job_id}"), ]; match Options::try_parse_from(correct_arguments) { Ok(got) => assert_eq!(got, expected), @@ -1405,6 +1438,7 @@ mod tests { query: QueryOptions { batch_interval_start: Some(1_000_000), batch_interval_duration: Some(1_000), + idempotency_key: None, }, }; let task_id_encoded = URL_SAFE_NO_PAD.encode(task_id.get_encoded().unwrap()); diff --git a/tools/tests/cmd/collect.trycmd b/tools/tests/cmd/collect.trycmd index 3e5848b83..74cf36b5e 100644 --- a/tools/tests/cmd/collect.trycmd +++ b/tools/tests/cmd/collect.trycmd @@ -69,4 +69,10 @@ Collect Request Parameters (Time Interval): --batch-interval-duration Duration of the collection batch interval, in seconds +Collect Request Parameters (Leader Selected): + --idempotency-key + Idempotency key for collection job creation, as base64 encoded bytes. + + If not set, a random idempotency key is used. + ``` From 72aea190f5cf9cea279cf56336c44aca8960ee70 Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Tue, 13 Oct 2026 10:21:36 -0400 Subject: [PATCH 4/6] interop/integration: server selected job IDs Adds support to interop binaries and integration tests for server selected collection job IDs and leader-selected batch mode idempotency keys. Re-enables previously neutered integration tests in CI. See --- .github/workflows/ci-build.yml | 14 ++-- integration_tests/tests/integration/common.rs | 7 +- .../tests/integration/simulation/arbitrary.rs | 71 +++++++++--------- .../tests/integration/simulation/model.rs | 21 +++--- .../integration/simulation/reproduction.rs | 73 ++++++++----------- .../tests/integration/simulation/run.rs | 63 +++++----------- .../src/commands/janus_interop_collector.rs | 20 ++--- 7 files changed, 114 insertions(+), 155 deletions(-) diff --git a/.github/workflows/ci-build.yml b/.github/workflows/ci-build.yml index ddeeba44b..983005197 100644 --- a/.github/workflows/ci-build.yml +++ b/.github/workflows/ci-build.yml @@ -90,13 +90,13 @@ jobs: - name: Build (Docker-specific tests) # This matches the `cargo test` command invoked by `cargo xtask test-docker`, except for `--no-run`. run: cargo test --profile ci --locked --package=janus_interop_binaries --package=janus_integration_tests --features=testcontainer --no-run - #- name: Test (Docker) - # id: test-docker - # env: - # RUST_LOG: info - # JANUS_E2E_LOGS_PATH: ${{ github.workspace }}/test-logs - # DIVVIUP_TS_INTEROP_CONTAINER: ${{ steps.default-input-values.outputs.divviup_ts_interop_container }} - # run: cargo run --package xtask --profile=ci -- test-docker --profile=ci --locked + - name: Test (Docker) + id: test-docker + env: + RUST_LOG: info + JANUS_E2E_LOGS_PATH: ${{ github.workspace }}/test-logs + DIVVIUP_TS_INTEROP_CONTAINER: ${{ steps.default-input-values.outputs.divviup_ts_interop_container }} + run: cargo run --package xtask --profile=ci -- test-docker --profile=ci --locked - name: Upload container logs # Only get logs on test failure if: failure() && (steps.test.outcome == 'failure' || steps.test-docker.outcome == 'failure') diff --git a/integration_tests/tests/integration/common.rs b/integration_tests/tests/integration/common.rs index 92d36f183..9021deef9 100644 --- a/integration_tests/tests/integration/common.rs +++ b/integration_tests/tests/integration/common.rs @@ -15,7 +15,7 @@ use janus_integration_tests::{ }; use janus_messages::{ Duration, Interval, Query, Time, TimePrecision, - batch_mode::{self, LeaderSelected}, + batch_mode::{self, LeaderSelected, LeaderSelectedQueryConfig}, problem_type::DapProblemType, }; use prio::{ @@ -325,13 +325,12 @@ where BatchMode::LeaderSelected { .. } => { let mut requests = 0; + let idempotency_key: LeaderSelectedQueryConfig = random(); let collection = loop { requests += 1; let collection_res = collect_generic::<_, LeaderSelected>( &collector, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(idempotency_key.clone()), aggregation_parameter, ) .await; diff --git a/integration_tests/tests/integration/simulation/arbitrary.rs b/integration_tests/tests/integration/simulation/arbitrary.rs index 31638d476..60bffcfc5 100644 --- a/integration_tests/tests/integration/simulation/arbitrary.rs +++ b/integration_tests/tests/integration/simulation/arbitrary.rs @@ -9,13 +9,15 @@ use std::cmp::max; use chrono::TimeDelta; use janus_aggregator_core::task::AggregationMode; use janus_core::time::{DateTimeExt, TimeDeltaExt, TimeExt}; -use janus_messages::{CollectionJobId, Duration, Interval, Time, TimePrecision}; +use janus_messages::{ + Duration, Interval, Time, TimePrecision, batch_mode::LeaderSelectedQueryConfig, +}; use quickcheck::{Arbitrary, Gen, empty_shrinker}; use rand::random; use crate::simulation::{ START_TIME, - model::{Config, Input, Op}, + model::{Config, Input, Op, Query}, }; impl Arbitrary for Config { @@ -86,8 +88,8 @@ pub(super) struct KeyRotatorInput(pub(super) Input); struct Context { current_time: Time, time_precision: TimePrecision, - started_collection_job_ids: Vec, - polled_collection_job_ids: Vec, + started_collection_job_queries: Vec, + polled_collection_job_queries: Vec, } impl Context { @@ -95,8 +97,8 @@ impl Context { Self { current_time: START_TIME.to_time(&config.time_precision), time_precision: config.time_precision, - started_collection_job_ids: Vec::new(), - polled_collection_job_ids: Vec::new(), + started_collection_job_queries: Vec::new(), + polled_collection_job_queries: Vec::new(), } } @@ -108,14 +110,9 @@ impl Context { .add_timedelta(amount, &self.time_precision) .unwrap() } - Op::CollectorStart { - collection_job_id, - query: _, - } => self.started_collection_job_ids.push(*collection_job_id), - Op::CollectorPoll { collection_job_id } - if !self.polled_collection_job_ids.contains(collection_job_id) => - { - self.polled_collection_job_ids.push(*collection_job_id); + Op::CollectorStart { query } => self.started_collection_job_queries.push(query.clone()), + Op::CollectorPoll { query } if !self.polled_collection_job_queries.contains(query) => { + self.polled_collection_job_queries.push(query.clone()); } _ => {} } @@ -187,14 +184,8 @@ fn shrink_op(op: &Op) -> Box> { | Op::CollectionJobDriver | Op::CollectionJobDriverRequestError | Op::CollectionJobDriverResponseError - | Op::CollectorStart { - collection_job_id: _, - query: _, - } - | Op::CollectorPoll { - collection_job_id: _, - } => empty_shrinker(), - + | Op::CollectorStart { .. } + | Op::CollectorPoll { .. } => empty_shrinker(), Op::Upload { report_time: _, count: 0 | 1, @@ -297,21 +288,30 @@ fn arbitrary_collector_start_op_time_interval(g: &mut Gen, context: &Context) -> ]) .unwrap(); Op::CollectorStart { - collection_job_id: random(), - query: super::model::Query::TimeInterval( - Interval::new(start, duration_fn(g, context)).unwrap(), - ), + query: Query::TimeInterval(Interval::new(start, duration_fn(g, context)).unwrap()), + } +} + +/// Generate a collect poll operation for time interval batch mode. +fn arbitrary_collector_poll_op_time_interval(g: &mut Gen, context: &Context) -> Op { + Op::CollectorPoll { + query: g + .choose(&context.started_collection_job_queries) + .cloned() + .unwrap_or_else(|| { + Query::TimeInterval(Interval::minimal(Time::from_time_precision_units(0)).unwrap()) + }), } } -/// Generate a collect poll operation. -fn arbitrary_collector_poll_op(g: &mut Gen, context: &Context) -> Op { +/// Generate a collect poll operation for leader selected batch mode. +fn arbitrary_collector_poll_op_leader_selected(g: &mut Gen, context: &Context) -> Op { Op::CollectorPoll { - collection_job_id: g - .choose(&context.started_collection_job_ids) - .copied() + query: g + .choose(&context.started_collection_job_queries) + .cloned() .unwrap_or_else(|| { - CollectionJobId::try_from([0u8; CollectionJobId::LEN].as_slice()).unwrap() + Query::LeaderSelected(LeaderSelectedQueryConfig::new(vec![0u8; 8]).unwrap()) }), } } @@ -438,7 +438,7 @@ fn arbitrary_op_time_interval(g: &mut Gen, context: &Context, choices: &[OpKind] OpKind::CollectionJobDriverRequestError => Op::CollectionJobDriverRequestError, OpKind::CollectionJobDriverResponseError => Op::CollectionJobDriverResponseError, OpKind::CollectorStart => arbitrary_collector_start_op_time_interval(g, context), - OpKind::CollectorPoll => arbitrary_collector_poll_op(g, context), + OpKind::CollectorPoll => arbitrary_collector_poll_op_time_interval(g, context), } } @@ -491,10 +491,9 @@ fn arbitrary_op_leader_selected(g: &mut Gen, context: &Context, choices: &[OpKin OpKind::CollectionJobDriverRequestError => Op::CollectionJobDriverRequestError, OpKind::CollectionJobDriverResponseError => Op::CollectionJobDriverResponseError, OpKind::CollectorStart => Op::CollectorStart { - collection_job_id: random(), - query: super::model::Query::LeaderSelected, + query: Query::LeaderSelected(random()), }, - OpKind::CollectorPoll => arbitrary_collector_poll_op(g, context), + OpKind::CollectorPoll => arbitrary_collector_poll_op_leader_selected(g, context), } } diff --git a/integration_tests/tests/integration/simulation/model.rs b/integration_tests/tests/integration/simulation/model.rs index e1a0baa44..f90dae95a 100644 --- a/integration_tests/tests/integration/simulation/model.rs +++ b/integration_tests/tests/integration/simulation/model.rs @@ -1,6 +1,8 @@ use chrono::TimeDelta; use janus_aggregator_core::task::AggregationMode; -use janus_messages::{CollectionJobId, Duration, Interval, Time, TimePrecision}; +use janus_messages::{ + Duration, Interval, Time, TimePrecision, batch_mode::LeaderSelectedQueryConfig, +}; #[derive(Debug, Clone)] pub(super) struct Input { @@ -121,21 +123,22 @@ pub(super) enum Op { /// Forward all requests, but drop the responses, and return some sort of error. CollectionJobDriverResponseError, - /// The collector sends a collection request to the leader. It remembers the collection job ID. - CollectorStart { - collection_job_id: CollectionJobId, - query: Query, - }, + /// The collector sends a collection request to the leader. + /// + /// It remembers a mapping of queries to the collection job IDs. + CollectorStart { query: Query }, /// The collector sends a request to the leader to poll an existing collection job. - CollectorPoll { collection_job_id: CollectionJobId }, + /// + /// The collection job matching the query is looked up in the collector state. + CollectorPoll { query: Query }, } /// Representation of a DAP query used in a collection job. -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub(super) enum Query { /// A time interval query, parameterized with a batch interval. TimeInterval(Interval), /// A leader-selected query. - LeaderSelected, + LeaderSelected(LeaderSelectedQueryConfig), } diff --git a/integration_tests/tests/integration/simulation/reproduction.rs b/integration_tests/tests/integration/simulation/reproduction.rs index d27e9203c..fe6ab6e0c 100644 --- a/integration_tests/tests/integration/simulation/reproduction.rs +++ b/integration_tests/tests/integration/simulation/reproduction.rs @@ -21,8 +21,14 @@ fn successful_collection_time_interval() { install_test_trace_subscriber(); initialize_rustls(); - let collection_job_id = random(); let time_precision = TimePrecision::from_seconds(3600); + let query = Query::TimeInterval( + Interval::new( + Time::from_seconds_since_epoch(1_699_999_200, &time_precision), + Duration::from_seconds(3600, &time_precision), + ) + .unwrap(), + ); let input = Input { is_leader_selected: false, config: Config { @@ -58,7 +64,6 @@ fn successful_collection_time_interval() { Op::LeaderAggregationJobDriver, Op::LeaderGarbageCollector, Op::CollectorStart { - collection_job_id, query: Query::TimeInterval( Interval::new( Time::from_seconds_since_epoch(1_699_999_200, &time_precision), @@ -68,7 +73,9 @@ fn successful_collection_time_interval() { ), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { + query: query.clone(), + }, Op::Upload { report_time: START_TIME.to_time(&time_precision), count: 4, @@ -76,17 +83,10 @@ fn successful_collection_time_interval() { Op::AggregationJobCreator, Op::LeaderAggregationJobDriver, Op::CollectorStart { - collection_job_id, - query: Query::TimeInterval( - Interval::new( - Time::from_seconds_since_epoch(1_699_999_200, &time_precision), - Duration::from_seconds(3600, &time_precision), - ) - .unwrap(), - ), + query: query.clone(), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { query }, ]), }; let (test_result, metrics) = Simulation::run_with_metrics(input); @@ -129,8 +129,8 @@ fn successful_collection_leader_selected() { install_test_trace_subscriber(); initialize_rustls(); - let collection_job_id = random(); let time_precision = TimePrecision::from_seconds(3600); + let query = Query::LeaderSelected(random()); let input = Input { is_leader_selected: true, config: Config { @@ -166,11 +166,12 @@ fn successful_collection_leader_selected() { Op::LeaderAggregationJobDriver, Op::LeaderGarbageCollector, Op::CollectorStart { - collection_job_id, - query: Query::LeaderSelected, + query: query.clone(), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { + query: query.clone(), + }, Op::Upload { report_time: START_TIME.to_time(&time_precision), count: 4, @@ -178,11 +179,10 @@ fn successful_collection_leader_selected() { Op::AggregationJobCreator, Op::LeaderAggregationJobDriver, Op::CollectorStart { - collection_job_id, - query: Query::LeaderSelected, + query: query.clone(), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { query }, ]), }; let (test_result, metrics) = Simulation::run_with_metrics(input); @@ -225,8 +225,14 @@ fn successful_collection_asynchronous() { install_test_trace_subscriber(); initialize_rustls(); - let collection_job_id = random(); let time_precision = TimePrecision::from_seconds(3600); + let query = Query::TimeInterval( + Interval::new( + Time::from_seconds_since_epoch(1_699_999_200, &time_precision), + Duration::from_seconds(3600, &time_precision), + ) + .unwrap(), + ); let input = Input { is_leader_selected: false, config: Config { @@ -277,17 +283,12 @@ fn successful_collection_asynchronous() { Op::LeaderAggregationJobDriver, Op::LeaderGarbageCollector, Op::CollectorStart { - collection_job_id, - query: Query::TimeInterval( - Interval::new( - Time::from_seconds_since_epoch(1_699_999_200, &time_precision), - Duration::from_seconds(3600, &time_precision), - ) - .unwrap(), - ), + query: query.clone(), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { + query: query.clone(), + }, Op::Upload { report_time: START_TIME.to_time(&time_precision), count: 4, @@ -300,17 +301,10 @@ fn successful_collection_asynchronous() { }, Op::LeaderAggregationJobDriver, Op::CollectorStart { - collection_job_id, - query: Query::TimeInterval( - Interval::new( - Time::from_seconds_since_epoch(1_699_999_200, &time_precision), - Duration::from_seconds(3600, &time_precision), - ) - .unwrap(), - ), + query: query.clone(), }, Op::CollectionJobDriver, - Op::CollectorPoll { collection_job_id }, + Op::CollectorPoll { query }, ]), }; let (test_result, metrics) = Simulation::run_with_metrics(input); @@ -489,7 +483,6 @@ fn repro_abandoned_aggregation_job_batch_mismatch() { install_test_trace_subscriber(); initialize_rustls(); - let collection_job_id = random(); let time_precision = TimePrecision::from_seconds(1000); let input = Input { is_leader_selected: false, @@ -525,7 +518,6 @@ fn repro_abandoned_aggregation_job_batch_mismatch() { }, Op::LeaderAggregationJobDriver, Op::CollectorStart { - collection_job_id, query: Query::TimeInterval( Interval::new( START_TIME.to_time(&time_precision), @@ -571,7 +563,6 @@ fn repro_helper_accumulate_on_retried_request() { }, Op::LeaderAggregationJobDriver, Op::CollectorStart { - collection_job_id: random(), query: Query::TimeInterval( Interval::new( START_TIME.to_time(&time_precision), diff --git a/integration_tests/tests/integration/simulation/run.rs b/integration_tests/tests/integration/simulation/run.rs index 815ae259a..273de38b6 100644 --- a/integration_tests/tests/integration/simulation/run.rs +++ b/integration_tests/tests/integration/simulation/run.rs @@ -150,16 +150,11 @@ impl Simulation { .execute_collection_job_driver_response_error() .await } - Op::CollectorStart { - collection_job_id, - query, - } => { - simulation - .execute_collector_start(collection_job_id, query) - .await + Op::CollectorStart { query } => { + simulation.execute_collector_start(query).await } - Op::CollectorPoll { collection_job_id } => { - simulation.execute_collector_poll(collection_job_id).await + Op::CollectorPoll { query } => { + simulation.execute_collector_poll(query).await } }; info!(elapsed = ?start.elapsed(), "finished operation"); @@ -430,44 +425,39 @@ impl Simulation { result } - async fn execute_collector_start( - &mut self, - collection_job_id: &CollectionJobId, - query: &Query, - ) -> ControlFlow { + async fn execute_collector_start(&mut self, query: &Query) -> ControlFlow { match query { Query::TimeInterval(interval) => { - let query = janus_messages::Query::new_time_interval(*interval); match self .components .collector - .collection(query, &()) + .collection(janus_messages::Query::new_time_interval(*interval), &()) .start() .await { Ok(collection_job) => { self.state .collection_jobs_time_interval - .insert(*collection_job_id, collection_job); + .insert(query.clone(), collection_job); } Err(error) => info!(?error, "collector error"), } } - Query::LeaderSelected => { - let query = janus_messages::Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ); + Query::LeaderSelected(idempotency_key) => { match self .components .collector - .collection(query, &()) + .collection( + janus_messages::Query::new_leader_selected(idempotency_key.clone()), + &(), + ) .start() .await { Ok(collection_job) => { self.state .collection_jobs_leader_selected - .insert(*collection_job_id, collection_job); + .insert(query.clone(), collection_job); } Err(error) => info!(?error, "collector error"), } @@ -476,15 +466,8 @@ impl Simulation { ControlFlow::Continue(()) } - async fn execute_collector_poll( - &mut self, - collection_job_id: &CollectionJobId, - ) -> ControlFlow { - if let Some(collection_job) = self - .state - .collection_jobs_time_interval - .get(collection_job_id) - { + async fn execute_collector_poll(&mut self, query: &Query) -> ControlFlow { + if let Some(collection_job) = self.state.collection_jobs_time_interval.get(query) { let result = self.components.collector.poll_once(collection_job).await; match result { Ok(PollResult::CollectionResult(collection)) => { @@ -494,7 +477,7 @@ impl Simulation { let old_opt = self .state .aggregate_results_time_interval - .insert(*collection_job_id, collection); + .insert(*collection_job.collection_job_id(), collection); if let Some(old_collection) = old_opt { if report_count != old_collection.report_count() || &interval != old_collection.interval() @@ -508,11 +491,7 @@ impl Simulation { Ok(PollResult::NotReady(_)) => {} Err(error) => info!(?error, "collector error"), } - } else if let Some(collection_job) = self - .state - .collection_jobs_leader_selected - .get(collection_job_id) - { + } else if let Some(collection_job) = self.state.collection_jobs_leader_selected.get(query) { let result = self.components.collector.poll_once(collection_job).await; match result { Ok(PollResult::CollectionResult(collection)) => { @@ -522,7 +501,7 @@ impl Simulation { let old_opt = self .state .aggregate_results_leader_selected - .insert(*collection_job_id, collection); + .insert(*collection_job.collection_job_id(), collection); if let Some(old_collection) = old_opt { if report_count != old_collection.report_count() || &interval != old_collection.interval() @@ -566,10 +545,8 @@ pub(super) struct State { pub(super) runtime_manager: TestRuntimeManager<&'static str>, pub(super) vdaf_instance: VdafInstance, pub(super) vdaf: Prio3Histogram, - pub(super) collection_jobs_time_interval: - HashMap>, - pub(super) collection_jobs_leader_selected: - HashMap>, + pub(super) collection_jobs_time_interval: HashMap>, + pub(super) collection_jobs_leader_selected: HashMap>, pub(super) aggregate_results_time_interval: HashMap>>, pub(super) aggregate_results_leader_selected: HashMap>>, pub(super) next_measurement: usize, diff --git a/interop_binaries/src/commands/janus_interop_collector.rs b/interop_binaries/src/commands/janus_interop_collector.rs index 72f442a19..bb8b722ab 100644 --- a/interop_binaries/src/commands/janus_interop_collector.rs +++ b/interop_binaries/src/commands/janus_interop_collector.rs @@ -437,9 +437,7 @@ async fn handle_collection_start( handle_collect_generic( http_client, task_state, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), vdaf, &agg_param, |result| AggregationResult::Number(NumberAsString((*result).into())), @@ -453,9 +451,7 @@ async fn handle_collection_start( handle_collect_generic( http_client, task_state, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), vdaf, &agg_param, |result| AggregationResult::Number(NumberAsString(u128::from(*result))), @@ -477,9 +473,7 @@ async fn handle_collection_start( handle_collect_generic( http_client, task_state, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), vdaf, &agg_param, |result| { @@ -507,9 +501,7 @@ async fn handle_collection_start( handle_collect_generic( http_client, task_state, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), vdaf, &agg_param, |result: &Vec| { @@ -538,9 +530,7 @@ async fn handle_collection_start( handle_collect_generic( http_client, task_state, - Query::new_leader_selected( - janus_messages::batch_mode::LeaderSelectedQueryConfig::test_default(), - ), + Query::new_leader_selected(random()), vdaf, &agg_param, |result| { From c7ea143f472c224795eb2480a18dab4cf5f7243a Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Mon, 19 Oct 2026 13:58:17 -0400 Subject: [PATCH 5/6] clarify batch aggregation scrubbing Adds some doccomments on items related to scrubbed batch aggregations, as well as tidying up comments in the DB schema so they don't scroll off the right side of one's editor window. --- aggregator_core/src/datastore/models.rs | 22 ++++++++++---- db/00000000000001_initial_schema.up.sql | 39 ++++++++++++++++++------- 2 files changed, 45 insertions(+), 16 deletions(-) diff --git a/aggregator_core/src/datastore/models.rs b/aggregator_core/src/datastore/models.rs index 844abfec5..e7b96a4be 100644 --- a/aggregator_core/src/datastore/models.rs +++ b/aggregator_core/src/datastore/models.rs @@ -1717,18 +1717,30 @@ impl> Eq { } -// ReportAggregationStateCode exists alongside the public ReportAggregationState because there is no -// apparent way to denote a Postgres enum literal without deriving FromSql/ToSql on a Rust enum -// type, but it is not possible to derive FromSql/ToSql on a non-C-style enum. +/// ReportAggregationStateCode corresponds to the `batch_aggregation_state` enum. +/// +/// It exists alongside the public ReportAggregationState because there is no apparent way to denote +/// a Postgres enum literal without deriving FromSql/ToSql on a Rust enum type, but it is not +/// possible to derive FromSql/ToSql on a non-C-style enum. #[derive(Debug, Clone, Copy, PartialEq, Eq, FromSql, ToSql)] #[postgres(name = "batch_aggregation_state")] pub(super) enum BatchAggregationStateCode { - /// This batch aggregation has not been collected & permits further aggregation. + /// This batch aggregation has not been collected. + /// + /// Report shares may still be aggregated into it. #[postgres(name = "AGGREGATING")] Aggregating, - /// This batch aggregation has been collected & no longer permits aggregation. + /// This batch aggregation has been collected. + /// + /// It contains a batch aggregate share which may be merged into an aggregate share to service a + /// collection job or aggregate share request. No more report shares may be aggregated into it. #[postgres(name = "COLLECTED")] Collected, + /// This batch aggregation has been scrubbed. + /// + /// Report shares have been aggregated into it and its batch aggregate has been merged into an + /// aggregate share. It has been scrubbed of data to save space. No more report shares may be + /// aggregated into it. #[postgres(name = "SCRUBBED")] Scrubbed, } diff --git a/db/00000000000001_initial_schema.up.sql b/db/00000000000001_initial_schema.up.sql index 9fc91009e..3e24014cb 100644 --- a/db/00000000000001_initial_schema.up.sql +++ b/db/00000000000001_initial_schema.up.sql @@ -375,23 +375,40 @@ CREATE TYPE BATCH_AGGREGATION_STATE AS ENUM( -- Information on aggregation for a single batch. This information may be incremental if the VDAF -- supports incremental aggregation. Each batch's aggregation is sharded via the `ord` column. CREATE TABLE batch_aggregations( - id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, -- artificial ID, internal-only - task_id BIGINT NOT NULL, -- the task ID - batch_identifier BYTEA NOT NULL, -- encoded batch mode-specific batch identifier (corresponds to identifier in BatchSelector) + -- artificial ID, internal-only + id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, + -- the task ID + task_id BIGINT NOT NULL, + -- encoded batch mode-specific batch identifier (corresponds to identifier in BatchSelector) + batch_identifier BYTEA NOT NULL, -- batch interval, as an INT8RANGE, in time precision units. -- Populated only for time-interval tasks. Will always match batch_identifier. batch_interval INT8RANGE, - aggregation_param BYTEA NOT NULL, -- the aggregation parameter (opaque VDAF message) - ord BIGINT NOT NULL, -- the index of this batch aggregation shard, over (task ID, batch_identifier, aggregation_param). + -- the aggregation parameter (opaque VDAF message) + aggregation_param BYTEA NOT NULL, + -- the index of this batch aggregation shard, over (task ID, batch_identifier, + -- aggregation_param). + ord BIGINT NOT NULL, -- The minimal interval, in time precision units, containing all of client timestamps included -- in this batch aggregation. client_timestamp_interval INT8RANGE NOT NULL, - state BATCH_AGGREGATION_STATE NOT NULL, -- the current state of this batch aggregation - aggregate_share BYTEA, -- the (possibly-incremental) aggregate share; populated unless report_count is 0 or the batch aggregation has been scrubbed. - report_count BIGINT, -- the (possibly-incremental) client report count; populated unless the batch aggregation has been scrubbed. - checksum BYTEA, -- the (possibly-incremental) checksum; populated unless the batch aggregation has been scrubbed. - aggregation_jobs_created BIGINT, -- the number of aggregation jobs that have been created for this batch; populated unless the batch aggregation has been scrubbed. - aggregation_jobs_terminated BIGINT, -- the number of aggregation jobs that have been terminated (finished or abandoned) for this batch; populated unless the batch aggregation has been scrubbed. + -- the current state of this batch aggregation + state BATCH_AGGREGATION_STATE NOT NULL, + -- the (possibly-incremental) aggregate share; populated unless report_count is 0 or the batch + -- aggregation has been scrubbed. + aggregate_share BYTEA, + -- the (possibly-incremental) client report count; populated unless the batch aggregation has + -- been scrubbed. + report_count BIGINT, + -- the (possibly-incremental) checksum; populated unless the batch aggregation has been + -- scrubbed. + checksum BYTEA, + -- the number of aggregation jobs that have been created for this batch; populated unless the + -- batch aggregation has been scrubbed. + aggregation_jobs_created BIGINT, + -- the number of aggregation jobs that have been terminated (finished or abandoned) for this + -- batch; populated unless the batch aggregation has been scrubbed. + aggregation_jobs_terminated BIGINT, -- creation/update records created_at TIMESTAMP WITH TIME ZONE NOT NULL, -- when the row was created From 602eae9e6b3687c98735737cefe6291d1ecd7af4 Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Mon, 19 Oct 2026 14:20:48 -0400 Subject: [PATCH 6/6] Tests for deleted/expired collection job creation Adds test cases for trying to re-create collection jobs that have been either deleted explicitly by a DELETE or implicitly by their constituent data expiring. Along the way, we do some pretty significant refactoring of `CollectionJobTest` so that various test setup routines become methods on that type. --- .../src/aggregator/collection_job_tests.rs | 623 ++++++++++++------ .../http_handlers/tests/collection_job.rs | 503 +++++++------- 2 files changed, 700 insertions(+), 426 deletions(-) diff --git a/aggregator/src/aggregator/collection_job_tests.rs b/aggregator/src/aggregator/collection_job_tests.rs index 6c27913c7..4dc0ec9d5 100644 --- a/aggregator/src/aggregator/collection_job_tests.rs +++ b/aggregator/src/aggregator/collection_job_tests.rs @@ -29,7 +29,7 @@ use janus_core::{ }; use janus_messages::{ AggregateShareAad, AggregationJobStep, BatchId, CollectionJobId, CollectionJobReq, - CollectionJobResp, Interval, MediaType, Query, Role, TaskId, Time, TimePrecision, + CollectionJobResp, Duration, Interval, MediaType, Query, Role, TaskId, Time, TimePrecision, batch_mode::{ BatchMode as BatchModeTrait, LeaderSelected, LeaderSelectedQueryConfig, TimeInterval, }, @@ -58,6 +58,79 @@ pub(crate) struct CollectionJobTestCase { } impl CollectionJobTestCase { + pub(crate) async fn setup(role: Role, batch_mode: BatchMode) -> Self { + Self::setup_with_report_expiry_age(role, batch_mode, None).await + } + + pub(crate) async fn setup_with_report_expiry_age( + role: Role, + batch_mode: BatchMode, + report_expiry_age: Option, + ) -> Self { + install_test_trace_subscriber(); + + let task = TaskBuilder::new( + batch_mode, + AggregationMode::Synchronous, + VdafInstance::Fake { rounds: 1 }, + ) + .with_report_expiry_age(report_expiry_age) + .build(); + let role_task = task.view_for_role(role).unwrap(); + let clock = MockClock::default(); + let ephemeral_datastore = ephemeral_datastore().await; + let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await); + + datastore.put_aggregator_task(&role_task).await.unwrap(); + datastore.put_hpke_key().await.unwrap(); + + let router = AggregatorHandlerBuilder::new( + Arc::clone(&datastore), + TestRuntime::default(), + &noop_meter(), + Config { + batch_aggregation_shard_count: BATCH_AGGREGATION_SHARD_COUNT, + ..Default::default() + }, + ) + .await + .unwrap() + .build() + .unwrap(); + + CollectionJobTestCase { + task, + clock, + router, + datastore, + _ephemeral_datastore: ephemeral_datastore, + } + } + + pub(super) async fn setup_leader_selected_current_batch( + &mut self, + ) -> (BatchId, BatchId, Interval) { + // Fill the datastore with the necessary data so that there are two outstanding batches to be + // collected. + let time = self + .clock + .now() + .to_time(self.task.time_precision()); + let batch_id_1 = self + .setup_leader_selected_batch(time, self.task.min_batch_size() + 1) + .await; + let batch_id_2 = self + .setup_leader_selected_batch(time, self.task.min_batch_size() + 1) + .await; + let client_timestamp_interval = Interval::minimal(time).unwrap(); + + (batch_id_1, batch_id_2, client_timestamp_interval) + } + + pub(super) fn clock(&self) -> &MockClock { + &self.clock + } + pub(super) async fn post_collection_job_with_auth_token( &self, request: &CollectionJobReq, @@ -124,6 +197,29 @@ impl CollectionJobTestCase { .await } + pub(super) async fn delete_collection_job( + &self, + collection_job_id: &CollectionJobId, + ) -> http::Response { + self.router + .clone() + .oneshot( + Request::builder() + .method("DELETE") + .uri( + self.task + .collection_job_uri(&collection_job_id) + .unwrap() + .path(), + ) + .with_authentication_token(self.task.collector_auth_token()) + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap() + } + /// Seed the database with reports and various aggregation artifacts from aggregating them in a /// leader-selected batch. pub(super) async fn setup_leader_selected_batch( @@ -252,148 +348,70 @@ impl CollectionJobTestCase { .await .unwrap() } -} - -/// Parse task ID and collection job ID out of response location header. -/// -/// Panics if header is missing from response or value is malformed. -pub(super) fn ids_from_collection_job_response( - response: &http::Response, -) -> (TaskId, CollectionJobId) { - response - .headers() - .ids_from_collection_job_location() - .unwrap() - .unwrap() -} -pub(crate) async fn setup_collection_job_test_case( - role: Role, - batch_mode: BatchMode, -) -> CollectionJobTestCase { - install_test_trace_subscriber(); + pub(super) async fn simulate_successful_collection_job_time_interval( + &mut self, + report_time: Time, + queried_interval: Interval, + ) -> (CollectionJobReq, CollectionJobId) { + self.setup_time_interval_batch(report_time).await; - let task = TaskBuilder::new( - batch_mode, - AggregationMode::Synchronous, - VdafInstance::Fake { rounds: 1 }, - ) - .build(); - let role_task = task.view_for_role(role).unwrap(); - let clock = MockClock::default(); - let ephemeral_datastore = ephemeral_datastore().await; - let datastore = Arc::new(ephemeral_datastore.datastore(clock.clone()).await); - - datastore.put_aggregator_task(&role_task).await.unwrap(); - datastore.put_hpke_key().await.unwrap(); - - let router = AggregatorHandlerBuilder::new( - Arc::clone(&datastore), - TestRuntime::default(), - &noop_meter(), - Config { - batch_aggregation_shard_count: BATCH_AGGREGATION_SHARD_COUNT, - ..Default::default() - }, - ) - .await - .unwrap() - .build() - .unwrap(); - - CollectionJobTestCase { - task, - clock, - router, - datastore, - _ephemeral_datastore: ephemeral_datastore, - } -} + let aggregation_param = dummy::AggregationParam::default(); + let leader_aggregate_share = dummy::AggregateShare(0); + let helper_aggregate_share = dummy::AggregateShare(1); -async fn setup_leader_selected_current_batch_collection_job_test_case() --> (CollectionJobTestCase, BatchId, BatchId, Interval) { - let test_case = setup_collection_job_test_case( - Role::Leader, - BatchMode::LeaderSelected { - batch_time_window_size: None, - }, - ) - .await; - - // Fill the datastore with the necessary data so that there are two outstanding batches to be - // collected. - let time = test_case - .clock - .now() - .to_time(test_case.task.time_precision()); - let batch_id_1 = test_case - .setup_leader_selected_batch(time, test_case.task.min_batch_size() + 1) - .await; - let batch_id_2 = test_case - .setup_leader_selected_batch(time, test_case.task.min_batch_size() + 1) - .await; - let client_timestamp_interval = Interval::minimal(time).unwrap(); - - (test_case, batch_id_1, batch_id_2, client_timestamp_interval) -} - -#[tokio::test] -async fn collection_job_success_leader_selected() { - // This test drives two current batch collection jobs to completion, verifying that distinct - // collection job IDs are minted and distinct batch IDs are collected each time. Then, we - // attempt to collect another current batch, which must fail as there are no more outstanding - // batches. - let (test_case, batch_id_1, batch_id_2, spanned_interval) = - setup_leader_selected_current_batch_collection_job_test_case().await; - - let first_idempotency_key: LeaderSelectedQueryConfig = random(); - let mut saw_batch_id_1 = false; - let mut saw_batch_id_2 = false; - let mut collection_job_ids = vec![]; - let mut collection_job_resps = vec![]; - let vdaf = dummy::Vdaf::new(1); - let leader_aggregate_share = dummy::AggregateShare(0); - let helper_aggregate_share = dummy::AggregateShare(1); - let aggregation_param = dummy::AggregationParam::default(); - - for idempotency_key in [first_idempotency_key.clone(), random()] { let request = CollectionJobReq::new( - Query::new_leader_selected(idempotency_key), + Query::new_time_interval(queried_interval), aggregation_param.get_encoded().unwrap(), ); - let mut response = test_case.post_collection_job(&request).await; + + let mut response = self.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::OK); let (task_id, collection_job_id) = ids_from_collection_job_response(&response); - collection_job_ids.push(collection_job_id); - assert_eq!(test_case.task.id(), &task_id); + assert_eq!(&task_id, self.task.id()); + + let got_collection_job = self + .datastore + .run_unnamed_tx(|tx| { + let task_id = *self.task.id(); + + Box::pin(async move { + Ok(tx + .get_collection_job(&dummy::Vdaf::new(1), &task_id, &collection_job_id) + .await + .unwrap() + .unwrap()) + }) + }) + .await + .unwrap(); + + let want_collection_job = CollectionJob::<0, TimeInterval, dummy::Vdaf>::new( + *self.task.id(), + collection_job_id, + // This is chosen by the leader, we have to copy it + *got_collection_job.aggregate_share_id(), + Query::new_time_interval(queried_interval), + aggregation_param, + queried_interval, + CollectionJobState::Start, + ); + + assert_eq!(want_collection_job, got_collection_job); assert!(take_response_body(&mut response).await.is_empty()); - let mut response = test_case.get_collection_job(&collection_job_id).await; + let mut response = self.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::OK); assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); assert!(take_response_body(&mut response).await.is_empty()); - // Update the collection job with the aggregate shares. collection job should now be - // complete. - let batch_id = test_case - .datastore + // Update the collection job with the aggregate shares and some aggregation jobs. collection + // job should now be complete. + self.datastore .run_unnamed_tx(|tx| { - let task = test_case.task.clone(); - let vdaf = vdaf.clone(); + let task = self.task.clone(); let helper_aggregate_share_bytes = helper_aggregate_share.get_encoded().unwrap(); - let request = request.clone(); Box::pin(async move { - let collection_job = tx - .get_collection_job::<0, LeaderSelected, dummy::Vdaf>( - &vdaf, - task.id(), - &collection_job_id, - ) - .await - .unwrap() - .unwrap(); - let batch_id = *collection_job.batch_identifier(); - let encrypted_helper_aggregate_share = hpke::seal( task.collector_hpke_keypair().config(), &HpkeApplicationInfo::new( @@ -402,43 +420,50 @@ async fn collection_job_success_leader_selected() { &Role::Collector, ), &helper_aggregate_share_bytes, - &AggregateShareAad::new(*task.id(), task.task_configuration(), request) - .get_encoded() - .unwrap(), + &AggregateShareAad::new( + *task.id(), + task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), + ) + .get_encoded() + .unwrap(), ) .unwrap(); - tx.update_collection_job::<0, LeaderSelected, dummy::Vdaf>( - &collection_job.with_state(CollectionJobState::Finished { - report_count: task.min_batch_size() + 1, - client_timestamp_interval: spanned_interval, + let collection_job = tx + .get_collection_job::<0, TimeInterval, dummy::Vdaf>( + &dummy::Vdaf::new(1), + task.id(), + &collection_job_id, + ) + .await + .unwrap() + .unwrap() + .with_state(CollectionJobState::Finished { + report_count: 12, + client_timestamp_interval: queried_interval, encrypted_helper_aggregate_share, leader_aggregate_share, - }), - ) - .await - .unwrap(); + }); - Ok(batch_id) + tx.update_collection_job::<0, TimeInterval, dummy::Vdaf>(&collection_job) + .await + .unwrap(); + Ok(()) }) }) .await .unwrap(); - if batch_id.eq(&batch_id_1) { - saw_batch_id_1 = true; - } else if batch_id.eq(&batch_id_2) { - saw_batch_id_2 = true; - } else { - panic!("unexpected batch ID"); - } - - let mut response = test_case.get_collection_job(&collection_job_id).await; + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::OK); assert_eq!( response.headers().get("content-type").unwrap(), CollectionJobResp::MEDIA_TYPE, ); - let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; let ( report_count, @@ -447,31 +472,29 @@ async fn collection_job_success_leader_selected() { helper_encrypted_aggregate_share, ) = assert_matches!( collect_resp.clone(), - CollectionJobResp { + CollectionJobResp{ report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share, .. - } => ( - report_count, - interval, - leader_encrypted_agg_share, - helper_encrypted_agg_share - ) + } => (report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share) ); - assert_eq!(report_count, test_case.task.min_batch_size() + 1); - assert_eq!(interval, spanned_interval); + assert_eq!(report_count, 12); + assert_eq!(interval, queried_interval); let decrypted_leader_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), + self.task.collector_hpke_keypair(), &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), &leader_encrypted_aggregate_share, &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - request.clone(), + *self.task.id(), + self.task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), ) .get_encoded() .unwrap(), @@ -483,13 +506,16 @@ async fn collection_job_success_leader_selected() { ); let decrypted_helper_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), + self.task.collector_hpke_keypair(), &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), &helper_encrypted_aggregate_share, &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - request.clone(), + *self.task.id(), + self.task.task_configuration(), + CollectionJobReq::new( + Query::new_time_interval(queried_interval), + aggregation_param.get_encoded().unwrap(), + ), ) .get_encoded() .unwrap(), @@ -500,20 +526,231 @@ async fn collection_job_success_leader_selected() { dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() ); - collection_job_resps.push(collect_resp); + (request, collection_job_id) + } + + pub(super) async fn simulate_successful_collection_job_leader_selected( + &mut self, + batch_id_1: BatchId, + batch_id_2: BatchId, + spanned_interval: Interval, + ) -> (CollectionJobReq, CollectionJobId) { + let mut saw_batch_id_1 = false; + let mut saw_batch_id_2 = false; + let mut collection_job_ids = vec![]; + let mut collection_job_resps = vec![]; + let vdaf = dummy::Vdaf::new(1); + let leader_aggregate_share = dummy::AggregateShare(0); + let helper_aggregate_share = dummy::AggregateShare(1); + let aggregation_param = dummy::AggregationParam::default(); + let requests = [ + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + CollectionJobReq::new( + Query::new_leader_selected(random()), + aggregation_param.get_encoded().unwrap(), + ), + ]; + + for request in requests.clone() { + let mut response = self.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::OK); + let (task_id, collection_job_id) = ids_from_collection_job_response(&response); + collection_job_ids.push(collection_job_id); + assert_eq!(self.task.id(), &task_id); + assert!(take_response_body(&mut response).await.is_empty()); + + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::OK); + assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); + assert!(take_response_body(&mut response).await.is_empty()); + + // Update the collection job with the aggregate shares. collection job should now be + // complete. + let batch_id = self + .datastore + .run_unnamed_tx(|tx| { + let task = self.task.clone(); + let vdaf = vdaf.clone(); + let helper_aggregate_share_bytes = + helper_aggregate_share.get_encoded().unwrap(); + let request = request.clone(); + Box::pin(async move { + let collection_job = tx + .get_collection_job::<0, LeaderSelected, dummy::Vdaf>( + &vdaf, + task.id(), + &collection_job_id, + ) + .await + .unwrap() + .unwrap(); + let batch_id = *collection_job.batch_identifier(); + + let encrypted_helper_aggregate_share = hpke::seal( + task.collector_hpke_keypair().config(), + &HpkeApplicationInfo::new( + &Label::AggregateShare, + &Role::Helper, + &Role::Collector, + ), + &helper_aggregate_share_bytes, + &AggregateShareAad::new(*task.id(), task.task_configuration(), request) + .get_encoded() + .unwrap(), + ) + .unwrap(); + + tx.update_collection_job::<0, LeaderSelected, dummy::Vdaf>( + &collection_job.with_state(CollectionJobState::Finished { + report_count: task.min_batch_size() + 1, + client_timestamp_interval: spanned_interval, + encrypted_helper_aggregate_share, + leader_aggregate_share, + }), + ) + .await + .unwrap(); + + Ok(batch_id) + }) + }) + .await + .unwrap(); + + if batch_id.eq(&batch_id_1) { + saw_batch_id_1 = true; + } else if batch_id.eq(&batch_id_2) { + saw_batch_id_2 = true; + } else { + panic!("unexpected batch ID"); + } + + let mut response = self.get_collection_job(&collection_job_id).await; + assert_eq!( + response.headers().get("content-type").unwrap(), + CollectionJobResp::MEDIA_TYPE, + ); + + let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; + let ( + report_count, + interval, + leader_encrypted_aggregate_share, + helper_encrypted_aggregate_share, + ) = assert_matches!( + collect_resp.clone(), + CollectionJobResp { + report_count, + interval, + leader_encrypted_agg_share, + helper_encrypted_agg_share, + .. + } => ( + report_count, + interval, + leader_encrypted_agg_share, + helper_encrypted_agg_share + ) + ); + + assert_eq!(report_count, self.task.min_batch_size() + 1); + assert_eq!(interval, spanned_interval); + + let decrypted_leader_aggregate_share = hpke::open( + self.task.collector_hpke_keypair(), + &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), + &leader_encrypted_aggregate_share, + &AggregateShareAad::new( + *self.task.id(), + self.task.task_configuration(), + request.clone(), + ) + .get_encoded() + .unwrap(), + ) + .unwrap(); + assert_eq!( + leader_aggregate_share, + dummy::AggregateShare::get_decoded(decrypted_leader_aggregate_share.as_ref()) + .unwrap() + ); + + let decrypted_helper_aggregate_share = hpke::open( + self.task.collector_hpke_keypair(), + &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), + &helper_encrypted_aggregate_share, + &AggregateShareAad::new( + *self.task.id(), + self.task.task_configuration(), + request.clone(), + ) + .get_encoded() + .unwrap(), + ) + .unwrap(); + assert_eq!( + helper_aggregate_share, + dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()) + .unwrap() + ); + + collection_job_resps.push(collect_resp); + } + + assert!(saw_batch_id_1 && saw_batch_id_2); + // Collection job IDs should be distinct + assert_eq!(collection_job_ids.len(), 2); + assert_ne!(collection_job_ids[0], collection_job_ids[1]); + + (requests[0].clone(), collection_job_ids[0]) } +} + +/// Parse task ID and collection job ID out of response location header. +/// +/// Panics if header is missing from response or value is malformed. +pub(super) fn ids_from_collection_job_response( + response: &http::Response, +) -> (TaskId, CollectionJobId) { + response + .headers() + .ids_from_collection_job_location() + .unwrap() + .unwrap() +} - assert!(saw_batch_id_1 && saw_batch_id_2); - // Collection job IDs should be distinct - assert_eq!(collection_job_ids.len(), 2); - assert_ne!(collection_job_ids[0], collection_job_ids[1]); +#[tokio::test] +async fn collection_job_success_leader_selected() { + // This test drives two current batch collection jobs to completion, verifying that distinct + // collection job IDs are minted and distinct batch IDs are collected each time. Then, we + // attempt to collect another current batch, which must fail as there are no more outstanding + // batches. + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; // We have run the two ready batches to completion. Further attempts to collect current batch // ought to fail. let mut response = test_case .post_collection_job(&CollectionJobReq::new( Query::new_leader_selected(random()), - aggregation_param.get_encoded().unwrap(), + dummy::AggregationParam::default().get_encoded().unwrap(), )) .await; assert_eq!(response.status(), StatusCode::BAD_REQUEST); @@ -529,14 +766,10 @@ async fn collection_job_success_leader_selected() { // Re-send the very first collection job creation request. We should get the representation of // the completed collection job without needing a further GET. - let request = CollectionJobReq::new( - Query::new_leader_selected(first_idempotency_key), - aggregation_param.get_encoded().unwrap(), - ); let mut response = test_case.post_collection_job(&request).await; assert_eq!(response.status(), StatusCode::OK); - let (_, collection_job_id) = ids_from_collection_job_response(&response); - assert_eq!(collection_job_id, collection_job_ids[0]); + let (_, collection_job_id_again) = ids_from_collection_job_response(&response); + assert_eq!(collection_job_id, collection_job_id_again); assert_eq!( response.headers().get("content-type").unwrap(), CollectionJobResp::MEDIA_TYPE, @@ -546,7 +779,7 @@ async fn collection_job_success_leader_selected() { #[tokio::test] async fn collection_job_post_idempotence_time_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // Seed DB with reports, aggregations, etc. corresponding to the two time intervals used later // in the test test_case @@ -637,7 +870,7 @@ async fn collection_job_post_idempotence_time_interval() { async fn collection_job_post_idempotence_time_interval_mutate_aggregation_param() { // Create a collection job for some time interval, then try to create another, varying the // aggregation parameter. It should fail due to batch overlap. - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_time_precision_units(0)) .await; @@ -670,8 +903,14 @@ async fn collection_job_post_idempotence_time_interval_mutate_aggregation_param( #[tokio::test] async fn collection_job_post_idempotence_leader_selected() { - let (test_case, batch_id_1, batch_id_2, _) = - setup_leader_selected_current_batch_collection_job_test_case().await; + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, _) = test_case.setup_leader_selected_current_batch().await; let mut seen_batch_id = None; let mut collection_job_id = None; @@ -769,9 +1008,16 @@ async fn collection_job_post_idempotence_leader_selected() { #[tokio::test] async fn collection_job_post_idempotence_leader_selected_mutate_aggregation_param() { + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + test_case.setup_leader_selected_current_batch().await; // Create a collection job for some idempotency key, then try to create another, varying the // aggregation parameter. It should fail because the batch is marked as aggregated. - let (test_case, _, _, _) = setup_leader_selected_current_batch_collection_job_test_case().await; let idempotency_key: LeaderSelectedQueryConfig = random(); @@ -803,8 +1049,15 @@ async fn collection_job_post_idempotence_leader_selected_mutate_aggregation_para #[tokio::test] async fn collection_job_post_idempotence_leader_selected_no_extra_reports() { - let (test_case, _batch_id_1, _batch_id_2, _) = - setup_leader_selected_current_batch_collection_job_test_case().await; + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + + test_case.setup_leader_selected_current_batch().await; let request_1 = CollectionJobReq::new( Query::new_leader_selected(random()), @@ -843,7 +1096,7 @@ async fn collection_job_post_idempotence_leader_selected_no_extra_reports() { #[tokio::test] async fn collection_job_batch_mode_misaligned() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, diff --git a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs index 9c1f9fde9..5c20fd3f2 100644 --- a/aggregator/src/aggregator/http_handlers/tests/collection_job.rs +++ b/aggregator/src/aggregator/http_handlers/tests/collection_job.rs @@ -1,39 +1,29 @@ -use assert_matches::assert_matches; use axum::body::Body; +use chrono::TimeDelta; use http::{Request, StatusCode}; -use janus_aggregator_core::{ - batch_mode::AccumulableBatchMode, - datastore::models::{CollectionJob, CollectionJobState}, - task::{AggregationMode, BatchMode, test_util::TaskBuilder}, -}; +use janus_aggregator_core::task::{AggregationMode, BatchMode, test_util::TaskBuilder}; use janus_core::{ auth_tokens::test_util::WithAuthenticationToken, - hpke::{self, HpkeApplicationInfo, Label}, + time::{Clock, DateTimeExt, TimeExt}, vdaf::VdafInstance, }; use janus_messages::{ - AggregateShareAad, CollectionJobExtension, CollectionJobExtensionType, CollectionJobId, - CollectionJobReq, CollectionJobResp, Duration, Interval, MediaType, Query, Role, Time, - batch_mode::TimeInterval, -}; -use prio::{ - codec::{Decode, Encode}, - vdaf::dummy, + CollectionJobExtension, CollectionJobExtensionType, CollectionJobId, CollectionJobReq, + CollectionJobResp, Duration, Interval, MediaType, Query, Role, Time, batch_mode::TimeInterval, }; +use prio::{codec::Encode, vdaf::dummy}; use rand::random; use serde_json::json; use tower::ServiceExt; use crate::aggregator::{ - collection_job_tests::{ids_from_collection_job_response, setup_collection_job_test_case}, - http_handlers::test_util::{ - HttpHandlerTest, decode_response_body, take_problem_details, take_response_body, - }, + collection_job_tests::{CollectionJobTestCase, ids_from_collection_job_response}, + http_handlers::test_util::{HttpHandlerTest, decode_response_body, take_problem_details}, }; #[tokio::test] async fn collection_job_post_request_to_helper() { - let test_case = setup_collection_job_test_case(Role::Helper, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Helper, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -70,7 +60,7 @@ async fn collection_job_post_request_to_helper() { #[tokio::test] async fn collection_job_post_request_invalid_batch_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -109,7 +99,7 @@ async fn collection_job_post_request_invalid_batch_interval() { #[tokio::test] async fn collection_job_post_request_invalid_aggregation_parameter() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -146,7 +136,7 @@ async fn collection_job_post_request_invalid_aggregation_parameter() { #[tokio::test] async fn collection_job_post_request_unsupported_extension() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // No collection job extension types are defined, so any extension is unsupported. let request = CollectionJobReq::new( @@ -181,7 +171,7 @@ async fn collection_job_post_request_unsupported_extension() { #[tokio::test] async fn collection_job_post_request_invalid_extension_order() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; // Extension types 0x0002 then 0x0001 are not in strictly increasing order. let request = CollectionJobReq::new( @@ -271,7 +261,7 @@ async fn collection_job_post_request_invalid_batch_size() { #[tokio::test] async fn collection_job_post_request_unauthenticated() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -310,7 +300,7 @@ async fn collection_job_post_request_unauthenticated() { #[tokio::test] async fn collection_job_get_request_unauthenticated_collection_jobs() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -358,188 +348,14 @@ async fn collection_job_get_request_unauthenticated_collection_jobs() { #[tokio::test] async fn collection_job_success_time_interval() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; - test_case - .setup_time_interval_batch(Time::from_seconds_since_epoch( - 0, - test_case.task.time_precision(), - )) - .await; - - let batch_interval = - TimeInterval::to_batch_identifier(&(), &Time::from_time_precision_units(0)).unwrap(); - - let aggregation_param = dummy::AggregationParam::default(); - let leader_aggregate_share = dummy::AggregateShare(0); - let helper_aggregate_share = dummy::AggregateShare(1); - - let request = CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ); - - let mut response = test_case.post_collection_job(&request).await; - assert_eq!(response.status(), StatusCode::OK); - let (task_id, collection_job_id) = ids_from_collection_job_response(&response); - assert_eq!(&task_id, test_case.task.id()); - - let got_collection_job = test_case - .datastore - .run_unnamed_tx(|tx| { - let task_id = *test_case.task.id(); - - Box::pin(async move { - Ok(tx - .get_collection_job(&dummy::Vdaf::new(1), &task_id, &collection_job_id) - .await - .unwrap() - .unwrap()) - }) - }) - .await - .unwrap(); - - let want_collection_job = CollectionJob::<0, TimeInterval, dummy::Vdaf>::new( - *test_case.task.id(), - collection_job_id, - // This is chosen by the leader, we have to copy it - *got_collection_job.aggregate_share_id(), - Query::new_time_interval(batch_interval), - aggregation_param, - batch_interval, - CollectionJobState::Start, - ); - - assert_eq!(want_collection_job, got_collection_job); - assert!(take_response_body(&mut response).await.is_empty()); - - let mut response = test_case.get_collection_job(&collection_job_id).await; - assert_eq!(response.status(), StatusCode::OK); - assert!(response.headers().get(http::header::CONTENT_TYPE).is_none()); - assert!(take_response_body(&mut response).await.is_empty()); - - // Update the collection job with the aggregate shares and some aggregation jobs. collection - // job should now be complete. - test_case - .datastore - .run_unnamed_tx(|tx| { - let task = test_case.task.clone(); - let helper_aggregate_share_bytes = helper_aggregate_share.get_encoded().unwrap(); - Box::pin(async move { - let encrypted_helper_aggregate_share = hpke::seal( - task.collector_hpke_keypair().config(), - &HpkeApplicationInfo::new( - &Label::AggregateShare, - &Role::Helper, - &Role::Collector, - ), - &helper_aggregate_share_bytes, - &AggregateShareAad::new( - *task.id(), - task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), - ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - - let collection_job = tx - .get_collection_job::<0, TimeInterval, dummy::Vdaf>( - &dummy::Vdaf::new(1), - task.id(), - &collection_job_id, - ) - .await - .unwrap() - .unwrap() - .with_state(CollectionJobState::Finished { - report_count: 12, - client_timestamp_interval: batch_interval, - encrypted_helper_aggregate_share, - leader_aggregate_share, - }); - - tx.update_collection_job::<0, TimeInterval, dummy::Vdaf>(&collection_job) - .await - .unwrap(); - Ok(()) - }) - }) - .await - .unwrap(); + let mut test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; - let mut response = test_case.get_collection_job(&collection_job_id).await; - - assert_eq!(response.status(), StatusCode::OK); - assert_eq!( - response.headers().get("content-type").unwrap(), - CollectionJobResp::MEDIA_TYPE, - ); - let collect_resp: CollectionJobResp = decode_response_body(&mut response).await; - let ( - report_count, - interval, - leader_encrypted_aggregate_share, - helper_encrypted_aggregate_share, - ) = assert_matches!( - collect_resp.clone(), - CollectionJobResp{ - report_count, - interval, - leader_encrypted_agg_share, - helper_encrypted_agg_share, - .. - } => (report_count, interval, leader_encrypted_agg_share, helper_encrypted_agg_share) - ); - - assert_eq!(report_count, 12); - assert_eq!(interval, batch_interval); - - let decrypted_leader_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), - &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Leader, &Role::Collector), - &leader_encrypted_aggregate_share, - &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), - ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - assert_eq!( - leader_aggregate_share, - dummy::AggregateShare::get_decoded(decrypted_leader_aggregate_share.as_ref()).unwrap() - ); - - let decrypted_helper_aggregate_share = hpke::open( - test_case.task.collector_hpke_keypair(), - &HpkeApplicationInfo::new(&Label::AggregateShare, &Role::Helper, &Role::Collector), - &helper_encrypted_aggregate_share, - &AggregateShareAad::new( - *test_case.task.id(), - test_case.task.task_configuration(), - CollectionJobReq::new( - Query::new_time_interval(batch_interval), - aggregation_param.get_encoded().unwrap(), - ), + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval( + Time::from_time_precision_units(10), + Interval::minimal(Time::from_time_precision_units(10)).unwrap(), ) - .get_encoded() - .unwrap(), - ) - .unwrap(); - assert_eq!( - helper_aggregate_share, - dummy::AggregateShare::get_decoded(decrypted_helper_aggregate_share.as_ref()).unwrap() - ); + .await; // Re-send the very first collection job creation request. We should get the representation of // the completed collection job without needing a further GET. @@ -556,7 +372,7 @@ async fn collection_job_success_time_interval() { #[tokio::test] async fn collection_job_get_request_no_such_collection_job() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -574,7 +390,7 @@ async fn collection_job_get_request_no_such_collection_job() { #[tokio::test] async fn collection_job_post_request_batch_queried_multiple_times() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -612,7 +428,7 @@ async fn collection_job_post_request_batch_queried_multiple_times() { #[tokio::test] async fn collection_job_post_request_batch_overlap() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -657,7 +473,7 @@ async fn collection_job_post_request_batch_overlap() { #[tokio::test] async fn delete_collection_job() { - let test_case = setup_collection_job_test_case(Role::Leader, BatchMode::TimeInterval).await; + let test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; let batch_interval = test_case .setup_time_interval_batch(Time::from_seconds_since_epoch( 0, @@ -666,19 +482,7 @@ async fn delete_collection_job() { .await; // Try to delete a collection job that doesn't exist - let response = test_case - .router - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri(test_case.task.collection_job_uri(&random()).unwrap().path()) - .with_authentication_token(test_case.task.collector_auth_token()) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); + let response = test_case.delete_collection_job(&random()).await; assert_eq!(response.status(), StatusCode::NOT_FOUND); // Create a collection job @@ -692,28 +496,245 @@ async fn delete_collection_job() { let (_, collection_job_id) = ids_from_collection_job_response(&response); // Cancel the job - let response = test_case - .router - .clone() - .oneshot( - Request::builder() - .method("DELETE") - .uri( - test_case - .task - .collection_job_uri(&collection_job_id) - .unwrap() - .path(), - ) - .with_authentication_token(test_case.task.collector_auth_token()) - .body(Body::empty()) - .unwrap(), - ) - .await - .unwrap(); + let response = test_case.delete_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::NO_CONTENT); // Get the job again let response = test_case.get_collection_job(&collection_job_id).await; assert_eq!(response.status(), StatusCode::NO_CONTENT); } + +#[tokio::test] +async fn recreate_deleted_collection_job_time_interval() { + let mut test_case = CollectionJobTestCase::setup(Role::Leader, BatchMode::TimeInterval).await; + + // Run a collection to completion, then delete the job. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval( + Time::from_time_precision_units(10), + Interval::minimal(Time::from_time_precision_units(10)).unwrap(), + ) + .await; + let response = test_case.delete_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job should yield NO CONTENT again. + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job with a mutated request should fail. + let mutant_request = CollectionJobReq::new( + request.query().clone(), + dummy::AggregationParam(17).get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&mutant_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_deleted_collection_job_leader_selected() { + let mut test_case = CollectionJobTestCase::setup( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + ) + .await; + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + + // Run a collection to completion, then delete the job. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; + + // Put another outstanding batch in the datastore to ensure subsequent job creation requests + // aren't failing just because there's no reports to aggregate. + test_case + .setup_leader_selected_batch( + test_case + .clock() + .now() + .to_time(test_case.task.time_precision()), + test_case.task.min_batch_size() + 1, + ) + .await; + + let response = test_case.delete_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job with an identical request should fail. + let response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::NO_CONTENT); + + // Re-creating the job with a mutated request should fail. + let mutant_request = CollectionJobReq::new( + request.query().clone(), + dummy::AggregationParam(17).get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&mutant_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_collection_job_time_interval() { + let expiry = Duration::from_time_precision_units(20); + let mut test_case = CollectionJobTestCase::setup_with_report_expiry_age( + Role::Leader, + BatchMode::TimeInterval, + Some(expiry), + ) + .await; + + let now = Time::from_seconds_since_epoch( + test_case.clock().now().as_seconds_since_epoch(), + test_case.task.time_precision(), + ); + + let report_time = now + .sub_duration(&Duration::from_time_precision_units(3)) + .unwrap(); + let queried_interval = Interval::new( + now.sub_duration(&Duration::from_time_precision_units(10)) + .unwrap(), + Duration::from_time_precision_units(10), + ) + .unwrap(); + + // Run a collection to completion. + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_time_interval(report_time, queried_interval) + .await; + + // Advance clock far enough for the collection job to expire. + test_case.clock().advance(TimeDelta::seconds( + expiry.as_seconds(test_case.task.time_precision()) as i64, + )); + let response = test_case.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + // Re-creating the job with an identical request should fail. + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); + + // Re-creating the job with a mutated request should fail. + let mutant_request = CollectionJobReq::new( + request.query().clone(), + dummy::AggregationParam(17).get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&mutant_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +} + +#[tokio::test] +async fn recreate_expired_collection_job_leader_selected() { + let expiry = Duration::from_time_precision_units(20); + let mut test_case = CollectionJobTestCase::setup_with_report_expiry_age( + Role::Leader, + BatchMode::LeaderSelected { + batch_time_window_size: None, + }, + Some(expiry), + ) + .await; + + let (batch_id_1, batch_id_2, spanned_interval) = + test_case.setup_leader_selected_current_batch().await; + + let (request, collection_job_id) = test_case + .simulate_successful_collection_job_leader_selected( + batch_id_1, + batch_id_2, + spanned_interval, + ) + .await; + + // Advance clock far enough for the collection job to expire. + test_case.clock().advance(TimeDelta::seconds( + 2 * expiry.as_seconds(test_case.task.time_precision()) as i64, + )); + let response = test_case.get_collection_job(&collection_job_id).await; + assert_eq!(response.status(), StatusCode::NOT_FOUND); + + // Put another outstanding batch in the datastore to ensure subsequent job creation requests + // aren't failing just because there's no reports to aggregate. + test_case + .setup_leader_selected_batch( + test_case + .clock() + .now() + .to_time(test_case.task.time_precision()), + test_case.task.min_batch_size() + 1, + ) + .await; + + // Re-creating the job should fail. + let mut response = test_case.post_collection_job(&request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); + + // Re-creating the job with a mutated request should fail. + let mutant_request = CollectionJobReq::new( + request.query().clone(), + dummy::AggregationParam(17).get_encoded().unwrap(), + ); + let mut response = test_case.post_collection_job(&mutant_request).await; + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!( + take_problem_details(&mut response).await, + json!({ + "status": StatusCode::BAD_REQUEST.as_u16(), + "type": "urn:ietf:params:ppm:dap:error:batchOverlap", + "title": "The queried batch overlaps with a previously queried batch.", + "taskid": format!("{}", test_case.task.id()), + }) + ); +}