Merge pull request #834 from taikiy/event_rng

make random event generator generate "min_events_per_user..=max_events_per_user" number of events
private-attribution · Nov 8, 2023 · cb0fe9e · cb0fe9e
2 parents 4dd9554 + 62e5139
commit cb0fe9e
Show file tree

Hide file tree

Showing 4 changed files with 8,575 additions and 23 deletions.
diff --git a/benches/oneshot/ipa.rs b/benches/oneshot/ipa.rs
@@ -1,5 +1,5 @@
 use std::{
-    num::{NonZeroU32, NonZeroUsize},
+    num::{NonZeroU32, NonZeroU64, NonZeroUsize},
     time::Instant,
 };
 
@@ -111,16 +111,36 @@ async fn run(args: Args) -> Result<(), Error> {
         q = args.query_size
     );
     let rng = StdRng::seed_from_u64(seed);
+    let (user_count, min_events_per_user, max_events_per_user, query_size) =
+        if args.oprf && cfg!(feature = "step-trace") {
+            // For the steps collection, OPRF mode requires a single user with the same number
+            // of dynamic steps as defined for `UserNthRowStep::Row`.
+            (
+                NonZeroU64::new(1).unwrap(),
+                NonZeroU32::new(64).unwrap(),
+                NonZeroU32::new(64).unwrap(),
+                64,
+            )
+        } else {
+            (
+                EventGeneratorConfig::default().user_count,
+                EventGeneratorConfig::default().min_events_per_user,
+                NonZeroU32::new(args.records_per_user).unwrap(),
+                args.query_size,
+            )
+        };
     let raw_data = EventGenerator::with_config(
         rng,
         EventGeneratorConfig {
+            user_count,
             max_trigger_value: NonZeroU32::try_from(args.max_trigger_value).unwrap(),
             max_breakdown_key: NonZeroU32::try_from(args.breakdown_keys).unwrap(),
-            max_events_per_user: NonZeroU32::try_from(args.records_per_user).unwrap(),
+            min_events_per_user,
+            max_events_per_user,
             ..Default::default()
         },
     )
-    .take(args.query_size)
+    .take(query_size)
     .collect::<Vec<_>>();
 
     let order = if args.oprf {

diff --git a/src/protocol/ipa/mod.rs b/src/protocol/ipa/mod.rs
@@ -782,6 +782,7 @@ pub mod tests {
         const MAX_BREAKDOWN_KEY: u32 = 32;
         const MAX_TRIGGER_VALUE: u32 = 5;
         const NUM_USERS: u32 = 8;
+        const MIN_RECORDS_PER_USER: u32 = 1;
         const MAX_RECORDS_PER_USER: u32 = 8;
         const NUM_MULTI_BITS: u32 = 3;
         const ATTRIBUTION_WINDOW_SECONDS: Option<NonZeroU32> = NonZeroU32::new(86_400);
@@ -803,6 +804,7 @@ pub mod tests {
                 u64::from(NUM_USERS),
                 MAX_TRIGGER_VALUE,
                 MAX_BREAKDOWN_KEY,
+                MIN_RECORDS_PER_USER,
                 MAX_RECORDS_PER_USER,
             ),
         )