Repository: cloudflare/daphne
Pull request #614 (2024-06-07 15:32:57)

Allow costumization of load test request patterns


Patch:
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index 74b51d3..5bb52fb 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -39,7 +39,7 @@ use daphne::{
     EarlyReportStateInitialized,
 };
 use daphne_service_utils::http_headers;
-use futures::{StreamExt, TryStreamExt};
+use futures::{future::OptionFuture, StreamExt, TryStreamExt};
 use prio::codec::{Decode, ParameterizedEncode};
 use prometheus::{Encoder, HistogramVec, IntCounterVec, TextEncoder};
 use rand::{rngs, Rng};
@@ -50,12 +50,93 @@ use std::{
     env,
     ops::Range,
     path::PathBuf,
+    sync::atomic::{AtomicUsize, Ordering},
     time::{Duration, Instant, SystemTime},
 };
-use tokio::sync::Barrier;
+use tokio::sync::{Barrier, Semaphore};
 use tracing::{info, instrument};
 use url::Url;
 
+/// Stride controls the rate at which requests are made during an aggregation job.
+///
+/// For example, for `len` = 10 and `wait_time` = 1s, every one second 10 requests will be made
+/// simultaneously.
+#[derive(Clone, Copy, Debug)]
+pub struct LoadControlStride {
+    pub len: usize,
+    pub wait_time: Duration,
+}
+
+#[derive(Clone, Copy, Debug)]
+pub struct LoadControlParams {
+    max_concurrent_requests: Option<usize>,
+    min_requests_before_starting: usize,
+    stride: Option<LoadControlStride>,
+}
+
+impl LoadControlParams {
+    pub fn new<M, B>(
+        max_concurrent_requests: M,
+        min_requests_before_starting: usize,
+        stride: B,
+    ) -> Self
+    where
+        M: Into<Option<usize>>,
+        B: Into<Option<LoadControlStride>>,
+    {
+        Self {
+            max_concurrent_requests: max_concurrent_requests.into(),
+            min_requests_before_starting,
+            stride: stride.into(),
+        }
+    }
+
+    pub fn max_requests_before_starting(
+        reports_per_batch: usize,
+        reports_per_agg_job: usize,
+    ) -> usize {
+        // total reports / job + (0 if reports divide by reports_per_agg_job else 1)
+        reports_per_batch / reports_per_agg_job
+            + usize::from(reports_per_batch % reports_per_agg_job != 0)
+    }
+}
+
+struct LoadControl {
+    max_concurrent_requests: Option<Semaphore>,
+    min_requests_before_starting: Barrier,
+    stride: Option<LoadControlStride>,
+    current_batch_len: AtomicUsize,
+}
+
+impl From<LoadControlParams> for LoadControl {
+    fn from(params: LoadControlParams) -> Self {
+        Self {
+            max_concurrent_requests: params.max_concurrent_requests.map(Semaphore::new),
+            min_requests_before_starting: Barrier::new(params.min_requests_before_starting),
+            stride: params.stride,
+            current_batch_len: AtomicUsize::default(),
+        }
+    }
+}
+
+impl LoadControl {
+    #[must_use]
+    pub async fn wait(&self) -> Option<tokio::sync::SemaphorePermit> {
+        let id = self.current_batch_len.fetch_add(1, Ordering::SeqCst);
+        self.min_requests_before_starting.wait().await;
+        if let Some(LoadControlStride { len, wait_time }) = self.stride {
+            let batch_number = u32::try_from(id.checked_div(len).unwrap()).unwrap();
+            tokio::time::sleep(wait_time * batch_number).await;
+        }
+        OptionFuture::from(
+            self.max_concurrent_requests
+                .as_ref()
+                .map(|s| async { s.acquire().await.unwrap() }),
+        )
+        .await
+    }
+}
+
 struct TestMetrics {
     aggregated: IntCounterVec,
     test_durations: HistogramVec,
@@ -129,6 +210,7 @@ pub struct Test {
     /// The path to the hpke signing certificate, which can be used to verify the hpke config
     /// signature.
     hpke_signing_certificate_path: Option<PathBuf>,
+    load_control: LoadControlParams,
 }
 
 pub struct TestOptions {
@@ -176,6 +258,7 @@ impl Test {
         helper_url: Url,
         vdaf_verify_init: &str,
         vdaf_config: VdafConfig,
+        load_control: LoadControlParams,
     ) -> Result<Self> {
         let vdaf_verify_init = <[u8; 32]>::try_from(
             hex::decode(vdaf_verify_init)
@@ -218,6 +301,7 @@ impl Test {
             metrics,
             vdaf_config,
             hpke_signing_certificate_path: None,
+            load_control,
         })
     }
 
@@ -226,6 +310,7 @@ impl Test {
         vdaf_config: VdafConfig,
         hpke_signing_certificate_path: Option<PathBuf>,
         re_use_http_client: bool,
+        load_control: LoadControlParams,
     ) -> Result<Self> {
         const LEADER_BEARER_TOKEN_VAR: &str = "LEADER_BEARER_TOKEN";
         const LEADER_TLS_CLIENT_CERT_VAR: &str = "LEADER_TLS_CLIENT_CERT";
@@ -267,6 +352,7 @@ impl Test {
             helper_url,
             &vdaf_verify_init,
             vdaf_config,
+            load_control,
         )?;
         if let Some(token) = leader_bearer_token {
             test = test.with_bearer_token(token);
@@ -406,9 +492,7 @@ impl Test {
     where
         F: Fn(usize) -> ReportGenerator,
     {
-        let job_count = report_count / reports_per_agg_job
-            + usize::from(report_count % reports_per_agg_job != 0);
-        let barrier = Barrier::new(job_count);
+        let load_control = LoadControl::from(self.load_control);
         let (count, out_shares_for_batch, agg_job_durations) = futures::stream::iter(
             distribute_reports_in_chunks(report_count, reports_per_agg_job),
         )
@@ -419,7 +503,7 @@ impl Test {
                 agg_job_index,
                 report_generator(reports_for_agg_job),
                 part_batch_sel,
-                &barrier,
+                &load_control,
             )
         })
         .buffer_unordered((report_count / reports_per_agg_job) + 1)
@@ -450,7 +534,7 @@ impl Test {
         // agg_job_index is kept
         reports_for_agg_job,
         part_batch_sel,
-        barrier,
+        load_control,
     ))]
     pub async fn run_agg_job(
         &self,
@@ -458,7 +542,7 @@ impl Test {
         agg_job_index: usize,
         reports_for_agg_job: ReportGenerator,
         part_batch_sel: &PartialBatchSelector,
-        barrier: &Barrier,
+        load_control: &LoadControl,
     ) -> anyhow::Result<(DapAggregateSpan<DapAggregateShare>, Duration)> {
         let report_count = reports_for_agg_job.len();
         info!(report_count, "Starting aggregation job");
@@ -500,7 +584,8 @@ impl Test {
         ))?;
 
         // wait for all agg jobs to be ready to fire.
-        barrier.wait().await;
+        info!("Reports generated, waiting for other tasks...");
+        let _guard = load_control.wait().await;
         info!("Starting AggregationJobInitReq");
         let start = Instant::now();
         let resp = send(
diff --git a/crates/dapf/src/main.rs b/crates/dapf/src/main.rs
index ac883f2..e08e78f 100644
--- a/crates/dapf/src/main.rs
+++ b/crates/dapf/src/main.rs
@@ -4,7 +4,7 @@
 use anyhow::{anyhow, Context, Result};
 use clap::{builder::PossibleValue, Parser, Subcommand, ValueEnum};
 use dapf::{
-    acceptance::{load_testing, TestOptions},
+    acceptance::{load_testing, LoadControlParams, LoadControlStride, TestOptions},
     deduce_dap_version_from_url, HttpClientExt,
 };
 use daphne::{
@@ -24,7 +24,7 @@ use std::{
     path::PathBuf,
     process::Command,
     str::FromStr,
-    time::SystemTime,
+    time::{Duration, SystemTime},
 };
 use tracing::level_filters::LevelFilter;
 use tracing_subscriber::EnvFilter;
@@ -236,6 +236,8 @@ enum Action {
         reports_per_batch: usize,
         #[arg(env, default_value_t = 10)]
         reports_per_agg_job: usize,
+        #[command(flatten)]
+        load_control: LoadControlParamsCli,
     },
     /// Perform multiple aggregation jobs against a helper.
     ///
@@ -253,6 +255,8 @@ enum Action {
         helper_url: Url,
         #[arg(long, env)]
         no_reuse_http_client: bool,
+        #[command(flatten)]
+        load_control: LoadControlParamsCli,
         #[command(subcommand)]
         params: LoadTestParameters,
     },
@@ -273,6 +277,47 @@ enum LoadTestParameters {
     },
 }
 
+#[derive(Debug, Clone, Parser)]
+struct LoadControlParamsCli {
+    #[arg(long, env)]
+    stride_wait_time: Option<u64>,
+    #[arg(long, env)]
+    stride_len: Option<usize>,
+    #[arg(long, env)]
+    min_requests_before_starting: Option<usize>,
+    #[arg(long, env)]
+    max_concurrent_requests: Option<usize>,
+}
+
+impl LoadControlParamsCli {
+    fn into_params(
+        self,
+        reports_per_agg_job: usize,
+        reports_per_batch: usize,
+    ) -> LoadControlParams {
+        LoadControlParams::new(
+            self.max_concurrent_requests,
+            self.min_requests_before_starting.unwrap_or_else(|| {
+                LoadControlParams::max_requests_before_starting(
+                    reports_per_batch,
+                    reports_per_agg_job,
+                )
+            }),
+            match (self.stride_len, self.stride_wait_time) {
+                (Some(len), Some(wait_time)) => Some(LoadControlStride {
+                    len,
+                    wait_time: Duration::from_millis(wait_time),
+                }),
+                (Some(len), None) => Some(LoadControlStride {
+                    len,
+                    wait_time: Duration::from_secs(1),
+                }),
+                (None, _) => None,
+            },
+        )
+    }
+}
+
 #[derive(Clone, Debug)]
 struct KemAlg(HpkeKemId);
 
@@ -631,14 +676,18 @@ async fn main() -> Result<()> {
             hpke_signing_certificate_path,
             reports_per_batch,
             reports_per_agg_job,
+            load_control,
         } => {
             tracing::info!("using vdaf: {:?}", vdaf_config);
+            let load_control = load_control.into_params(reports_per_agg_job, reports_per_batch);
+            tracing::info!("using load params: {load_control:?}");
 
             let t = dapf::acceptance::Test::from_env(
                 helper_url,
                 vdaf_config.into_vdaf(),
                 hpke_signing_certificate_path,
                 true,
+                load_control,
             )?;
 
             let res = t
@@ -655,14 +704,21 @@ async fn main() -> Result<()> {
         Action::LoadTest {
             helper_url,
             no_reuse_http_client,
+            load_control,
             params: LoadTestParameters::Multiple,
         } => {
-            load_testing::execute_multiple_combinations(helper_url, !no_reuse_http_client).await;
+            load_testing::execute_multiple_combinations(
+                helper_url,
+                !no_reuse_http_client,
+                load_control.into_params(usize::MAX, 1),
+            )
+            .await;
             Ok(())
         }
         Action::LoadTest {
             helper_url,
             no_reuse_http_client,
+            load_control,
             params:
                 LoadTestParameters::Single {
                     reports_per_batch,
@@ -676,6 +732,7 @@ async fn main() -> Result<()> {
                 reports_per_batch,
                 reports_per_agg_job,
                 !no_reuse_http_client,
+                load_control.into_params(reports_per_agg_job, reports_per_batch),
             )
             .await;
             Ok(())


Test patch:
diff --git a/crates/dapf/src/acceptance/load_testing.rs b/crates/dapf/src/acceptance/load_testing.rs
index 18ed260..e8816a9 100644
--- a/crates/dapf/src/acceptance/load_testing.rs
+++ b/crates/dapf/src/acceptance/load_testing.rs
@@ -23,6 +23,8 @@ use futures::StreamExt;
 use rand::{thread_rng, Rng};
 use url::Url;
 
+use super::LoadControlParams;
+
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
 struct MeasurementParameters {
     reports_per_batch: usize,
@@ -332,13 +334,18 @@ async fn execute(t: &Test, test_config: &TestOptions) -> TestResults {
     TestResults { tests }
 }
 
-pub async fn execute_multiple_combinations(helper_url: Url, re_use_http_client: bool) {
+pub async fn execute_multiple_combinations(
+    helper_url: Url,
+    re_use_http_client: bool,
+    load_control: LoadControlParams,
+) {
     // This vdaf config is later replaced on each of the runs
     let mut t = Test::from_env(
         helper_url,
         VdafConfig::Prio2 { dimension: 44 },
         None,
         re_use_http_client,
+        load_control,
     )
     .expect("env to be present");
 
@@ -350,6 +357,7 @@ pub async fn execute_multiple_combinations(helper_url: Url, re_use_http_client:
 
     let mut measurments = HashMap::new();
 
+    let min_requests_before_starting = t.load_control.min_requests_before_starting;
     loop {
         for params @ MeasurementParameters {
             reports_per_batch,
@@ -362,12 +370,17 @@ pub async fn execute_multiple_combinations(helper_url: Url, re_use_http_client:
             test_config.reports_per_agg_job = reports_per_agg_job;
             t.vdaf_config = vdaf_config;
             test_config.measurement = Some(t.gen_measurement().unwrap());
+            t.load_control.min_requests_before_starting = std::cmp::max(
+                min_requests_before_starting,
+                reports_per_batch / reports_per_agg_job,
+            );
 
             println!("===== Performance Report =====");
             println!("parameters:");
             println!("\t- reports_per_batch:   {reports_per_batch}");
             println!("\t- reports_per_agg_job: {reports_per_agg_job}");
             println!("\t- vdaf_config:         {vdaf_config:?}");
+            println!("\t- load_control:        {:?}", t.load_control);
 
             let results = execute(&t, &test_config).await;
 
@@ -389,11 +402,18 @@ pub async fn execute_single_combination_from_env(
     reports_per_batch: usize,
     reports_per_agg_job: usize,
     re_use_http_client: bool,
+    load_control: LoadControlParams,
 ) {
     const VERSION: daphne::DapVersion = daphne::DapVersion::Latest;
 
-    let t = Test::from_env(helper_url, vdaf_config, None, re_use_http_client)
-        .expect("env to be present");
+    let t = Test::from_env(
+        helper_url,
+        vdaf_config,
+        None,
+        re_use_http_client,
+        load_control,
+    )
+    .expect("env to be present");
 
     let system_now = now();
 
@@ -407,6 +427,7 @@ pub async fn execute_single_combination_from_env(
         println!("\t- reports_per_batch:   {reports_per_batch}");
         println!("\t- reports_per_agg_job: {reports_per_agg_job}");
         println!("\t- vdaf_config:         {:?}", t.vdaf_config);
+        println!("\t- load_control:        {load_control:?}");
 
         let batch_id = BatchId(thread_rng().gen());
         let part_batch_sel = PartialBatchSelector::FixedSizeByBatchId { batch_id };