From f3951dd7ddb007c1680796623bfe4fb02a59e70a Mon Sep 17 00:00:00 2001 From: Tim Geoghegan Date: Fri, 9 Oct 2026 16:33:38 -0400 Subject: [PATCH 1/2] 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 2/2] 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| {