Some quick benching

Sushisource · Sushisource · commit 1d4bafbec6be · 2025-06-24T15:59:36.000-07:00
diff --git a/core/Cargo.toml b/core/Cargo.toml
@@ -91,7 +91,7 @@ path = "../fsm"
 assert_matches = "1.4"
 bimap = "0.6.1"
 clap = { version = "4.0", features = ["derive"] }
-criterion = "0.5"
+criterion = { version = "0.6", features = ["async", "async_tokio"] }
 rstest = "0.25"
 temporal-sdk-core-test-utils = { path = "../test-utils" }
 temporal-sdk = { path = "../sdk" }
diff --git a/core/benches/workflow_replay.rs b/core/benches/workflow_replay.rs
@@ -1,17 +1,27 @@
-use criterion::{Criterion, criterion_group, criterion_main};
+use criterion::{BatchSize, Criterion, criterion_group, criterion_main};
 use futures_util::StreamExt;
-use std::time::Duration;
+use std::{
+    sync::{Arc, mpsc},
+    thread,
+    time::Duration,
+};
 use temporal_sdk::{WfContext, WorkflowFunction};
-use temporal_sdk_core::replay::HistoryForReplay;
+use temporal_sdk_core::{CoreRuntime, replay::HistoryForReplay};
+use temporal_sdk_core_api::telemetry::metrics::{
+    MetricKeyValue, MetricParametersBuilder, NewAttributes,
+};
 use temporal_sdk_core_protos::DEFAULT_WORKFLOW_TYPE;
-use temporal_sdk_core_test_utils::{canned_histories, replay_sdk_worker};
+use temporal_sdk_core_test_utils::{
+    DONT_AUTO_INIT_INTEG_TELEM, canned_histories, prom_metrics, replay_sdk_worker,
+};
 
 pub fn criterion_benchmark(c: &mut Criterion) {
     let tokio_runtime = tokio::runtime::Builder::new_current_thread()
         .enable_time()
         .build()
         .unwrap();
     let _g = tokio_runtime.enter();
+    DONT_AUTO_INIT_INTEG_TELEM.set(true);
 
     let num_timers = 10;
     let t = canned_histories::long_sequential_timers(num_timers as usize);
@@ -21,14 +31,17 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     );
 
     c.bench_function("Small history replay", |b| {
-        b.iter(|| {
-            tokio_runtime.block_on(async {
+        b.to_async(&tokio_runtime).iter_batched(
+            || {
                 let func = timers_wf(num_timers);
-                let mut worker = replay_sdk_worker([hist.clone()]);
+                (func, replay_sdk_worker([hist.clone()]))
+            },
+            |(func, mut worker)| async move {
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
                 worker.run().await.unwrap();
-            })
-        })
+            },
+            BatchSize::SmallInput,
+        )
     });
 
     let num_tasks = 50;
@@ -39,18 +52,104 @@ pub fn criterion_benchmark(c: &mut Criterion) {
     );
 
     c.bench_function("Large payloads history replay", |b| {
-        b.iter(|| {
-            tokio_runtime.block_on(async {
+        b.to_async(&tokio_runtime).iter_batched(
+            || {
                 let func = big_signals_wf(num_tasks);
-                let mut worker = replay_sdk_worker([hist.clone()]);
+                (func, replay_sdk_worker([hist.clone()]))
+            },
+            |(func, mut worker)| async move {
                 worker.register_wf(DEFAULT_WORKFLOW_TYPE, func);
                 worker.run().await.unwrap();
-            })
-        })
+            },
+            BatchSize::SmallInput,
+        )
     });
 }
 
-criterion_group!(benches, criterion_benchmark);
+pub fn bench_metrics(c: &mut Criterion) {
+    DONT_AUTO_INIT_INTEG_TELEM.set(true);
+    let tokio_runtime = tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .unwrap();
+    let _tokio = tokio_runtime.enter();
+    let (mut telemopts, addr, _aborter) = prom_metrics(None);
+    telemopts.logging = None;
+    let rt = CoreRuntime::new_assume_tokio(telemopts).unwrap();
+    let meter = rt.telemetry().get_metric_meter().unwrap();
+
+    c.bench_function("Record with new attributes on each call", move |b| {
+        b.iter_batched(
+            || {
+                let c = meter.counter(
+                    MetricParametersBuilder::default()
+                        .name("c")
+                        .build()
+                        .unwrap(),
+                );
+                let h = meter.histogram(
+                    MetricParametersBuilder::default()
+                        .name("h")
+                        .build()
+                        .unwrap(),
+                );
+                let g = meter.gauge(
+                    MetricParametersBuilder::default()
+                        .name("g")
+                        .build()
+                        .unwrap(),
+                );
+
+                let vals = [1, 2, 3, 4, 5];
+                let labels = ["l1", "l2"];
+
+                let (start_tx, start_rx) = mpsc::channel();
+                let start_rx = Arc::new(std::sync::Mutex::new(start_rx));
+
+                let mut thread_handles = Vec::new();
+                for _ in 0..3 {
+                    let c = c.clone();
+                    let h = h.clone();
+                    let g = g.clone();
+                    let meter = meter.clone();
+                    let start_rx = start_rx.clone();
+
+                    let handle = thread::spawn(move || {
+                        // Wait for start signal
+                        let _ = start_rx.lock().unwrap().recv();
+
+                        for _ in 1..=100 {
+                            for &val in &vals {
+                                for &label in &labels {
+                                    let attribs = meter.new_attributes(NewAttributes::from(vec![
+                                        MetricKeyValue::new("label", label),
+                                    ]));
+                                    c.add(val, &attribs);
+                                    h.record(val, &attribs);
+                                    g.record(val, &attribs);
+                                }
+                            }
+                        }
+                    });
+                    thread_handles.push(handle);
+                }
+
+                (start_tx, thread_handles)
+            },
+            |(start_tx, thread_handles)| {
+                for _ in 0..3 {
+                    let _ = start_tx.send(());
+                }
+                for handle in thread_handles {
+                    let _ = handle.join();
+                }
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+criterion_group!(benches, criterion_benchmark, bench_metrics);
 criterion_main!(benches);
 
 fn timers_wf(num_timers: u32) -> WorkflowFunction {
diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs
@@ -18,6 +18,7 @@ use parking_lot::Mutex;
 use prost::Message;
 use rand::Rng;
 use std::{
+    cell::Cell,
     convert::TryFrom,
     env,
     future::Future,
@@ -171,8 +172,15 @@ pub async fn history_from_proto_binary(path_from_root: &str) -> Result<History,
     Ok(History::decode(&*bytes)?)
 }
 
+thread_local! {
+    /// Can be set true to disable auto-initialization of integ-test telemetry.
+    pub static DONT_AUTO_INIT_INTEG_TELEM: Cell<bool> = const { Cell::new(false) };
+}
 static INTEG_TESTS_RT: std::sync::OnceLock<CoreRuntime> = std::sync::OnceLock::new();
-pub fn init_integ_telem() -> &'static CoreRuntime {
+pub fn init_integ_telem() {
+    if DONT_AUTO_INIT_INTEG_TELEM.get() {
+        return;
+    }
     INTEG_TESTS_RT.get_or_init(|| {
         let telemetry_options = get_integ_telem_options();
         let rt =
@@ -181,7 +189,7 @@ pub fn init_integ_telem() -> &'static CoreRuntime {
             let _ = tracing::subscriber::set_global_default(sub);
         }
         rt
-    })
+    });
 }
 
 /// Implements a builder pattern to help integ tests initialize core and create workflows
diff --git a/tests/heavy_tests.rs b/tests/heavy_tests.rs
@@ -12,13 +12,15 @@ use std::{
 };
 use temporal_client::{GetWorkflowResultOpts, WfClientExt, WorkflowClientTrait, WorkflowOptions};
 use temporal_sdk::{ActContext, ActivityOptions, WfContext, WorkflowResult};
-use temporal_sdk_core::{ResourceBasedTuner, ResourceSlotOptions};
+use temporal_sdk_core::{CoreRuntime, ResourceBasedTuner, ResourceSlotOptions};
 use temporal_sdk_core_api::worker::PollerBehavior;
 use temporal_sdk_core_protos::{
     coresdk::{AsJsonPayloadExt, workflow_commands::ActivityCancellationType},
     temporal::api::enums::v1::WorkflowIdReusePolicy,
 };
-use temporal_sdk_core_test_utils::{CoreWfStarter, rand_6_chars, workflows::la_problem_workflow};
+use temporal_sdk_core_test_utils::{
+    CoreWfStarter, prom_metrics, rand_6_chars, workflows::la_problem_workflow,
+};
 
 mod fuzzy_workflow;
 
@@ -183,7 +185,9 @@ async fn workflow_load() {
     const SIGNAME: &str = "signame";
     let num_workflows = 500;
     let wf_name = "workflow_load";
-    let mut starter = CoreWfStarter::new("workflow_load");
+    let (telemopts, _, _aborter) = prom_metrics(None);
+    let rt = CoreRuntime::new_assume_tokio(telemopts).unwrap();
+    let mut starter = CoreWfStarter::new_with_runtime("workflow_load", rt);
     starter
         .worker_config
         .max_outstanding_workflow_tasks(5_usize)
