crypto: Add a bench for the key claiming process

matrix_sdk_crypto/benches/crypto_bench.rs

@@ -1,12 +1,13 @@
 use std::convert::TryFrom;
 
 use criterion::{
-    async_executor::FuturesExecutor, criterion_group, criterion_main, BenchmarkId, Criterion,
+    async_executor::FuturesExecutor, criterion_group, criterion_main, BatchSize, BenchmarkId,
-    Throughput,
+    Criterion, Throughput,
 };
 
+use futures::executor::block_on;
 use matrix_sdk_common::{
-    api::r0::keys::get_keys,
+    api::r0::keys::{claim_keys, get_keys},
     identifiers::{user_id, DeviceIdBox, UserId},
     uuid::Uuid,
 };
@@ -29,7 +30,14 @@ fn keys_query_response() -> get_keys::Response {
     get_keys::Response::try_from(data).expect("Can't parse the keys upload response")
 }
 
-pub fn receive_keys_query(c: &mut Criterion) {
+fn keys_claim_response() -> claim_keys::Response {
+    let data = include_bytes!("./keys_claim.json");
+    let data: Value = serde_json::from_slice(data).unwrap();
+    let data = response_from_file(&data);
+    claim_keys::Response::try_from(data).expect("Can't parse the keys upload response")
+}
+
+pub fn keys_query(c: &mut Criterion) {
     let machine = OlmMachine::new(&alice_id(), &alice_device_id());
     let response = keys_query_response();
     let uuid = Uuid::new_v4();
@@ -42,11 +50,14 @@ pub fn receive_keys_query(c: &mut Criterion) {
         + response.self_signing_keys.len()
         + response.user_signing_keys.len();
 
-    let mut group = c.benchmark_group("key query throughput");
+    let mut group = c.benchmark_group("Keys querying");
     group.throughput(Throughput::Elements(count as u64));
 
     group.bench_with_input(
-        BenchmarkId::new("key_query", "150 devices key query response parsing"),
+        BenchmarkId::new(
+            "Keys querying",
+            "150 device keys parsing and signature checking",
+        ),
         &response,
         |b, response| {
             b.to_async(FuturesExecutor)
@@ -56,5 +67,39 @@ pub fn receive_keys_query(c: &mut Criterion) {
     group.finish()
 }
 
-criterion_group!(benches, receive_keys_query);
+pub fn keys_claiming(c: &mut Criterion) {
+    let keys_query_response = keys_query_response();
+    let uuid = Uuid::new_v4();
+
+    let response = keys_claim_response();
+
+    let count = response
+        .one_time_keys
+        .values()
+        .fold(0, |acc, d| acc + d.len());
+
+    let mut group = c.benchmark_group("Keys claiming throughput");
+    group.throughput(Throughput::Elements(count as u64));
+
+    let name = format!("{} one-time keys claiming and session creation", count);
+
+    group.bench_with_input(
+        BenchmarkId::new("One-time keys claiming", &name),
+        &response,
+        |b, response| {
+            b.iter_batched(
+                || {
+                    let machine = OlmMachine::new(&alice_id(), &alice_device_id());
+                    block_on(machine.mark_request_as_sent(&uuid, &keys_query_response)).unwrap();
+                    machine
+                },
+                move |machine| block_on(machine.mark_request_as_sent(&uuid, response)).unwrap(),
+                BatchSize::SmallInput,
+            )
+        },
+    );
+    group.finish()
+}
+
+criterion_group!(benches, keys_query, keys_claiming);
 criterion_main!(benches);