From 860e046501220bb87ca41471bec70648e22269a1 Mon Sep 17 00:00:00 2001
From: Alex Koshelev <akoshelev@users.noreply.github.com>
Date: Tue, 20 Feb 2024 15:16:12 -0800
Subject: [PATCH] Enable Tokio tasks profiler (#947)

* Enable Tokio tasks profiler

Requires `logging::setup()`, but otherwise works smoothly by running [`tokio-console`](https://github.com/tokio-rs/console) client that consumes metrics emitted by IPA benchmarks and shows per-task and combined stats

* Pull tokio/tracing off the main dependency graph
---
 ipa-core/Cargo.toml                  |  5 +++
 ipa-core/src/test_fixture/logging.rs | 54 ++++++++++++++++------------
 2 files changed, 37 insertions(+), 22 deletions(-)

diff --git a/ipa-core/Cargo.toml b/ipa-core/Cargo.toml
index c4097c165..033a56931 100644
--- a/ipa-core/Cargo.toml
+++ b/ipa-core/Cargo.toml
@@ -63,6 +63,10 @@ descriptive-gate = []
 compact-gate = ["ipa-macros/compact-gate"]
 # Enable using more than one thread for protocol execution. Most of the parallelism occurs at parallel/seq_join operations
 multi-threading = ["async-scoped"]
+# Enable tokio task profiling. Requires tokio_unstable flag to be passed to the compiler.
+# RUSTFLAGS="--cfg tokio_unstable" cargo run ... --features="tokio-console ...".
+# Note that if there are other flags enabled on your platform in .cargo/config.toml, you need to include them as well.
+tokio-console = ["console-subscriber", "tokio/tracing"]
 
 # Standalone aggregation protocol. We use IPA infra for communication
 # but it has nothing to do with IPA.
@@ -88,6 +92,7 @@ bytes = "1.4"
 clap = { version = "4.3.2", optional = true, features = ["derive"] }
 comfy-table = { version = "7.0", optional = true }
 config = "0.14"
+console-subscriber = {  version = "0.2", optional = true }
 criterion = { version = "0.5.1", optional = true, default-features = false, features = [
     "async_tokio",
     "plotters",
diff --git a/ipa-core/src/test_fixture/logging.rs b/ipa-core/src/test_fixture/logging.rs
index b8553eef0..8490387f8 100644
--- a/ipa-core/src/test_fixture/logging.rs
+++ b/ipa-core/src/test_fixture/logging.rs
@@ -1,10 +1,4 @@
-use std::{str::FromStr, sync::Once};
-
-use metrics_tracing_context::MetricsLayer;
-use tracing::Level;
-use tracing_subscriber::{
-    filter::Directive, fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
-};
+use std::sync::Once;
 
 /// Set up logging for IPA
 ///
@@ -14,21 +8,37 @@ pub fn setup() {
     static INIT: Once = Once::new();
 
     INIT.call_once(|| {
-        let default_directive = if let Some(crate_name) = option_env!("CARGO_CRATE_NAME") {
-            // only print IPA crate logging by default
-            Directive::from_str(&format!("{crate_name}=INFO")).unwrap()
-        } else {
-            Level::INFO.into()
-        };
+        #[cfg(feature = "tokio-console")]
+        {
+            console_subscriber::init();
+        }
+
+        #[cfg(not(feature = "tokio-console"))]
+        {
+            use std::str::FromStr;
+
+            use metrics_tracing_context::MetricsLayer;
+            use tracing::Level;
+            use tracing_subscriber::{
+                filter::Directive, fmt, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter,
+            };
+
+            let default_directive = if let Some(crate_name) = option_env!("CARGO_CRATE_NAME") {
+                // only print IPA crate logging by default
+                Directive::from_str(&format!("{crate_name}=INFO")).unwrap()
+            } else {
+                Level::INFO.into()
+            };
 
-        tracing_subscriber::registry()
-            .with(
-                EnvFilter::builder()
-                    .with_default_directive(default_directive)
-                    .from_env_lossy(),
-            )
-            .with(fmt::layer())
-            .with(MetricsLayer::new())
-            .init();
+            tracing_subscriber::registry()
+                .with(
+                    EnvFilter::builder()
+                        .with_default_directive(default_directive)
+                        .from_env_lossy(),
+                )
+                .with(fmt::layer())
+                .with(MetricsLayer::new())
+                .init();
+        }
     });
 }