From 68e75208249f3ae61b5287f4b053ca7c7eb264f4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E8=97=8D+85CD?= <50108258+kwaa@users.noreply.github.com>
Date: Sun, 26 May 2024 23:00:42 +0800
Subject: [PATCH] feat: support `tokio-console` (#40)

* feat: support `tokio-console`

* chore: rename feature to `tokio-console`
---
 .cargo/config.toml |  13 +++-
 Cargo.lock         | 165 +++++++++++++++++++++++++++++++++++++++++++--
 Cargo.toml         |   7 ++
 src/main.rs        |  10 ++-
 4 files changed, 183 insertions(+), 12 deletions(-)

diff --git a/.cargo/config.toml b/.cargo/config.toml
index a4dbdeb..834c69b 100644
--- a/.cargo/config.toml
+++ b/.cargo/config.toml
@@ -5,10 +5,17 @@ lld = true
 # [target.x86_64-unknown-linux-gnu]
 # rustflags = ["-C", "linker=clang", "-C", "link-arg=-fuse-ld=lld"]
 
-# Enable UUID Unstable Features
-# https://docs.rs/uuid/1.4.0/uuid/index.html#unstable-features
 [target.'cfg(all())']
-rustflags = ["--cfg", "uuid_unstable"]
+rustflags = [
+    # UUID Unstable for Uuid::now_v7()
+    # https://docs.rs/uuid/1.4.0/uuid/index.html#unstable-features
+    "--cfg",
+    "uuid_unstable",
+    # Tokio Unstable for `console` feature
+    # https://github.com/tokio-rs/console#instrumenting-your-program
+    "--cfg",
+    "tokio_unstable",
+]
 
 # x86_64-v3 (Intel Haswell / AMD Excavator and above)
 # https://en.wikipedia.org/wiki/X86-64#Microarchitecture_levels
diff --git a/Cargo.lock b/Cargo.lock
index d2e3a5e..ff4f7d8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -614,6 +614,43 @@ dependencies = [
  "crossbeam-utils",
 ]
 
+[[package]]
+name = "console-api"
+version = "0.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fd326812b3fd01da5bb1af7d340d0d555fd3d4b641e7f1dfcf5962a902952787"
+dependencies = [
+ "futures-core",
+ "prost",
+ "prost-types",
+ "tonic",
+ "tracing-core",
+]
+
+[[package]]
+name = "console-subscriber"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7481d4c57092cd1c19dd541b92bdce883de840df30aa5d03fd48a3935c01842e"
+dependencies = [
+ "console-api",
+ "crossbeam-channel",
+ "crossbeam-utils",
+ "futures-task",
+ "hdrhistogram",
+ "humantime",
+ "prost-types",
+ "serde",
+ "serde_json",
+ "thread_local",
+ "tokio",
+ "tokio-stream",
+ "tonic",
+ "tracing",
+ "tracing-core",
+ "tracing-subscriber",
+]
+
 [[package]]
 name = "const-oid"
 version = "0.9.6"
@@ -1287,7 +1324,7 @@ dependencies = [
  "futures-sink",
  "futures-util",
  "http 0.2.12",
- "indexmap",
+ "indexmap 2.2.6",
  "slab",
  "tokio",
  "tokio-util",
@@ -1327,6 +1364,7 @@ name = "hatsu"
 version = "0.2.0-beta.8"
 dependencies = [
  "activitypub_federation",
+ "console-subscriber",
  "dotenvy",
  "hatsu_api_apub",
  "hatsu_apub",
@@ -1557,6 +1595,19 @@ dependencies = [
  "utoipa",
 ]
 
+[[package]]
+name = "hdrhistogram"
+version = "7.5.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d"
+dependencies = [
+ "base64 0.21.7",
+ "byteorder",
+ "flate2",
+ "nom",
+ "num-traits",
+]
+
 [[package]]
 name = "headers"
 version = "0.3.9"
@@ -1765,6 +1816,12 @@ dependencies = [
  "uuid",
 ]
 
+[[package]]
+name = "humantime"
+version = "2.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
+
 [[package]]
 name = "hyper"
 version = "0.14.28"
@@ -1789,6 +1846,18 @@ dependencies = [
  "want",
 ]
 
+[[package]]
+name = "hyper-timeout"
+version = "0.4.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bbb958482e8c7be4bc3cf272a766a2b0bf1a6755e7a6ae777f017a31d11b13b1"
+dependencies = [
+ "hyper",
+ "pin-project-lite",
+ "tokio",
+ "tokio-io-timeout",
+]
+
 [[package]]
 name = "hyper-tls"
 version = "0.5.0"
@@ -1841,6 +1910,16 @@ dependencies = [
  "unicode-normalization",
 ]
 
+[[package]]
+name = "indexmap"
+version = "1.9.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99"
+dependencies = [
+ "autocfg",
+ "hashbrown 0.12.3",
+]
+
 [[package]]
 name = "indexmap"
 version = "2.2.6"
@@ -2620,6 +2699,38 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "prost"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "deb1435c188b76130da55f17a466d252ff7b1418b2ad3e037d127b94e3411f29"
+dependencies = [
+ "bytes",
+ "prost-derive",
+]
+
+[[package]]
+name = "prost-derive"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "81bddcdb20abf9501610992b6759a4c888aef7d1a7247ef75e2404275ac24af1"
+dependencies = [
+ "anyhow",
+ "itertools",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.65",
+]
+
+[[package]]
+name = "prost-types"
+version = "0.12.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9091c90b0a32608e984ff2fa4091273cbdd755d54935c51d520887f4a1dbd5b0"
+dependencies = [
+ "prost",
+]
+
 [[package]]
 name = "ptr_meta"
 version = "0.1.4"
@@ -3333,7 +3444,7 @@ version = "1.0.117"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "455182ea6142b14f93f4bc5320a2b31c1f266b66a4a5c858b013302a5d8cbfc3"
 dependencies = [
- "indexmap",
+ "indexmap 2.2.6",
  "itoa",
  "ryu",
  "serde",
@@ -3571,7 +3682,7 @@ dependencies = [
  "futures-util",
  "hashlink",
  "hex",
- "indexmap",
+ "indexmap 2.2.6",
  "log",
  "memchr",
  "once_cell",
@@ -4015,6 +4126,7 @@ dependencies = [
  "signal-hook-registry",
  "socket2",
  "tokio-macros",
+ "tracing",
  "windows-sys 0.48.0",
 ]
 
@@ -4050,6 +4162,16 @@ dependencies = [
  "tracing",
 ]
 
+[[package]]
+name = "tokio-io-timeout"
+version = "1.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "30b74022ada614a1b4834de765f9bb43877f910cc8ce4be40e89042c9223a8bf"
+dependencies = [
+ "pin-project-lite",
+ "tokio",
+]
+
 [[package]]
 name = "tokio-macros"
 version = "2.2.0"
@@ -4122,7 +4244,7 @@ version = "0.21.1"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "6a8534fd7f78b5405e860340ad6575217ce99f38d4d5c8f2442cb5ecb50090e1"
 dependencies = [
- "indexmap",
+ "indexmap 2.2.6",
  "toml_datetime",
  "winnow",
 ]
@@ -4133,12 +4255,39 @@ version = "0.22.13"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c127785850e8c20836d49732ae6abfa47616e60bf9d9f57c43c250361a9db96c"
 dependencies = [
- "indexmap",
+ "indexmap 2.2.6",
  "serde",
  "serde_spanned",
  "toml_datetime",
 ]
 
+[[package]]
+name = "tonic"
+version = "0.10.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d560933a0de61cf715926b9cac824d4c883c2c43142f787595e48280c40a1d0e"
+dependencies = [
+ "async-stream",
+ "async-trait",
+ "axum",
+ "base64 0.21.7",
+ "bytes",
+ "h2",
+ "http 0.2.12",
+ "http-body 0.4.6",
+ "hyper",
+ "hyper-timeout",
+ "percent-encoding",
+ "pin-project",
+ "prost",
+ "tokio",
+ "tokio-stream",
+ "tower",
+ "tower-layer",
+ "tower-service",
+ "tracing",
+]
+
 [[package]]
 name = "tower"
 version = "0.4.13"
@@ -4147,9 +4296,13 @@ checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c"
 dependencies = [
  "futures-core",
  "futures-util",
+ "indexmap 1.9.3",
  "pin-project",
  "pin-project-lite",
+ "rand",
+ "slab",
  "tokio",
+ "tokio-util",
  "tower-layer",
  "tower-service",
  "tracing",
@@ -4385,7 +4538,7 @@ version = "4.2.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c5afb1a60e207dca502682537fefcfd9921e71d0b83e9576060f09abc6efab23"
 dependencies = [
- "indexmap",
+ "indexmap 2.2.6",
  "serde",
  "serde_json",
  "utoipa-gen",
diff --git a/Cargo.toml b/Cargo.toml
index 4901299..7d6c7d5 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -39,6 +39,10 @@ opt-level = "z"
 panic = "abort"
 strip = "symbols"
 
+[features]
+default = ["tokio-console"]
+tokio-console = ["dep:console-subscriber", "tokio/tracing"]
+
 [workspace]
 members = [
     ".",
@@ -134,3 +138,6 @@ tokio-graceful-shutdown = { workspace = true }
 tracing = { workspace = true }
 tracing-error = { workspace = true }
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
+
+# optional-dependencies
+console-subscriber = { version = "0.2", optional = true }
diff --git a/src/main.rs b/src/main.rs
index 1524ba8..34e89b1 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -19,10 +19,14 @@ use tracing_subscriber::prelude::*;
 async fn main() -> Result<(), AppError> {
     setup_panic!(metadata!().homepage("https://github.com/importantimport/hatsu/issues"));
 
-    tracing_subscriber::registry()
+    let registry = tracing_subscriber::registry()
         .with(tracing_subscriber::fmt::layer())
-        .with(tracing_error::ErrorLayer::default())
-        .init();
+        .with(tracing_error::ErrorLayer::default());
+
+    #[cfg(feature = "tokio-console")]
+    let registry = registry.with(console_subscriber::spawn());
+
+    registry.init();
 
     tracing::info!("loading environment variables");
     if dotenvy::dotenv().is_err() {