identity benchmarks

dev/flamegraph

@@ -12,4 +12,4 @@ else
   XMTP_FLAMEGRAPH=trace cargo bench --no-fail-fast --features bench -- "$1"
 fi
 
-inferno-flamegraph > tracing-flamegraph.svg
+cat xmtp_mls/tracing.foldeed | inferno-flamegraph > tracing-flamegraph.svg

xmtp_api_grpc/src/lib.rs

@@ -22,7 +22,7 @@ mod utils {
             }
 
             async fn create_dev() -> Self {
-                crate::Client::create("https://grpc.dev.xmtp.network:443", false)
+                crate::Client::create("https://grpc.dev.xmtp.network:443", true)
                     .await
                     .unwrap()
             }

xmtp_mls/Cargo.toml

@@ -175,3 +175,9 @@ required-features = ["bench"]
 harness = false
 name = "crypto"
 required-features = ["bench"]
+
+[[bench]]
+harness = false
+name = "identity"
+required-features = ["bench"]
+

xmtp_mls/benches/group_limit.rs

@@ -4,13 +4,16 @@
 //! may be used to generate a flamegraph of execution from tracing logs.
 use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput};
 use std::{collections::HashMap, sync::Arc};
-use tokio::runtime::{Builder, Handle, Runtime};
+use tokio::runtime::{Builder, Runtime};
 use tracing::{trace_span, Instrument};
 use xmtp_mls::{
     builder::ClientBuilder,
     groups::GroupMetadataOptions,
     utils::{
-        bench::{create_identities_if_dont_exist, init_logging, Identity, BENCH_ROOT_SPAN},
+        bench::{
+            bench_async_setup, create_identities_if_dont_exist, init_logging, Identity,
+            BENCH_ROOT_SPAN,
+        },
         test::TestClient,
     },
     Client,
@@ -52,17 +55,6 @@ fn setup() -> (Arc<BenchClient>, Vec<Identity>, Runtime) {
     (client, identities, runtime)
 }
 
-/// criterion `batch_iter` surrounds the closure in a `Runtime.block_on` despite being a sync
-/// function, even in the async 'to_async` setup. Therefore we do this (only _slightly_) hacky
-/// workaround to allow us to async setup some groups.
-fn bench_async_setup<F, T, Fut>(fun: F) -> T
-where
-    F: Fn() -> Fut,
-    Fut: futures::future::Future<Output = T>,
-{
-    tokio::task::block_in_place(move || Handle::current().block_on(async move { fun().await }))
-}
-
 fn add_to_empty_group(c: &mut Criterion) {
     init_logging();
     let mut benchmark_group = c.benchmark_group("add_to_empty_group");

xmtp_mls/benches/identity.rs

@@ -0,0 +1,120 @@
+use crate::tracing::Instrument;
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use ethers::signers::LocalWallet;
+use tokio::runtime::{Builder, Runtime};
+use xmtp_id::{
+    associations::{
+        builder::SignatureRequest,
+        generate_inbox_id,
+        unverified::{UnverifiedRecoverableEcdsaSignature, UnverifiedSignature},
+    },
+    InboxOwner,
+};
+use xmtp_mls::utils::{bench::init_logging, test::TestClient as TestApiClient};
+use xmtp_mls::{
+    client::Client,
+    identity::IdentityStrategy,
+    utils::bench::{bench_async_setup, BENCH_ROOT_SPAN},
+};
+use xmtp_proto::api_client::XmtpTestClient;
+
+type BenchClient = Client<TestApiClient>;
+
+#[macro_use]
+extern crate tracing;
+
+fn setup() -> Runtime {
+    Builder::new_multi_thread()
+        .enable_time()
+        .enable_io()
+        .thread_name("xmtp-bencher")
+        .build()
+        .unwrap()
+}
+
+async fn new_client() -> (BenchClient, LocalWallet) {
+    let nonce = 1;
+    let wallet = xmtp_cryptography::utils::generate_local_wallet();
+    let inbox_id = generate_inbox_id(&wallet.get_address(), &nonce).unwrap();
+
+    let dev = std::env::var("DEV_GRPC");
+    let is_dev_network = matches!(dev, Ok(d) if d == "true" || d == "1");
+
+    let api_client = if is_dev_network {
+        tracing::info!("Using Dev GRPC");
+        <TestApiClient as XmtpTestClient>::create_dev().await
+    } else {
+        tracing::info!("Using Local GRPC");
+        <TestApiClient as XmtpTestClient>::create_local().await
+    };
+
+    let client = BenchClient::builder(IdentityStrategy::new(
+        inbox_id,
+        wallet.get_address(),
+        nonce,
+        None,
+    ));
+
+    let client = client
+        .temp_store()
+        .await
+        .api_client(api_client)
+        .build()
+        .await
+        .unwrap();
+
+    (client, wallet)
+}
+
+async fn ecdsa_signature(client: &BenchClient, owner: impl InboxOwner) -> SignatureRequest {
+    let mut signature_request = client.context().signature_request().unwrap();
+    let signature_text = signature_request.signature_text();
+    let unverified_signature = UnverifiedSignature::RecoverableEcdsa(
+        UnverifiedRecoverableEcdsaSignature::new(owner.sign(&signature_text).unwrap().into()),
+    );
+    signature_request
+        .add_signature(unverified_signature, client.scw_verifier())
+        .await
+        .unwrap();
+
+    signature_request
+}
+
+fn register_identity_eoa(c: &mut Criterion) {
+    init_logging();
+
+    let runtime = setup();
+
+    let mut benchmark_group = c.benchmark_group("register_identity");
+    benchmark_group.sample_size(10);
+    benchmark_group.bench_function("register_identity", |b| {
+        let span = trace_span!(BENCH_ROOT_SPAN);
+        b.to_async(&runtime).iter_batched(
+            || {
+                bench_async_setup(|| async {
+                    let (client, wallet) = new_client().await;
+                    let signature_request = ecdsa_signature(&client, wallet).await;
+
+                    (client, signature_request, span.clone())
+                })
+            },
+            |(client, request, span)| async move {
+                client
+                    .register_identity(request)
+                    .instrument(span)
+                    .await
+                    .unwrap()
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    benchmark_group.finish();
+}
+
+criterion_group!(
+    name = identity;
+    config = Criterion::default().sample_size(10);
+    targets = register_identity_eoa
+);
+criterion_main!(identity);

xmtp_mls/src/utils/bench.rs

@@ -65,6 +65,18 @@ pub fn init_logging() {
     })
 }
 
+/// criterion `batch_iter` surrounds the closure in a `Runtime.block_on` despite being a sync
+/// function, even in the async 'to_async` setup. Therefore we do this (only _slightly_) hacky
+/// workaround to allow us to async setup some groups.
+pub fn bench_async_setup<F, T, Fut>(fun: F) -> T
+where
+    F: Fn() -> Fut,
+    Fut: futures::future::Future<Output = T>,
+{
+    use tokio::runtime::Handle;
+    tokio::task::block_in_place(move || Handle::current().block_on(async move { fun().await }))
+}
+
 /// Filters for only spans where the root span name is "bench"
 pub struct BenchFilter;
 

xmtp_proto/Cargo.toml

@@ -14,6 +14,7 @@ serde = { workspace = true }
 async-trait = "0.1"
 hex.workspace = true
 openmls_rust_crypto = { workspace = true, optional = true }
+tracing.workspace = true 
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
 tonic = { workspace = true }
@@ -43,4 +44,4 @@ proto_full = ["xmtp-identity","xmtp-identity-api-v1","xmtp-identity-associations
 "xmtp-xmtpv4-envelopes" = ["xmtp-identity-associations","xmtp-mls-api-v1"]
 "xmtp-xmtpv4-message_api" = ["xmtp-xmtpv4-envelopes"]
 "xmtp-xmtpv4-payer_api" = ["xmtp-xmtpv4-envelopes"]
-## @@protoc_insertion_point(features)
+## @@protoc_insertion_point(features)

xmtp_proto/src/api_client.rs

@@ -311,11 +311,13 @@ impl<T> XmtpMlsStreams for Box<T>
 where
     T: XmtpMlsStreams + Sync + ?Sized,
 {
-    type GroupMessageStream<'a> = <T as XmtpMlsStreams>::GroupMessageStream<'a>
+    type GroupMessageStream<'a>
+        = <T as XmtpMlsStreams>::GroupMessageStream<'a>
     where
         Self: 'a;
 
-    type WelcomeMessageStream<'a> = <T as XmtpMlsStreams>::WelcomeMessageStream<'a>
+    type WelcomeMessageStream<'a>
+        = <T as XmtpMlsStreams>::WelcomeMessageStream<'a>
     where
         Self: 'a;