Repository: cloudflare/daphne
Pull request #675 (2024-09-10 09:47:41)

Rewrite benchmarks and other small refactors
- Rewrite the aggregation benchmark
- Stop using a global variable to control replay protection in tests and dapf
- Make more fields in InMemoryAggregator private
- Fix dapf http client
- Add dhat profiling to helper
- Add benchmark results



Patch:
diff --git a/Cargo.lock b/Cargo.lock
index d4ffb51..d5a1c4c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -426,7 +426,7 @@ dependencies = [
  "bitflags 2.6.0",
  "cexpr",
  "clang-sys",
- "itertools 0.10.5",
+ "itertools 0.12.1",
  "lazy_static",
  "lazycell",
  "log",
@@ -962,6 +962,7 @@ dependencies = [
  "config",
  "daphne",
  "daphne-service-utils",
+ "dhat",
  "futures",
  "hex",
  "hpke-rs",
@@ -1133,6 +1134,22 @@ dependencies = [
  "powerfmt",
 ]
 
+[[package]]
+name = "dhat"
+version = "0.3.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98cd11d84628e233de0ce467de10b8633f4ddaecafadefc86e13b84b8739b827"
+dependencies = [
+ "backtrace",
+ "lazy_static",
+ "mintex",
+ "parking_lot",
+ "rustc-hash 1.1.0",
+ "serde",
+ "serde_json",
+ "thousands",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -2065,6 +2082,12 @@ dependencies = [
  "adler",
 ]
 
+[[package]]
+name = "mintex"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9bec4598fddb13cc7b528819e697852653252b760f1228b7642679bf2ff2cd07"
+
 [[package]]
 name = "mio"
 version = "1.0.2"
@@ -3424,6 +3447,15 @@ version = "1.3.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64"
 
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "signature"
 version = "2.2.0"
@@ -3618,6 +3650,12 @@ dependencies = [
  "syn 2.0.76",
 ]
 
+[[package]]
+name = "thousands"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "3bf63baf9f5039dadc247375c29eb13706706cfde997d0330d05aa63a77d8820"
+
 [[package]]
 name = "thread_local"
 version = "1.1.8"
@@ -3717,6 +3755,7 @@ dependencies = [
  "libc",
  "mio",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
  "tokio-macros",
  "windows-sys 0.52.0",
diff --git a/Cargo.toml b/Cargo.toml
index c98a71e..101c79e 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -47,6 +47,7 @@ config = "0.13.4"
 constcat = "0.5.0"
 criterion = { version = "0.5.1", features = ["async_tokio"] }
 deepsize = { version = "0.2.0" }
+dhat = "0.3.3"
 futures = "0.3.30"
 getrandom = "0.2.15"
 headers = "0.4"
diff --git a/crates/dapf/src/acceptance/mod.rs b/crates/dapf/src/acceptance/mod.rs
index 14b1dd8..e9e1a5b 100644
--- a/crates/dapf/src/acceptance/mod.rs
+++ b/crates/dapf/src/acceptance/mod.rs
@@ -38,7 +38,7 @@ use daphne::{
     vdaf::VdafConfig,
     DapAggregateShare, DapAggregateSpan, DapAggregationParam, DapError, DapMeasurement,
     DapQueryConfig, DapTaskConfig, DapTaskParameters, DapVersion, EarlyReportStateConsumed,
-    EarlyReportStateInitialized,
+    EarlyReportStateInitialized, ReplayProtection,
 };
 use daphne_service_utils::http_headers;
 use futures::{future::OptionFuture, StreamExt, TryStreamExt};
@@ -169,6 +169,9 @@ pub struct Test {
     /// signature.
     hpke_signing_certificate_path: Option<PathBuf>,
     load_control: LoadControlParams,
+    /// Replay reports when generating the test aggregation. This improves test speed but requires
+    /// the test target to have replay protection disabled.
+    replay_reports: bool,
 }
 
 pub struct TestOptions {
@@ -216,6 +219,7 @@ impl Test {
         vdaf_verify_init: &str,
         vdaf_config: VdafConfig,
         load_control: LoadControlParams,
+        replay_reports: bool,
     ) -> Result<Self> {
         let vdaf_verify_init = <[u8; 32]>::try_from(
             hex::decode(vdaf_verify_init)
@@ -258,6 +262,7 @@ impl Test {
             vdaf_config,
             hpke_signing_certificate_path: None,
             load_control,
+            replay_reports,
         })
     }
 
@@ -270,6 +275,7 @@ impl Test {
     ) -> Result<Self> {
         const LEADER_BEARER_TOKEN_VAR: &str = "LEADER_BEARER_TOKEN";
         const VDAF_VERIFY_INIT_VAR: &str = "VDAF_VERIFY_INIT";
+        let replay_reports = std::env::var("REPLAY_REPORTS").unwrap_or_default() == "1";
 
         let leader_bearer_token = env::var(LEADER_BEARER_TOKEN_VAR).ok();
 
@@ -282,6 +288,7 @@ impl Test {
             &vdaf_verify_init,
             vdaf_config,
             load_control,
+            replay_reports,
         )?;
         if let Some(token) = leader_bearer_token {
             test = test.with_bearer_token(token);
@@ -485,7 +492,7 @@ impl Test {
         let agg_job_id = AggregationJobId(rngs::OsRng.gen());
         let report_count = reports_for_agg_job.len();
         let (agg_job_state, agg_job_init_req) = task_config
-            .produce_agg_job_req(
+            .test_produce_agg_job_req(
                 fake_leader_hpke_receiver_config,
                 self,
                 task_id,
@@ -493,6 +500,11 @@ impl Test {
                 &DapAggregationParam::Empty,
                 reports_for_agg_job,
                 self.metrics(),
+                if self.replay_reports {
+                    ReplayProtection::InsecureDisabled
+                } else {
+                    ReplayProtection::Enabled
+                },
             )
             .await
             .context("producing agg job init request")?;
@@ -713,6 +725,7 @@ impl Test {
                         version,
                         now.0,
                         vec![messages::Extension::Taskprov],
+                        self.replay_reports,
                     )
                 },
             )
diff --git a/crates/dapf/src/http_client.rs b/crates/dapf/src/http_client.rs
index 184f427..10364f1 100644
--- a/crates/dapf/src/http_client.rs
+++ b/crates/dapf/src/http_client.rs
@@ -150,7 +150,7 @@ impl HttpClient {
     }
 
     pub fn put<U: IntoUrl>(&self, url: U) -> RequestBuilder {
-        self.client().post(url)
+        self.client().put(url)
     }
 
     pub async fn get_hpke_config(
diff --git a/crates/dapf/src/main.rs b/crates/dapf/src/main.rs
index 977fade..b41bff5 100644
--- a/crates/dapf/src/main.rs
+++ b/crates/dapf/src/main.rs
@@ -399,10 +399,6 @@ async fn main() -> Result<()> {
         HttpClient::new_no_reuse(cli.enable_ssl_key_log_file)?
     };
 
-    if std::env::var("REPLAY_REPORTS").unwrap_or_default() == "1" {
-        daphne::testing::report_generator::replay_reports(true);
-    }
-
     match cli.action {
         Action::Leader(leader) => handle_leader_actions(leader, http_client).await,
         Action::Hpke(hpke) => handle_hpke_actions(hpke, http_client).await,
diff --git a/crates/daphne-server/Cargo.toml b/crates/daphne-server/Cargo.toml
index 696fd99..efec364 100644
--- a/crates/daphne-server/Cargo.toml
+++ b/crates/daphne-server/Cargo.toml
@@ -44,11 +44,13 @@ clap.workspace = true
 config.workspace = true
 daphne = { path = "../daphne", features = ["test-utils"] }
 daphne-service-utils = { path = "../daphne-service-utils", features = ["prometheus"] }
+dhat.workspace = true
 hpke-rs.workspace = true
 paste.workspace = true
 prometheus.workspace = true
 rand.workspace = true
 rcgen.workspace = true
+tokio = { workspace = true, features = ["signal"] }
 tracing-subscriber = { workspace = true, features = ["env-filter"] }
 webpki.workspace = true
 x509-parser.workspace = true
diff --git a/crates/daphne-server/examples/service.rs b/crates/daphne-server/examples/service.rs
index 3a8d925..0c3c355 100644
--- a/crates/daphne-server/examples/service.rs
+++ b/crates/daphne-server/examples/service.rs
@@ -89,8 +89,13 @@ struct Args {
     storage_proxy: Option<Url>,
 }
 
+#[global_allocator]
+static ALLOC: dhat::Alloc = dhat::Alloc;
+
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
+    let _profiler = dhat::Profiler::new_heap();
+
     // Parse the configuration from the command line arguments.
     let config = Config::try_from(Args::parse())?;
     println!("starting service with config:\n{config:#?}");
@@ -112,12 +117,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error + Sync + Send>> {
         .init();
 
     // hand the router to axum for it to run
-    axum::Server::bind(&std::net::SocketAddr::new(
+    let serve = axum::Server::bind(&std::net::SocketAddr::new(
         "0.0.0.0".parse().unwrap(),
         config.port,
     ))
-    .serve(router.into_make_service())
-    .await?;
+    .serve(router.into_make_service());
+
+    let ctrl_c = tokio::signal::ctrl_c();
+
+    tokio::select! {
+        _ = serve => {}
+        _ = ctrl_c => {}
+    }
 
     Ok(())
 }
diff --git a/crates/daphne-server/src/roles/mod.rs b/crates/daphne-server/src/roles/mod.rs
index 5f84160..0d460be 100644
--- a/crates/daphne-server/src/roles/mod.rs
+++ b/crates/daphne-server/src/roles/mod.rs
@@ -26,7 +26,7 @@ pub async fn fetch_replay_protection_override(kv: Kv<'_>) -> ReplayProtection {
         .unwrap_or_default(); // treat missing as false
     if skip_replay_protection {
         tracing::debug!("replay protection is disabled");
-        ReplayProtection::Disabled
+        ReplayProtection::InsecureDisabled
     } else {
         ReplayProtection::Enabled
     }
diff --git a/crates/daphne/benches/aggregation.rs b/crates/daphne/benches/aggregation.rs
index 3098589..9c5b922 100644
--- a/crates/daphne/benches/aggregation.rs
+++ b/crates/daphne/benches/aggregation.rs
@@ -1,215 +1,130 @@
 // Copyright (c) 2024 Cloudflare, Inc. All rights reserved.
 // SPDX-License-Identifier: BSD-3-Clause
 
+#![allow(clippy::cast_possible_truncation)]
+
 use std::{
-    num::NonZeroUsize,
-    time::{Duration, Instant, SystemTime},
+    hint::black_box,
+    iter::repeat,
+    time::{Duration, Instant},
 };
 
-use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use criterion::{criterion_group, criterion_main, Bencher, BenchmarkId, Criterion, Throughput};
 use daphne::{
-    auth::BearerToken,
-    hpke::{HpkeAeadId, HpkeConfig, HpkeKdfId, HpkeKemId, HpkeReceiverConfig},
-    messages::{BatchId, TaskId},
-    roles::{helper, DapAggregator},
-    testing::{report_generator::ReportGenerator, InMemoryAggregator},
-    vdaf::{VdafConfig, VdafVerifyKey},
-    DapRequest, DapTaskConfig, DapVersion,
+    hpke::HpkeKemId,
+    messages::AggregationJobInitReq,
+    testing::AggregationJobTest,
+    vdaf::{Prio3Config, VdafConfig},
+    DapAggregationParam, DapVersion,
 };
-use hpke_rs::HpkePublicKey;
-use prio::codec::ParameterizedEncode;
-use rand::{thread_rng, Rng};
-
-const VERSION: DapVersion = DapVersion::Draft09;
-const VERIFY_KEY: [u8; 32] = [0; 32];
-
-fn aggregate(c: &mut Criterion) {
-    let bearer_token = BearerToken::from("the-bearer-token");
-    let tasks = make_tasks();
-    let aggregator = make_aggregator(tasks.clone().into_iter(), bearer_token.clone());
-    let runtime = tokio::runtime::Runtime::new().unwrap();
-    for (task_id, task) in tasks {
-        aggregator.clear_storage();
-        let req = runtime.block_on(make_request(
-            &aggregator,
-            task_id,
-            &task,
-            bearer_token.clone(),
-        ));
-        let mut g = c.benchmark_group(vdaf_short_name(&task.vdaf));
-        g.throughput(Throughput::Elements(task.min_batch_size));
-        g.bench_with_input(format!("{}", task.min_batch_size), &(), |b, ()| {
-            b.to_async(&runtime).iter_custom(|iters| {
-                let aggregator = &aggregator;
-                let req = &req;
-                async move {
-                    let mut total = Duration::ZERO;
-                    for _ in 0..iters {
-                        let now = Instant::now();
-                        let ret = black_box(
-                            helper::handle_agg_job_req(aggregator, req, Default::default()).await,
-                        );
-                        total += now.elapsed();
-                        aggregator.clear_storage();
-                        drop(ret.unwrap());
-                    }
-                    total
-                }
-            });
-        });
-    }
-}
+use tokio::runtime::Runtime;
 
-fn make_tasks() -> Vec<(TaskId, DapTaskConfig)> {
-    [1, 10, 100, 1000, 10_000]
-        .into_iter()
-        .flat_map(|num_reports| {
-            [100, 1_000, 10_000, 100_000]
-                .into_iter()
-                .map(move |dimension| (num_reports, dimension))
-        })
-        .map(|(num_reports, dimension)| DapTaskConfig {
-            version: daphne::DapVersion::Draft09,
-            leader_url: "http://1.1.1.1".parse().unwrap(),
-            helper_url: "http://1.1.1.1".parse().unwrap(),
-            time_precision: 3600,
-            min_batch_size: num_reports,
-            query: daphne::DapQueryConfig::FixedSize {
-                max_batch_size: Some(num_reports),
-            },
-            vdaf: daphne::vdaf::VdafConfig::Prio3(
-                daphne::vdaf::Prio3Config::SumVecField64MultiproofHmacSha256Aes128 {
-                    bits: 1,
-                    length: dimension,
-                    chunk_length: 320,
-                    num_proofs: 2,
-                },
-            ),
-            not_after: (SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap()
-                + Duration::from_secs(5000))
-            .as_secs(),
-            not_before: SystemTime::now()
-                .duration_since(SystemTime::UNIX_EPOCH)
-                .unwrap()
-                .as_secs(),
-            vdaf_verify_key: VdafVerifyKey::L32(VERIFY_KEY),
-            collector_hpke_config: HpkeConfig {
-                id: 1,
-                kem_id: HpkeKemId::P256HkdfSha256,
-                kdf_id: HpkeKdfId::HkdfSha256,
-                aead_id: HpkeAeadId::Aes128Gcm,
-                public_key: HpkePublicKey::new(vec![]),
-            },
-            method: daphne::DapTaskConfigMethod::Unknown,
-            num_agg_span_shards: NonZeroUsize::new(1).unwrap(),
-        })
-        .map(|task| (TaskId(thread_rng().gen()), task))
-        .last()
-        .into_iter()
-        .collect()
-}
+macro_rules! function {
+    () => {{
+        fn f() {}
+        let name = std::any::type_name_of_val(&f);
 
-fn make_aggregator(
-    tasks: impl Iterator<Item = (TaskId, DapTaskConfig)>,
-    token: BearerToken,
-) -> InMemoryAggregator {
-    let hpke = HpkeReceiverConfig::gen(1, HpkeKemId::P256HkdfSha256).unwrap();
-    let collector_hpke = HpkeReceiverConfig::gen(2, HpkeKemId::P256HkdfSha256)
-        .unwrap()
-        .config;
-    InMemoryAggregator::new_helper(
-        tasks,
-        [hpke],
-        daphne::DapGlobalConfig {
-            max_batch_duration: u64::MAX,
-            min_batch_interval_start: u64::MAX,
-            max_batch_interval_end: u64::MAX,
-            supported_hpke_kems: vec![HpkeKemId::P256HkdfSha256],
-            allow_taskprov: false,
-            default_num_agg_span_shards: NonZeroUsize::new(1).unwrap(),
-        },
-        token.clone(),
-        collector_hpke,
-        prometheus::default_registry(),
-        VERIFY_KEY,
-        token,
-    )
+        // Find and cut the rest of the path
+        match &name[..name.len() - 3].rfind(':') {
+            Some(pos) => &name[pos + 1..name.len() - 3],
+            None => &name[..name.len() - 3],
+        }
+    }};
 }
 
-async fn make_request(
-    aggregator: &InMemoryAggregator,
-    task_id: TaskId,
-    task_config: &DapTaskConfig,
-    bearer_token: BearerToken,
-) -> DapRequest<BearerToken> {
-    let fake_leader_hpke = HpkeReceiverConfig::gen(2, HpkeKemId::P256HkdfSha256).unwrap();
-    let leader = InMemoryAggregator::new_leader(
-        [],
-        [fake_leader_hpke],
-        aggregator.global_config.clone(),
-        aggregator.leader_token.clone(),
-        aggregator.collector_token.clone(),
-        aggregator.collector_hpke_config.clone(),
-        &prometheus::Registry::new(),
-        VERIFY_KEY,
-        "taskprov_leader_token".into(),
-        None,
-        None,
+fn consume_reports_vary_vdaf_dimension(c: &mut Criterion) {
+    const NUM_REPORTS: u64 = 1000;
+    let vdaf_lengths = [10, 100, 1_000, 10_000, 100_000];
+    let mut test = AggregationJobTest::new(
+        &VdafConfig::Prio2 { dimension: 0 },
+        HpkeKemId::P256HkdfSha256,
+        DapVersion::Latest,
     );
-    let batch_id = BatchId(thread_rng().gen());
-
-    let hpke_config_list = [
-        leader.hpke_receiver_config_list[0].config.clone(),
-        aggregator.hpke_receiver_config_list[0].config.clone(),
-    ];
-
-    let payload = task_config
-        .produce_agg_job_req(
-            &leader,
-            &leader,
-            &task_id,
-            &daphne::messages::PartialBatchSelector::FixedSizeByBatchId { batch_id },
-            &daphne::DapAggregationParam::Empty,
-            ReportGenerator::new(
-                &task_config.vdaf,
-                &hpke_config_list,
-                task_id,
-                task_config.min_batch_size.try_into().unwrap(),
-                &task_config.vdaf.gen_measurement().unwrap(),
-                VERSION,
-                task_config.not_after - 10,
-                vec![],
-            ),
-            aggregator.metrics(),
-        )
-        .await
-        .unwrap()
-        .1;
-    assert_ne!(payload.prep_inits.len(), 0);
-    DapRequest {
-        version: daphne::DapVersion::Draft09,
-        media_type: Some(daphne::constants::DapMediaType::AggregationJobInitReq),
-        task_id: Some(task_id),
-        resource: daphne::DapResource::AggregationJob(daphne::messages::AggregationJobId(
-            thread_rng().gen(),
-        )),
-        payload: payload.get_encoded_with_param(&VERSION).unwrap(),
-        sender_auth: Some(bearer_token),
-        taskprov: None,
+    test.disable_replay_protection();
+
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+    let mut g = c.benchmark_group(function!());
+    for vdaf_length in vdaf_lengths {
+        let vdaf = VdafConfig::Prio3(Prio3Config::SumVecField64MultiproofHmacSha256Aes128 {
+            bits: 1,
+            length: vdaf_length,
+            chunk_length: 320,
+            num_proofs: 2,
+        });
+        test.change_vdaf(vdaf);
+        let reports = test
+            .produce_repeated_reports(vdaf.gen_measurement().unwrap())
+            .take(NUM_REPORTS as _);
+
+        let (_, init) =
+            runtime.block_on(test.produce_agg_job_req(&DapAggregationParam::Empty, reports));
+
+        g.throughput(Throughput::Bytes(vdaf_length as _));
+        g.bench_with_input(
+            BenchmarkId::new("consume_agg_job_req", vdaf_length),
+            &init,
+            |b, init| bench(b, &test, init, &runtime),
+        );
     }
 }
 
-fn vdaf_short_name(vdaf: &VdafConfig) -> String {
-    match vdaf {
-        VdafConfig::Prio3(daphne::vdaf::Prio3Config::SumVecField64MultiproofHmacSha256Aes128 {
-            length,
-            ..
-        }) => format!("prio3-sum-vec-multiproof-{length}"),
-        _ => todo!(),
+fn consume_reports_vary_num_reports(c: &mut Criterion) {
+    const VDAF: VdafConfig =
+        VdafConfig::Prio3(Prio3Config::SumVecField64MultiproofHmacSha256Aes128 {
+            bits: 1,
+            length: 1000,
+            chunk_length: 320,
+            num_proofs: 2,
+        });
+
+    let mut test = AggregationJobTest::new(&VDAF, HpkeKemId::P256HkdfSha256, DapVersion::Latest);
+    test.disable_replay_protection();
+
+    let runtime = tokio::runtime::Builder::new_current_thread()
+        .build()
+        .unwrap();
+
+    let mut g = c.benchmark_group(function!());
+    for report_counts in [10, 100, 1_000, 10_000] {
+        let reports = test
+            .produce_repeated_reports(VDAF.gen_measurement().unwrap())
+            .take(report_counts);
+
+        let (_, init) =
+            runtime.block_on(test.produce_agg_job_req(&DapAggregationParam::Empty, reports));
+
+        g.throughput(Throughput::Elements(report_counts as _));
+        g.bench_with_input(
+            BenchmarkId::new("consume_agg_job_req", report_counts),
+            &init,
+            |b, init| bench(b, &test, init, &runtime),
+        );
     }
 }
 
-criterion_group!(benches, aggregate);
+fn bench(
+    b: &mut Bencher,
+    test: &AggregationJobTest,
+    init: &AggregationJobInitReq,
+    runtime: &Runtime,
+) {
+    b.to_async(runtime).iter_custom(|iters| async move {
+        let mut total = Duration::ZERO;
+        for init in repeat(init).take(iters as _).cloned() {
+            let now = Instant::now();
+            let ret = black_box(test.handle_agg_job_req(init).await);
+            total += now.elapsed();
+            drop(ret);
+        }
+        total
+    });
+}
+
+criterion_group!(
+    benches,
+    consume_reports_vary_num_reports,
+    consume_reports_vary_vdaf_dimension
+);
 criterion_main!(benches);
diff --git a/crates/daphne/src/protocol/aggregator.rs b/crates/daphne/src/protocol/aggregator.rs
index 6a6e8ef..a1e4087 100644
--- a/crates/daphne/src/protocol/aggregator.rs
+++ b/crates/daphne/src/protocol/aggregator.rs
@@ -383,11 +383,11 @@ pub(crate) enum ReportProcessedStatus {
     Rejected(TransitionFailure),
 }
 
-#[derive(Default, Debug)]
+#[derive(Default, Debug, Clone, Copy)]
 pub enum ReplayProtection {
     #[default]
     Enabled,
-    Disabled,
+    InsecureDisabled,
 }
 
 impl ReplayProtection {
@@ -396,7 +396,7 @@ impl ReplayProtection {
     }
 
     pub const fn disabled(&self) -> bool {
-        matches!(self, ReplayProtection::Disabled)
+        matches!(self, ReplayProtection::InsecureDisabled)
     }
 }
 
@@ -414,6 +414,34 @@ impl DapTaskConfig {
         reports: S,
         metrics: &dyn DaphneMetrics,
     ) -> Result<(DapAggregationJobState, AggregationJobInitReq), DapError>
+    where
+        S: Stream<Item = Report>,
+    {
+        self.produce_agg_job_req_impl(
+            decrypter,
+            initializer,
+            task_id,
+            part_batch_sel,
+            agg_param,
+            reports,
+            metrics,
+            ReplayProtection::Enabled,
+        )
+        .await
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    async fn produce_agg_job_req_impl<S>(
+        &self,
+        decrypter: &impl HpkeDecrypter,
+        initializer: &impl DapReportInitializer,
+        task_id: &TaskId,
+        part_batch_sel: &PartialBatchSelector,
+        agg_param: &DapAggregationParam,
+        reports: S,
+        metrics: &dyn DaphneMetrics,
+        replay_protection: ReplayProtection,
+    ) -> Result<(DapAggregationJobState, AggregationJobInitReq), DapError>
     where
         S: Stream<Item = Report>,
     {
@@ -421,24 +449,20 @@ impl DapTaskConfig {
         let mut consumed_reports = Vec::with_capacity(report_count_hint);
         let mut helper_shares = Vec::with_capacity(report_count_hint);
         {
-            let mut processed = HashSet::with_capacity(report_count_hint);
+            let mut processed = replay_protection
+                .enabled()
+                .then(|| HashSet::with_capacity(report_count_hint));
             let mut reports = pin!(reports);
             while let Some(report) = reports.next().await {
-                let enforce_replay_check = {
-                    #[cfg(feature = "report-generator")]
-                    {
-                        !crate::testing::report_generator::generator_replaying_reports()
+                if let Some(processed) = &mut processed {
+                    if processed.contains(&report.report_metadata.id) {
+                        return Err(fatal_error!(
+                            err = "tried to process report sequence with non-unique report IDs",
+                            non_unique_id = %report.report_metadata.id,
+                        ));
                     }
-                    #[cfg(not(feature = "report-generator"))]
-                    true
-                };
-                if enforce_replay_check && processed.contains(&report.report_metadata.id) {
-                    return Err(fatal_error!(
-                        err = "tried to process report sequence with non-unique report IDs",
-                        non_unique_id = %report.report_metadata.id,
-                    ));
+                    processed.insert(report.report_metadata.id);
                 }
-                processed.insert(report.report_metadata.id);
 
                 let [leader_share, helper_share] = report.encrypted_input_shares;
 
@@ -529,6 +553,35 @@ impl DapTaskConfig {
         ))
     }
 
+    #[allow(clippy::too_many_arguments)]
+    #[cfg(any(test, feature = "test-utils"))]
+    pub async fn test_produce_agg_job_req<S>(
+        &self,
+        decrypter: &impl HpkeDecrypter,
+        initializer: &impl DapReportInitializer,
+        task_id: &TaskId,
+        part_batch_sel: &PartialBatchSelector,
+        agg_param: &DapAggregationParam,
+        reports: S,
+        metrics: &dyn DaphneMetrics,
+        replay_protection: ReplayProtection,
+    ) -> Result<(DapAggregationJobState, AggregationJobInitReq), DapError>
+    where
+        S: Stream<Item = Report>,
+    {
+        self.produce_agg_job_req_impl(
+            decrypter,
+            initializer,
+            task_id,
+            part_batch_sel,
+            agg_param,
+            reports,
+            metrics,
+            replay_protection,
+        )
+        .await
+    }
+
     /// Helper: Consume the `AggregationJobInitReq` sent by the Leader and return the initialized
     /// reports.
     pub async fn consume_agg_job_req(
@@ -542,21 +595,24 @@ impl DapTaskConfig {
         let num_reports = agg_job_init_req.prep_inits.len();
         let mut consumed_reports = Vec::with_capacity(num_reports);
         async {
-            let mut processed = HashSet::with_capacity(num_reports);
+            let mut processed = replay_protection
+                .enabled()
+                .then(|| HashSet::with_capacity(num_reports));
+
             for prep_init in agg_job_init_req.prep_inits {
-                if replay_protection.enabled()
-                    && processed.contains(&prep_init.report_share.report_metadata.id)
-                {
-                    return Err(DapAbort::InvalidMessage {
-                        detail: format!(
-                            "report ID {} appears twice in the same aggregation job",
-                            prep_init.report_share.report_metadata.id.to_base64url()
-                        ),
-                        task_id: *task_id,
+                if let Some(processed) = &mut processed {
+                    if processed.contains(&prep_init.report_share.report_metadata.id) {
+                        return Err(DapAbort::InvalidMessage {
+                            detail: format!(
+                                "report ID {} appears twice in the same aggregation job",
+                                prep_init.report_share.report_metadata.id.to_base64url()
+                            ),
+                            task_id: *task_id,
+                        }
+                        .into());
                     }
-                    .into());
+                    processed.insert(prep_init.report_share.report_metadata.id);
                 }
-                processed.insert(prep_init.report_share.report_metadata.id);
 
                 consumed_reports.push(
                     EarlyReportStateConsumed::consume(
diff --git a/crates/daphne/src/roles/leader/in_memory_leader.rs b/crates/daphne/src/roles/leader/in_memory_leader.rs
index 32baaab..a024dc7 100644
--- a/crates/daphne/src/roles/leader/in_memory_leader.rs
+++ b/crates/daphne/src/roles/leader/in_memory_leader.rs
@@ -22,6 +22,7 @@ use crate::{
 };
 
 #[derive(Default)]
+#[cfg_attr(feature = "test-utils", derive(deepsize::DeepSizeOf))]
 pub struct InMemoryLeaderState {
     work_queue: VecDeque<WorkItem>,
     per_task: HashMap<TaskId, MockLeaderMemoryPerTask>,
@@ -232,6 +233,7 @@ impl InMemoryLeaderState {
 }
 
 #[derive(Default)]
+#[cfg_attr(feature = "test-utils", derive(deepsize::DeepSizeOf))]
 struct MockLeaderMemoryPerTask {
     pending_reports: HashMap<DapBatchBucket, VecDeque<Report>>,
     coll_jobs: HashMap<CollectionJobId, DapCollectionJob>,
diff --git a/crates/daphne/src/roles/mod.rs b/crates/daphne/src/roles/mod.rs
index 1e77bc3..cb62c71 100644
--- a/crates/daphne/src/roles/mod.rs
+++ b/crates/daphne/src/roles/mod.rs
@@ -160,7 +160,7 @@ mod test {
             Extension, HpkeCiphertext, Interval, PartialBatchSelector, Query, Report, TaskId, Time,
             TransitionFailure, TransitionVar,
         },
-        roles::leader::WorkItem,
+        roles::{leader::WorkItem, DapAggregator},
         testing::InMemoryAggregator,
         vdaf::{Prio3Config, VdafConfig},
         DapAbort, DapAggregationJobState, DapAggregationParam, DapBatchBucket, DapCollectionJob,
@@ -501,7 +501,7 @@ mod test {
                     &part_batch_sel,
                     &agg_param,
                     futures::stream::iter(reports),
-                    &self.leader.metrics,
+                    self.leader.metrics(),
                 )
                 .await
                 .unwrap();
@@ -1625,8 +1625,8 @@ mod test {
         .to_config_with_taskprov(
             b"cool task".to_vec(),
             t.now,
-            &t.leader.taskprov_vdaf_verify_key_init,
-            &t.leader.collector_hpke_config,
+            t.leader.taskprov_vdaf_verify_key_init().unwrap(),
+            t.leader.taskprov_collector_hpke_config().unwrap(),
         )
         .unwrap();
 


Test patch:
diff --git a/crates/dapf/src/acceptance/load_testing.rs b/crates/dapf/src/acceptance/load_testing.rs
index c4bf4fe..d0266c6 100644
--- a/crates/dapf/src/acceptance/load_testing.rs
+++ b/crates/dapf/src/acceptance/load_testing.rs
@@ -449,6 +449,7 @@ pub async fn execute_single_combination_from_env(
                             VERSION,
                             system_now.0,
                             vec![messages::Extension::Taskprov],
+                            t.replay_reports,
                         )
                     },
                 )
diff --git a/crates/daphne/src/testing/mod.rs b/crates/daphne/src/testing/mod.rs
index 95de8d9..e558d21 100644
--- a/crates/daphne/src/testing/mod.rs
+++ b/crates/daphne/src/testing/mod.rs
@@ -29,7 +29,8 @@ use crate::{
     vdaf::VdafVerifyKey,
     DapAbort, DapAggregateResult, DapAggregateShare, DapAggregateSpan, DapAggregationJobState,
     DapAggregationParam, DapBatchBucket, DapCollectionJob, DapError, DapGlobalConfig,
-    DapMeasurement, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig, DapVersion, VdafConfig,
+    DapMeasurement, DapQueryConfig, DapRequest, DapResponse, DapTaskConfig, DapVersion,
+    ReplayProtection, VdafConfig,
 };
 use async_trait::async_trait;
 use deepsize::DeepSizeOf;
@@ -58,6 +59,8 @@ pub struct AggregationJobTest {
     pub(crate) client_hpke_config_list: [HpkeConfig; 2],
     pub(crate) collector_hpke_receiver_config: HpkeReceiverConfig,
 
+    replay_protection: ReplayProtection,
+
     // the current time
     pub(crate) now: Time,
     pub(crate) valid_report_range: Range<Time>,
@@ -194,11 +197,21 @@ impl AggregationJobTest {
                 method: Default::default(),
                 num_agg_span_shards: NonZeroUsize::new(3).unwrap(),
             },
+            replay_protection: ReplayProtection::Enabled,
             leader_registry,
             leader_metrics,
         }
     }
 
+    pub fn disable_replay_protection(&mut self) {
+        self.replay_protection = ReplayProtection::InsecureDisabled;
+    }
+
+    pub fn change_vdaf(&mut self, vdaf: VdafConfig) {
+        self.task_config.vdaf = vdaf;
+        self.task_config.vdaf_verify_key = vdaf.gen_verify_key();
+    }
+
     /// For each measurement, generate a report for the given task.
     ///
     /// Panics if a measurement is incompatible with the given VDAF.
@@ -222,16 +235,34 @@ impl AggregationJobTest {
         reports
     }
 
+    pub fn produce_repeated_reports(
+        &self,
+        measurement: DapMeasurement,
+    ) -> impl Iterator<Item = Report> + Clone {
+        std::iter::repeat(
+            self.task_config
+                .vdaf
+                .produce_report(
+                    &self.client_hpke_config_list,
+                    self.now,
+                    &self.task_id,
+                    measurement,
+                    self.task_config.version,
+                )
+                .unwrap(),
+        )
+    }
+
     /// Leader: Produce `AggregationJobInitReq`.
     ///
     /// Panics if the Leader aborts.
     pub async fn produce_agg_job_req(
         &self,
         agg_param: &DapAggregationParam,
-        reports: Vec<Report>,
+        reports: impl IntoIterator<Item = Report>,
     ) -> (DapAggregationJobState, AggregationJobInitReq) {
         self.task_config
-            .produce_agg_job_req(
+            .test_produce_agg_job_req(
                 &self.leader_hpke_receiver_config,
                 self,
                 &self.task_id,
@@ -239,6 +270,7 @@ impl AggregationJobTest {
                 agg_param,
                 futures::stream::iter(reports),
                 &self.leader_metrics,
+                self.replay_protection,
             )
             .await
             .unwrap()
@@ -262,7 +294,7 @@ impl AggregationJobTest {
                         self,
                         &self.task_id,
                         agg_job_init_req,
-                        Default::default(),
+                        self.replay_protection,
                     )
                     .await
                     .unwrap(),
@@ -542,48 +574,61 @@ impl InMemoryAggregateStore {
 /// An implementation of a DAP Aggregator without long-term storage. This is intended to be used
 /// for testing purposes only.
 pub struct InMemoryAggregator {
-    pub global_config: DapGlobalConfig,
-    pub(crate) tasks: Arc<Mutex<HashMap<TaskId, DapTaskConfig>>>,
-    pub hpke_receiver_config_list: Vec<HpkeReceiverConfig>,
-    pub leader_token: BearerToken,
-    pub collector_token: Option<BearerToken>, // Not set by Helper
-    pub(crate) leader_state_store: Arc<Mutex<InMemoryLeaderState>>,
-    helper_state_store: Arc<Mutex<HashMap<HelperStateInfo, DapAggregationJobState>>>,
+    pub(crate) global_config: DapGlobalConfig,
+    tasks: Mutex<HashMap<TaskId, DapTaskConfig>>,
+    pub hpke_receiver_config_list: Box<[HpkeReceiverConfig]>,
+    leader_token: BearerToken,
+    collector_token: Option<BearerToken>, // Not set by Helper
+    collector_hpke_config: HpkeConfig,
+
+    // aggregation state
+    leader_state_store: Arc<Mutex<InMemoryLeaderState>>,
     pub(crate) agg_store: Arc<Mutex<InMemoryAggregateStore>>,
-    pub collector_hpke_config: HpkeConfig,
-    pub metrics: DaphnePromMetrics,
-    pub(crate) audit_log: MockAuditLog,
+
+    // telemetry
+    metrics: DaphnePromMetrics,
+    pub audit_log: MockAuditLog,
 
     // taskprov
-    pub taskprov_vdaf_verify_key_init: [u8; 32],
-    pub taskprov_leader_token: BearerToken,
-    pub taskprov_collector_token: Option<BearerToken>, // Not set by Helper
+    taskprov_vdaf_verify_key_init: [u8; 32],
+    taskprov_leader_token: BearerToken,
+    taskprov_collector_token: Option<BearerToken>, // Not set by Helper
 
     // Leader: Reference to peer. Used to simulate HTTP requests from Leader to Helper, i.e.,
     // implement `DapLeader::send_http_post()` for `InMemoryAggregator`. Not set by the Helper.
-    pub peer: Option<Arc<InMemoryAggregator>>,
+    peer: Option<Arc<InMemoryAggregator>>,
 }
 
 impl DeepSizeOf for InMemoryAggregator {
     fn deep_size_of_children(&self, context: &mut deepsize::Context) -> usize {
-        self.global_config.deep_size_of_children(context)
-                + self.tasks.deep_size_of_children(context)
-                + self
-                    .hpke_receiver_config_list
-                    .deep_size_of_children(context)
-                + self.leader_token.deep_size_of_children(context)
-                + self.collector_token.deep_size_of_children(context)
-                + self.helper_state_store.deep_size_of_children(context)
-                + self.agg_store.deep_size_of_children(context)
-                + self.collector_hpke_config.deep_size_of_children(context)
-                // + self.metrics.deep_size_of_children(context)
-                // + self.audit_log.deep_size_of_children(context)
-                + self
-                    .taskprov_vdaf_verify_key_init
-                    .deep_size_of_children(context)
-                + self.taskprov_leader_token.deep_size_of_children(context)
-                + self.taskprov_collector_token.deep_size_of_children(context)
-                + self.peer.deep_size_of_children(context)
+        let Self {
+            global_config,
+            tasks,
+            hpke_receiver_config_list,
+            leader_token,
+            collector_token,
+            leader_state_store,
+            agg_store,
+            collector_hpke_config,
+            metrics: _,
+            audit_log: _,
+            taskprov_vdaf_verify_key_init,
+            taskprov_leader_token,
+            taskprov_collector_token,
+            peer,
+        } = self;
+        global_config.deep_size_of_children(context)
+            + tasks.deep_size_of_children(context)
+            + hpke_receiver_config_list.deep_size_of_children(context)
+            + leader_token.deep_size_of_children(context)
+            + collector_token.deep_size_of_children(context)
+            + leader_state_store.deep_size_of_children(context)
+            + agg_store.deep_size_of_children(context)
+            + collector_hpke_config.deep_size_of_children(context)
+            + taskprov_vdaf_verify_key_init.deep_size_of_children(context)
+            + taskprov_leader_token.deep_size_of_children(context)
+            + taskprov_collector_token.deep_size_of_children(context)
+            + peer.deep_size_of_children(context)
     }
 }
 
@@ -601,12 +646,11 @@ impl InMemoryAggregator {
     ) -> Self {
         Self {
             global_config,
-            tasks: Arc::new(Mutex::new(tasks.into_iter().collect())),
+            tasks: Mutex::new(tasks.into_iter().collect()),
             hpke_receiver_config_list: hpke_receiver_config_list.into_iter().collect(),
             leader_token,
             collector_token: None,
             leader_state_store: Default::default(),
-            helper_state_store: Default::default(),
             agg_store: Default::default(),
             collector_hpke_config,
             metrics: DaphnePromMetrics::register(registry).unwrap(),
@@ -630,16 +674,15 @@ impl InMemoryAggregator {
         taskprov_vdaf_verify_key_init: [u8; 32],
         taskprov_leader_token: BearerToken,
         taskprov_collector_token: impl Into<Option<BearerToken>>,
-        peer: impl Into<Option<Arc<Self>>>,
+        peer: Arc<Self>,
     ) -> Self {
         Self {
             global_config,
-            tasks: Arc::new(Mutex::new(tasks.into_iter().collect())),
+            tasks: Mutex::new(tasks.into_iter().collect()),
             hpke_receiver_config_list: hpke_receiver_config_list.into_iter().collect(),
             leader_token,
             collector_token: collector_token.into(),
             leader_state_store: Default::default(),
-            helper_state_store: Default::default(),
             agg_store: Default::default(),
             collector_hpke_config,
             metrics: DaphnePromMetrics::register(registry).unwrap(),
@@ -670,7 +713,6 @@ impl InMemoryAggregator {
 
     pub fn clear_storage(&self) {
         self.leader_state_store.lock().unwrap().delete_all();
-        self.helper_state_store.lock().unwrap().clear();
         self.agg_store.lock().unwrap().clear();
     }
 }
diff --git a/crates/daphne/src/testing/report_generator.rs b/crates/daphne/src/testing/report_generator.rs
index 0bd254e..be6c126 100644
--- a/crates/daphne/src/testing/report_generator.rs
+++ b/crates/daphne/src/testing/report_generator.rs
@@ -4,7 +4,7 @@
 use std::{
     pin::Pin,
     sync::{
-        atomic::{AtomicBool, AtomicUsize, Ordering},
+        atomic::{AtomicUsize, Ordering},
         Mutex, OnceLock,
     },
     task::Context,
@@ -56,14 +56,6 @@ impl Stream for ReportGenerator {
     }
 }
 
-static REPLAY_REPORTS: AtomicBool = AtomicBool::new(false);
-pub fn replay_reports(b: bool) {
-    REPLAY_REPORTS.store(b, Ordering::SeqCst);
-}
-pub fn generator_replaying_reports() -> bool {
-    REPLAY_REPORTS.load(Ordering::SeqCst)
-}
-
 impl ReportGenerator {
     #[allow(clippy::too_many_arguments)]
     pub fn new(
@@ -75,6 +67,7 @@ impl ReportGenerator {
         version: DapVersion,
         now: Time,
         extensions: Vec<messages::Extension>,
+        replay_reports: bool,
     ) -> Self {
         let (tx, rx) = mpsc::channel(4);
         rayon::spawn({
@@ -98,7 +91,7 @@ impl ReportGenerator {
                     // ----
 
                     static LAST_REPORT: OnceLock<Report> = OnceLock::new();
-                    let report = if REPLAY_REPORTS.load(Ordering::SeqCst) {
+                    let report = if replay_reports {
                         LAST_REPORT
                             .get_or_init(|| {
                                 vdaf.produce_report_with_extensions(