From e5c0c7dc4dc8bf12c412d489727b41e9e6eb2291 Mon Sep 17 00:00:00 2001
From: Arseny Sher <sher-ars@yandex.ru>
Date: Mon, 21 Nov 2022 10:37:00 +0400
Subject: [PATCH] Add tokio-console tracing as optional feature.

---
 Cargo.lock                    | 62 +++++++++++++++++++++++++++++++++++
 broker/Cargo.toml             |  5 +++
 broker/src/bin/neon_broker.rs |  7 ++--
 3 files changed, 72 insertions(+), 2 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index ffb9bb124c..a3290b81b0 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -569,6 +569,42 @@ dependencies = [
  "workspace_hack",
 ]
 
+[[package]]
+name = "console-api"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e57ff02e8ad8e06ab9731d5dc72dc23bef9200778eae1a89d555d8c42e5d4a86"
+dependencies = [
+ "prost 0.11.2",
+ "prost-types 0.11.2",
+ "tonic 0.8.2",
+ "tracing-core",
+]
+
+[[package]]
+name = "console-subscriber"
+version = "0.1.8"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "22a3a81dfaf6b66bce5d159eddae701e3a002f194d378cbf7be5f053c281d9be"
+dependencies = [
+ "console-api",
+ "crossbeam-channel",
+ "crossbeam-utils",
+ "futures",
+ "hdrhistogram",
+ "humantime",
+ "prost-types 0.11.2",
+ "serde",
+ "serde_json",
+ "thread_local",
+ "tokio",
+ "tokio-stream",
+ "tonic 0.8.2",
+ "tracing",
+ "tracing-core",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "const_format"
 version = "0.2.30"
@@ -1102,6 +1138,16 @@ version = "0.4.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "0ce7134b9999ecaf8bcd65542e436736ef32ddca1b3e06094cb6ec5755203b80"
 
+[[package]]
+name = "flate2"
+version = "1.0.24"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f82b0f4c27ad9f8bfd1f3208d882da2b09c301bc1c828fd3a00d0216d2fbbff6"
+dependencies = [
+ "crc32fast",
+ "miniz_oxide",
+]
+
 [[package]]
 name = "fnv"
 version = "1.0.7"
@@ -1344,6 +1390,19 @@ dependencies = [
  "ahash",
 ]
 
+[[package]]
+name = "hdrhistogram"
+version = "7.5.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7f19b9f54f7c7f55e31401bb647626ce0cf0f67b0004982ce815b3ee72a02aa8"
+dependencies = [
+ "base64",
+ "byteorder",
+ "flate2",
+ "nom",
+ "num-traits",
+]
+
 [[package]]
 name = "heapless"
 version = "0.7.16"
@@ -1912,6 +1971,7 @@ dependencies = [
  "async-stream",
  "bytes",
  "clap 4.0.15",
+ "console-subscriber",
  "futures-core",
  "futures-util",
  "git-version",
@@ -3790,10 +3850,12 @@ dependencies = [
  "mio",
  "num_cpus",
  "once_cell",
+ "parking_lot 0.12.1",
  "pin-project-lite",
  "signal-hook-registry",
  "socket2",
  "tokio-macros",
+ "tracing",
  "winapi",
 ]
 
diff --git a/broker/Cargo.toml b/broker/Cargo.toml
index 395ab8a275..8469768126 100644
--- a/broker/Cargo.toml
+++ b/broker/Cargo.toml
@@ -5,6 +5,9 @@ edition = "2021"
 
 [features]
 bench = []
+# for exploring with tokio-console. Note that tokio_unstable cfg must be enabled, i.e.
+# RUSTFLAGS="--cfg tokio_unstable" cargo build -r -p neon_broker --features console
+console = ["dep:console-subscriber", "tokio/full", "tokio/tracing"]
 
 [[bin]]
 name = "neon_broker_bench"
@@ -29,6 +32,8 @@ tokio = { version = "1.0", features = ["macros", "rt-multi-thread"] }
 tokio-stream = "0.1"
 tracing = "0.1.27"
 
+console-subscriber = {version = "0.1.8", optional = true }
+
 metrics = { path = "../libs/metrics" }
 utils = { path = "../libs/utils" }
 workspace_hack = { version = "0.1", path = "../workspace_hack" }
diff --git a/broker/src/bin/neon_broker.rs b/broker/src/bin/neon_broker.rs
index c2cea84c06..f144234f74 100644
--- a/broker/src/bin/neon_broker.rs
+++ b/broker/src/bin/neon_broker.rs
@@ -41,7 +41,6 @@ use neon_broker::neon_broker_proto::subscribe_safekeeper_info_request::Subscript
 use neon_broker::neon_broker_proto::{SafekeeperTimelineInfo, SubscribeSafekeeperInfoRequest};
 use neon_broker::{parse_proto_ttid, EitherBody, DEFAULT_LISTEN_ADDR};
 use utils::id::TenantTimelineId;
-use utils::logging::{self, LogFormat};
 use utils::project_git_version;
 
 project_git_version!(GIT_VERSION);
@@ -380,9 +379,13 @@ async fn http1_handler(
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args = Args::parse();
 
-    logging::init(LogFormat::from_config(&args.log_format)?)?;
+    #[cfg(not(feature = "console"))] // tokio-console inits tracing_subscriber on its own
+    utils::logging::init(utils::logging::LogFormat::from_config(&args.log_format)?)?;
     info!("version: {GIT_VERSION}");
 
+    #[cfg(feature = "console")]
+    console_subscriber::init();
+
     let registry = Registry {
         shared_state: Arc::new(RwLock::new(SharedState::new(args.chan_size))),
         chan_size: args.chan_size,