perf(store): Split store benchmarks (#243)

Author: john-z-yang

benches/store_bench.rs

@@ -11,7 +11,7 @@ use taskbroker::{
 };
 use tokio::task::JoinSet;
 
-async fn process_activations(num_activations: u32, num_workers: u32) {
+async fn get_pending_activations(num_activations: u32, num_workers: u32) {
     let url = if cfg!(feature = "bench-with-mnt-disk") {
         let mut rng = rand::thread_rng();
         format!(
@@ -48,23 +48,12 @@ async fn process_activations(num_activations: u32, num_workers: u32) {
         join_set.spawn(async move {
             let mut num_activations_processed = 0;
 
-            let mut pending = store
+            while store
                 .get_pending_activation(Some("namespace"))
                 .await
-                .unwrap();
-
-            while let Some(ref activation) = pending {
-                store
-                    .set_status(
-                        &activation.activation.id,
-                        InflightActivationStatus::Complete,
-                    )
-                    .await
-                    .unwrap();
-                pending = store
-                    .get_pending_activation(Some("namespace"))
-                    .await
-                    .unwrap();
+                .unwrap()
+                .is_some()
+            {
                 num_activations_processed += 1;
             }
             num_activations_processed
@@ -76,20 +65,77 @@ async fn process_activations(num_activations: u32, num_workers: u32) {
     );
 }
 
+async fn set_status(num_activations: u32, num_workers: u32) {
+    assert!(num_activations % num_workers == 0);
+
+    let url = if cfg!(feature = "bench-with-mnt-disk") {
+        let mut rng = rand::thread_rng();
+        format!(
+            "/mnt/disks/sqlite/{}-{}.sqlite",
+            Utc::now(),
+            rng.r#gen::<u64>()
+        )
+    } else {
+        generate_temp_filename()
+    };
+    let store = Arc::new(
+        InflightActivationStore::new(
+            &url,
+            InflightActivationStoreConfig {
+                max_processing_attempts: 1,
+            },
+        )
+        .await
+        .unwrap(),
+    );
+
+    for chunk in make_activations(num_activations).chunks(1024) {
+        store.store(chunk.to_vec()).await.unwrap();
+    }
+
+    assert_eq!(
+        store.count_pending_activations().await.unwrap(),
+        num_activations as usize
+    );
+
+    let mut join_set = JoinSet::new();
+    for i in 0..num_workers {
+        let store = store.clone();
+        join_set.spawn(async move {
+            for idx in 0..num_activations {
+                if idx % num_workers == i {
+                    store
+                        .set_status(&format!("id_{}", i), InflightActivationStatus::Complete)
+                        .await
+                        .unwrap();
+                }
+            }
+        });
+    }
+}
+
 fn store_bench(c: &mut Criterion) {
     let num_activations: u32 = 4_096;
     let num_workers = 64;
 
     c.benchmark_group("bench_InflightActivationStore")
         .sample_size(256)
         .throughput(criterion::Throughput::Elements(num_activations.into()))
-        .bench_function("process_activations", |b| {
+        .bench_function("get_pending_activation", |b| {
+            let runtime = tokio::runtime::Builder::new_multi_thread()
+                .enable_all()
+                .build()
+                .unwrap();
+            b.to_async(runtime)
+                .iter(|| get_pending_activations(num_activations, num_workers));
+        })
+        .bench_function("set_status", |b| {
             let runtime = tokio::runtime::Builder::new_multi_thread()
                 .enable_all()
                 .build()
                 .unwrap();
             b.to_async(runtime)
-                .iter(|| process_activations(num_activations, num_workers));
+                .iter(|| set_status(num_activations, num_workers));
         });
 }