diff --git a/ipa-core/src/bin/report_collector.rs b/ipa-core/src/bin/report_collector.rs index b80137047..fb41fcb9a 100644 --- a/ipa-core/src/bin/report_collector.rs +++ b/ipa-core/src/bin/report_collector.rs @@ -151,6 +151,11 @@ enum ReportCollectorCommand { /// Number of records to aggregate #[clap(long, short = 'n')] count: u32, + + // If set, use the specified fixed polling interval when running a query. + // Otherwise, use exponential backoff. + #[clap(long)] + set_fixed_polling_ms: Option, }, } @@ -264,6 +269,7 @@ async fn main() -> Result<(), Box> { ref encrypted_inputs, hybrid_query_config, count, + set_fixed_polling_ms, } => { hybrid( &args, @@ -271,6 +277,7 @@ async fn main() -> Result<(), Box> { clients, encrypted_inputs, count.try_into().expect("u32 should fit into usize"), + set_fixed_polling_ms, ) .await? } @@ -421,6 +428,7 @@ async fn hybrid( helper_clients: Vec<[IpaHttpClient; 3]>, encrypted_inputs: &EncryptedInputs, count: usize, + set_fixed_polling_ms: Option, ) -> Result<(), Box> { let query_type = QueryType::MaliciousHybrid(hybrid_query_config); @@ -471,6 +479,7 @@ async fn hybrid( helper_clients, query_id, hybrid_query_config, + set_fixed_polling_ms, ) .await; diff --git a/ipa-core/src/cli/playbook/hybrid.rs b/ipa-core/src/cli/playbook/hybrid.rs index c52547f6e..92d74b383 100644 --- a/ipa-core/src/cli/playbook/hybrid.rs +++ b/ipa-core/src/cli/playbook/hybrid.rs @@ -31,6 +31,7 @@ pub async fn run_hybrid_query_and_validate( clients: Vec<[IpaHttpClient; 3]>, query_id: QueryId, query_config: HybridQueryParams, + set_fixed_polling_ms: Option, ) -> HybridQueryResult where HV: SharedValue + U128Conversions, @@ -56,7 +57,11 @@ where let leader_clients = &clients[0]; - let mut delay = Duration::from_millis(125); + let (exponential_backoff, mut delay) = match set_fixed_polling_ms { + Some(specified_delay) => (false, Duration::from_millis(specified_delay)), + None => (true, Duration::from_millis(125)), + }; + loop { if try_join_all( leader_clients @@ -72,7 +77,9 @@ where } sleep(delay).await; - delay = min(Duration::from_secs(5), delay * 2); + if exponential_backoff { + delay = min(Duration::from_secs(5), delay * 2); + } // TODO: Add a timeout of some sort. Possibly, add some sort of progress indicator to // the status API so we can check whether the query is making progress. }