From 6a610b897d26334bb4a510aeec0c7bbc99f29250 Mon Sep 17 00:00:00 2001
From: Robbie Blaine <robbieblaine.rb@gmail.com>
Date: Fri, 17 Nov 2023 12:57:03 +0200
Subject: [PATCH] Add DB benchmark

Add a benchmark using `criterion` to benchmark DB performance,
specifically when the DB has a large number of profiles/items/items_tags
* Set `bench = false`
  * https://bheisler.github.io/criterion.rs/book/faq.html#cargo-bench-gives-unrecognized-option-errors-for-valid-command-line-options

Signed-off-by: Robbie Blaine <robbieblaine.rb@gmail.com>
---
 Cargo.toml           |   9 ++++
 benches/README.md    |  20 +++++++++
 benches/benchmark.rs | 105 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 134 insertions(+)
 create mode 100644 benches/README.md
 create mode 100644 benches/benchmark.rs

diff --git a/Cargo.toml b/Cargo.toml
index 4dd6561b..95244440 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,7 @@ rust-version = "1.65"
 name = "aries_askar"
 path = "src/lib.rs"
 crate-type = ["staticlib", "rlib", "cdylib"]
+bench = false
 
 [package.metadata.docs.rs]
 features = ["all_backends"]
@@ -62,3 +63,11 @@ features = ["any"]
 codegen-units = 1
 lto = true
 panic = "abort"
+
+[dev-dependencies]
+criterion = { version = "0.5", features = ["html_reports"] }
+rand = { version = "0.8" }
+
+[[bench]]
+name = "benchmark"
+harness = false
diff --git a/benches/README.md b/benches/README.md
new file mode 100644
index 00000000..999be265
--- /dev/null
+++ b/benches/README.md
@@ -0,0 +1,20 @@
+### Aries Askar Benchmarks
+
+Running `cargo bench` will run the benchmarks against an in-memory SQLite by default.
+
+To run against a Postgres, you need to set the `POSTGRES_URL` environment variable like so:
+```sh
+docker run --rm -p 5432:5432 --net aries --name aries-test-postgres -e POSTGRES_PASSWORD=mysecretpassword -d postgres
+POSTGRES_URL=postgres://postgres:mysecretpassword@localhost:5432/test-db cargo bench
+```
+
+To run comparison benchmarks:
+```sh
+git checkout main
+cargo bench -- --save-baseline main
+git checkout feature
+cargo bench -- --save-baseline feature
+
+# Compare `feature` (new) to `main` (original)
+cargo bench -- --load-baseline feature --baseline main
+```
diff --git a/benches/benchmark.rs b/benches/benchmark.rs
new file mode 100644
index 00000000..29f6d687
--- /dev/null
+++ b/benches/benchmark.rs
@@ -0,0 +1,105 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+use rand::{distributions::Alphanumeric, Rng};
+
+use aries_askar::{
+    future::block_on,
+    kms::{KeyAlg, LocalKey},
+    Store, StoreKeyMethod,
+};
+
+const ERR_RAW_KEY: &str = "Error creating raw store key";
+const ERR_SESSION: &str = "Error creating store session";
+const ERR_OPEN: &str = "Error opening test store instance";
+const ERR_REQ_ROW: &str = "Row required";
+const ERR_CLOSE: &str = "Error closing test store instance";
+
+const ROOT_SEED: [u8; 32] = [0x55; 32];
+
+// Bootstrap a DB for benchmarking
+fn initialize_database() -> Store {
+    block_on(async {
+        let db_url = match std::env::var("POSTGRES_URL") {
+            Ok(p) if !p.is_empty() => p,
+            _ => "sqlite://:memory:".to_string(),
+        };
+        let pass_key = Store::new_raw_key(Some(&ROOT_SEED)).expect(ERR_RAW_KEY);
+
+        Store::provision(
+            &db_url,
+            StoreKeyMethod::RawKey,
+            pass_key,
+            Some("askar-bench".to_string()),
+            true,
+        )
+        .await
+        .expect(ERR_OPEN)
+    })
+}
+
+// Inject `n` number of keys and profiles into the DB
+fn populate_database_keys_profiles(db: &Store, n: u64) {
+    block_on(async {
+        let mut conn = db.session(None).await.expect(ERR_SESSION);
+
+        for _ in 0..n {
+            let keypair =
+                LocalKey::generate(KeyAlg::Ed25519, false).expect("Error creating keypair");
+            let key_name = rand::thread_rng()
+                .sample_iter(&Alphanumeric)
+                .take(10)
+                .map(char::from)
+                .collect::<String>();
+            let metadata = rand::thread_rng()
+                .sample_iter(&Alphanumeric)
+                .take(10)
+                .map(char::from)
+                .collect::<String>();
+
+            conn.insert_key(&key_name, &keypair, Some(metadata.as_str()), None, None)
+                .await
+                .expect("Error inserting key");
+
+            let found = conn
+                .fetch_key(&key_name, false)
+                .await
+                .expect("Error fetching key")
+                .expect(ERR_REQ_ROW);
+            assert_eq!(found.algorithm(), Some(KeyAlg::Ed25519.as_str()));
+            assert_eq!(found.name(), key_name);
+            assert_eq!(found.metadata(), Some(metadata.as_str()));
+            assert!(found.is_local());
+            found.load_local_key().expect("Error loading key");
+
+            db.create_profile(None)
+                .await
+                .expect("Error creating profile");
+        }
+
+        drop(conn);
+    });
+}
+
+// Initialize a clean DB
+// Pre-populate the DB with 50,000 keys and profiles
+// Run the benchmark, adding 1 key and profile every iteration
+// Close the DB
+fn criterion_benchmarks(c: &mut Criterion) {
+    let db = initialize_database();
+    populate_database_keys_profiles(&db, 50_000);
+
+    c.bench_function("benchmark_database", |b| {
+        b.iter(|| {
+            let db = db.clone();
+            populate_database_keys_profiles(&db, 1);
+        })
+    });
+
+    block_on(async { db.close().await.expect(ERR_CLOSE) });
+}
+
+criterion_group!(
+    name = benchmarks;
+    config = Criterion::default().sample_size(1_000);
+    targets = criterion_benchmarks
+);
+criterion_main!(benchmarks);