Skip to content

Commit

Permalink
benchmarks for create_client
Browse files Browse the repository at this point in the history
  • Loading branch information
insipx committed Dec 11, 2024
1 parent 3df89c2 commit 8067f17
Show file tree
Hide file tree
Showing 7 changed files with 157 additions and 77 deletions.
77 changes: 77 additions & 0 deletions bindings_ffi/benches/create_client.rs
Original file line number Diff line number Diff line change
@@ -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);
1 change: 1 addition & 0 deletions common/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -36,3 +36,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"]
64 changes: 64 additions & 0 deletions common/src/bench.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,64 @@
static INIT: Once = Once::new();

static LOGGER: OnceCell<FlushGuard<std::io::BufWriter<std::fs::File>>> = 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<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;

impl<S> Filter<S> for BenchFilter
where
S: Subscriber + for<'lookup> LookupSpan<'lookup> + std::fmt::Debug,
for<'lookup> <S as LookupSpan<'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::<Vec<_>>().first() {
return s.name() == BENCH_ROOT_SPAN;
}
}
}
false
}
}
3 changes: 3 additions & 0 deletions common/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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::*;

Expand Down
18 changes: 8 additions & 10 deletions xmtp_mls/benches/group_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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];
Expand Down Expand Up @@ -50,7 +48,7 @@ fn setup() -> (Arc<BenchClient>, Vec<Identity>, 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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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);

Expand Down
6 changes: 3 additions & 3 deletions xmtp_mls/benches/identity.rs
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
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,
unverified::{UnverifiedRecoverableEcdsaSignature, UnverifiedSignature},
},
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;
Expand Down Expand Up @@ -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();

Expand Down
65 changes: 1 addition & 64 deletions xmtp_mls/src/utils/bench/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -35,65 +33,4 @@ pub enum BenchError {
Io(#[from] std::io::Error),
}

static INIT: Once = Once::new();

static LOGGER: OnceCell<FlushGuard<std::io::BufWriter<std::fs::File>>> = 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<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;

impl<S> Filter<S> for BenchFilter
where
S: Subscriber + for<'lookup> LookupSpan<'lookup> + std::fmt::Debug,
for<'lookup> <S as LookupSpan<'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::<Vec<_>>().first() {
return s.name() == BENCH_ROOT_SPAN;
}
}
}
false
}
}
pub use xmtp_common::bench::logger;

0 comments on commit 8067f17

Please sign in to comment.