From aba36f850b3fb39059212850eeeb821a8e951a39 Mon Sep 17 00:00:00 2001
From: Thomas James Yurek <tyurek@meta.com>
Date: Thu, 12 Dec 2024 11:18:32 -0800
Subject: [PATCH] init

---
 ipa-core/src/bin/report_collector.rs |  9 +++++++++
 ipa-core/src/cli/playbook/hybrid.rs  | 12 ++++++++++--
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/ipa-core/src/bin/report_collector.rs b/ipa-core/src/bin/report_collector.rs
index b80137047..64aee11f0 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.
+        #[arg(long, default_value_t = 0)]
+        set_fixed_polling_ms: u64,
     },
 }
 
@@ -264,6 +269,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
             ref encrypted_inputs,
             hybrid_query_config,
             count,
+            set_fixed_polling_ms,
         } => {
             hybrid(
                 &args,
@@ -271,6 +277,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
                 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<Helper>; 3]>,
     encrypted_inputs: &EncryptedInputs,
     count: usize,
+    set_fixed_polling_ms: u64,
 ) -> Result<(), Box<dyn Error>> {
     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..53a83b8fd 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<HV>(
     clients: Vec<[IpaHttpClient<Helper>; 3]>,
     query_id: QueryId,
     query_config: HybridQueryParams,
+    set_fixed_polling_ms: u64,
 ) -> HybridQueryResult
 where
     HV: SharedValue + U128Conversions,
@@ -55,8 +56,13 @@ where
     .unwrap();
 
     let leader_clients = &clients[0];
+    let exponential_backoff = set_fixed_polling_ms == 0;
+    let mut delay = if exponential_backoff {
+        Duration::from_millis(125)
+    } else {
+        Duration::from_millis(set_fixed_polling_ms)
+    };
 
-    let mut delay = Duration::from_millis(125);
     loop {
         if try_join_all(
             leader_clients
@@ -72,7 +78,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.
     }