From 29a7b753dc701f68836b61fa6187aafba5bbc47e Mon Sep 17 00:00:00 2001
From: Shaun Cox <shaunco@microsoft.com>
Date: Tue, 4 Jul 2023 09:28:18 -0500
Subject: [PATCH] Add benchmark for Context

---
 opentelemetry-sdk/Cargo.toml         |  4 ++
 opentelemetry-sdk/benches/context.rs | 94 ++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)
 create mode 100644 opentelemetry-sdk/benches/context.rs

diff --git a/opentelemetry-sdk/Cargo.toml b/opentelemetry-sdk/Cargo.toml
index 993afe5454..ab5e18a96c 100644
--- a/opentelemetry-sdk/Cargo.toml
+++ b/opentelemetry-sdk/Cargo.toml
@@ -51,6 +51,10 @@ rt-tokio = ["tokio", "tokio-stream"]
 rt-tokio-current-thread = ["tokio", "tokio-stream"]
 rt-async-std = ["async-std"]
 
+[[bench]]
+name = "context"
+harness = false
+
 [[bench]]
 name = "key_value_map"
 harness = false
diff --git a/opentelemetry-sdk/benches/context.rs b/opentelemetry-sdk/benches/context.rs
new file mode 100644
index 0000000000..eb15a8e3e0
--- /dev/null
+++ b/opentelemetry-sdk/benches/context.rs
@@ -0,0 +1,94 @@
+use std::fmt::Display;
+
+use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion};
+use futures_util::future::BoxFuture;
+use opentelemetry_api::{
+    trace::{TraceContextExt, Tracer, TracerProvider},
+    Context,
+};
+use opentelemetry_sdk::{
+    export::trace::{ExportResult, SpanData, SpanExporter},
+    trace as sdktrace,
+};
+use pprof::criterion::{Output, PProfProfiler};
+
+fn criterion_benchmark(c: &mut Criterion) {
+    benchmark_group(c, BenchmarkParameter::NoActiveSpan);
+    benchmark_group(c, BenchmarkParameter::WithActiveSpan);
+}
+
+fn benchmark_group(c: &mut Criterion, p: BenchmarkParameter) {
+    let _guard = match p {
+        BenchmarkParameter::NoActiveSpan => None,
+        BenchmarkParameter::WithActiveSpan => {
+            let (provider, tracer) = tracer();
+            let guard = Context::current_with_span(tracer.start("span")).attach();
+            Some((guard, provider))
+        }
+    };
+
+    let mut group = c.benchmark_group("context");
+
+    group.bench_function(BenchmarkId::new("baseline current()", p), |b| {
+        b.iter(|| {
+            black_box(Context::current());
+        })
+    });
+
+    group.bench_function(BenchmarkId::new("current().has_active_span()", p), |b| {
+        b.iter(|| {
+            black_box(Context::current().has_active_span());
+        })
+    });
+
+    group.bench_function(
+        BenchmarkId::new("map_current(|cx| cx.has_active_span())", p),
+        |b| {
+            b.iter(|| {
+                black_box(Context::map_current(|cx| cx.has_active_span()));
+            })
+        },
+    );
+
+    group.finish();
+}
+
+#[derive(Copy, Clone)]
+enum BenchmarkParameter {
+    NoActiveSpan,
+    WithActiveSpan,
+}
+
+impl Display for BenchmarkParameter {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match *self {
+            BenchmarkParameter::NoActiveSpan => write!(f, "no-active-span"),
+            BenchmarkParameter::WithActiveSpan => write!(f, "with-active-span"),
+        }
+    }
+}
+
+fn tracer() -> (sdktrace::TracerProvider, sdktrace::Tracer) {
+    let provider = sdktrace::TracerProvider::builder()
+        .with_config(sdktrace::config().with_sampler(sdktrace::Sampler::AlwaysOn))
+        .with_simple_exporter(NoopExporter)
+        .build();
+    let tracer = provider.tracer(module_path!());
+    (provider, tracer)
+}
+
+#[derive(Debug)]
+struct NoopExporter;
+
+impl SpanExporter for NoopExporter {
+    fn export(&mut self, _spans: Vec<SpanData>) -> BoxFuture<'static, ExportResult> {
+        Box::pin(futures_util::future::ready(Ok(())))
+    }
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
+    targets = criterion_benchmark
+}
+criterion_main!(benches);