diff --git a/bindings_ffi/benches/create_client.rs b/bindings_ffi/benches/create_client.rs new file mode 100644 index 000000000..a7b10709f --- /dev/null +++ b/bindings_ffi/benches/create_client.rs @@ -0,0 +1,77 @@ +use crate::tracing::Instrument; +use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; +use tokio::runtime::{Builder, Runtime}; +use xmtp_id::{ + associations::{ + builder::SignatureRequest, + unverified::{UnverifiedRecoverableEcdsaSignature, UnverifiedSignature}, + }, + InboxOwner, +}; +use xmtp_mls::utils::bench::{bench_async_setup, BenchClient, BENCH_ROOT_SPAN}; +use xmtp_mls::utils::bench::{clients, init_logging}; + +#[macro_use] +extern crate tracing; + +fn setup() -> Runtime { + Builder::new_multi_thread() + .enable_time() + .enable_io() + .thread_name("xmtp-bencher") + .build() + .unwrap() +} + +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_eoa", |b| { + let span = trace_span!(BENCH_ROOT_SPAN); + b.to_async(&runtime).iter_batched( + || { + bench_async_setup(|| async { + let (client, wallet) = clients::new_unregistered_client(false).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); diff --git a/common/Cargo.toml b/common/Cargo.toml index 5a1f774f8..d3fa43e64 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -33,3 +33,4 @@ tokio = { workspace = true, features = ["time", "macros", "rt-multi-thread", "sy [features] test-utils = ["dep:parking_lot", "dep:tracing-subscriber", "dep:tracing-wasm", "dep:console_error_panic_hook"] +bench = ["test-utils"] diff --git a/common/src/bench.rs b/common/src/bench.rs new file mode 100644 index 000000000..76b04374a --- /dev/null +++ b/common/src/bench.rs @@ -0,0 +1,64 @@ +static INIT: Once = Once::new(); + +static LOGGER: OnceCell>> = OnceCell::new(); + +pub const BENCH_ROOT_SPAN: &str = "xmtp-trace-bench"; + +/// initializes logging for benchmarks +/// - FMT logging is enabled by passing the normal `RUST_LOG` environment variable options. +/// - Generate a flamegraph from tracing data by passing `XMTP_FLAMEGRAPH=trace` +pub fn logger() { + INIT.call_once(|| { + let (flame_layer, guard) = FlameLayer::with_file("./tracing.folded").unwrap(); + let flame_layer = flame_layer + .with_threads_collapsed(true) + .with_module_path(true); + // .with_empty_samples(false); + + tracing_subscriber::registry() + .with(tracing_subscriber::fmt::layer().with_filter(EnvFilter::from_default_env())) + .with( + flame_layer + .with_filter(BenchFilter) + .with_filter(EnvFilter::from_env("XMTP_FLAMEGRAPH")), + ) + .init(); + + LOGGER.set(guard).unwrap(); + }) +} + +/// 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(fun: F) -> T +where + F: Fn() -> Fut, + Fut: futures::future::Future, +{ + 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; + +impl Filter for BenchFilter +where + S: Subscriber + for<'lookup> LookupSpan<'lookup> + std::fmt::Debug, + for<'lookup> >::Data: std::fmt::Debug, +{ + fn enabled(&self, meta: &Metadata<'_>, cx: &Context<'_, S>) -> bool { + if meta.name() == BENCH_ROOT_SPAN { + return true; + } + if let Some(id) = cx.current_span().id() { + if let Some(s) = cx.span_scope(id) { + if let Some(s) = s.from_root().take(1).collect::>().first() { + return s.name() == BENCH_ROOT_SPAN; + } + } + } + false + } +} diff --git a/common/src/lib.rs b/common/src/lib.rs index bc6664738..ec350386e 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -7,6 +7,9 @@ mod test; #[cfg(feature = "test-utils")] pub use test::*; +#[cfg(feature = "bench")] +pub mod bench; + pub mod retry; pub use retry::*; diff --git a/xmtp_mls/benches/group_limit.rs b/xmtp_mls/benches/group_limit.rs index ba3027eb3..ac2be6a5c 100755 --- a/xmtp_mls/benches/group_limit.rs +++ b/xmtp_mls/benches/group_limit.rs @@ -6,13 +6,11 @@ use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criteri use std::{collections::HashMap, sync::Arc}; use tokio::runtime::{Builder, Runtime}; use tracing::{trace_span, Instrument}; +use xmtp_common::bench::{self, bench_async_setup, BENCH_ROOT_SPAN}; use xmtp_mls::{ builder::ClientBuilder, groups::GroupMetadataOptions, - utils::bench::{ - bench_async_setup, create_identities_if_dont_exist, init_logging, BenchClient, Identity, - BENCH_ROOT_SPAN, - }, + utils::bench::{create_identities_if_dont_exist, BenchClient, Identity}, }; pub const IDENTITY_SAMPLES: [usize; 9] = [10, 20, 40, 80, 100, 200, 300, 400, 450]; @@ -50,7 +48,7 @@ fn setup() -> (Arc, Vec, Runtime) { } fn add_to_empty_group(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("add_to_empty_group"); benchmark_group.sample_size(SAMPLE_SIZE); @@ -88,7 +86,7 @@ fn add_to_empty_group(c: &mut Criterion) { } fn add_to_empty_group_by_inbox_id(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("add_to_empty_group_by_inbox_id"); benchmark_group.sample_size(SAMPLE_SIZE); @@ -130,7 +128,7 @@ fn add_to_empty_group_by_inbox_id(c: &mut Criterion) { } fn add_to_100_member_group_by_inbox_id(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("add_to_100_member_group_by_inbox_id"); benchmark_group.sample_size(SAMPLE_SIZE); @@ -188,7 +186,7 @@ fn add_to_100_member_group_by_inbox_id(c: &mut Criterion) { } fn remove_all_members_from_group(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("remove_all_members_from_group"); benchmark_group.sample_size(SAMPLE_SIZE); @@ -233,7 +231,7 @@ fn remove_all_members_from_group(c: &mut Criterion) { } fn remove_half_members_from_group(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("remove_half_members_from_group"); benchmark_group.sample_size(SAMPLE_SIZE); @@ -281,7 +279,7 @@ fn remove_half_members_from_group(c: &mut Criterion) { } fn add_1_member_to_group(c: &mut Criterion) { - init_logging(); + bench::logger(); let mut benchmark_group = c.benchmark_group("add_1_member_to_group"); benchmark_group.sample_size(SAMPLE_SIZE); diff --git a/xmtp_mls/benches/identity.rs b/xmtp_mls/benches/identity.rs index a7b10709f..8bb9de994 100644 --- a/xmtp_mls/benches/identity.rs +++ b/xmtp_mls/benches/identity.rs @@ -1,6 +1,7 @@ use crate::tracing::Instrument; use criterion::{criterion_group, criterion_main, BatchSize, Criterion}; use tokio::runtime::{Builder, Runtime}; +use xmtp_common::bench::{self, bench_async_setup, BenchClient, BENCH_ROOT_SPAN}; use xmtp_id::{ associations::{ builder::SignatureRequest, @@ -8,8 +9,7 @@ use xmtp_id::{ }, InboxOwner, }; -use xmtp_mls::utils::bench::{bench_async_setup, BenchClient, BENCH_ROOT_SPAN}; -use xmtp_mls::utils::bench::{clients, init_logging}; +use xmtp_mls::utils::bench::clients; #[macro_use] extern crate tracing; @@ -38,7 +38,7 @@ async fn ecdsa_signature(client: &BenchClient, owner: impl InboxOwner) -> Signat } fn register_identity_eoa(c: &mut Criterion) { - init_logging(); + bench::logger(); let runtime = setup(); diff --git a/xmtp_mls/src/utils/bench/mod.rs b/xmtp_mls/src/utils/bench/mod.rs index 5466c6dd9..ba2ac4894 100644 --- a/xmtp_mls/src/utils/bench/mod.rs +++ b/xmtp_mls/src/utils/bench/mod.rs @@ -20,8 +20,6 @@ use tracing_subscriber::{ EnvFilter, }; -pub const BENCH_ROOT_SPAN: &str = "xmtp-trace-bench"; - /// Re-export of functions in private modules for benchmarks pub mod re_export { pub use crate::hpke::encrypt_welcome; @@ -35,65 +33,4 @@ pub enum BenchError { Io(#[from] std::io::Error), } -static INIT: Once = Once::new(); - -static LOGGER: OnceCell>> = OnceCell::new(); - -/// initializes logging for benchmarks -/// - FMT logging is enabled by passing the normal `RUST_LOG` environment variable options. -/// - Generate a flamegraph from tracing data by passing `XMTP_FLAMEGRAPH=trace` -pub fn init_logging() { - INIT.call_once(|| { - let (flame_layer, guard) = FlameLayer::with_file("./tracing.folded").unwrap(); - let flame_layer = flame_layer - .with_threads_collapsed(true) - .with_module_path(true); - // .with_empty_samples(false); - - tracing_subscriber::registry() - .with(tracing_subscriber::fmt::layer().with_filter(EnvFilter::from_default_env())) - .with( - flame_layer - .with_filter(BenchFilter) - .with_filter(EnvFilter::from_env("XMTP_FLAMEGRAPH")), - ) - .init(); - - LOGGER.set(guard).unwrap(); - }) -} - -/// 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(fun: F) -> T -where - F: Fn() -> Fut, - Fut: futures::future::Future, -{ - 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; - -impl Filter for BenchFilter -where - S: Subscriber + for<'lookup> LookupSpan<'lookup> + std::fmt::Debug, - for<'lookup> >::Data: std::fmt::Debug, -{ - fn enabled(&self, meta: &Metadata<'_>, cx: &Context<'_, S>) -> bool { - if meta.name() == BENCH_ROOT_SPAN { - return true; - } - if let Some(id) = cx.current_span().id() { - if let Some(s) = cx.span_scope(id) { - if let Some(s) = s.from_root().take(1).collect::>().first() { - return s.name() == BENCH_ROOT_SPAN; - } - } - } - false - } -} +pub use xmtp_common::bench::logger;