From d524222a862018da5b508a06985935db50d737b0 Mon Sep 17 00:00:00 2001
From: cy <cy@cy7.sh>
Date: Sat, 26 Apr 2025 15:37:54 -0400
Subject: [PATCH] make tokio console optional and make it actually work

---
 Cargo.lock  | 53 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 Cargo.toml  |  4 ++++
 src/main.rs | 29 ++++++++++++++++++++++++++---
 3 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index bbd7534..a3e049c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1503,6 +1503,7 @@ dependencies = [
  "tokio",
  "tokio-util",
  "tracing",
+ "tracing-subscriber",
  "ulid",
  "url",
 ]
@@ -1526,6 +1527,16 @@ dependencies = [
  "memchr",
 ]
 
+[[package]]
+name = "nu-ansi-term"
+version = "0.46.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84"
+dependencies = [
+ "overload",
+ "winapi",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@@ -1649,6 +1660,12 @@ dependencies = [
  "vcpkg",
 ]
 
+[[package]]
+name = "overload"
+version = "0.1.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39"
+
 [[package]]
 name = "parking_lot"
 version = "0.12.3"
@@ -2684,6 +2701,17 @@ dependencies = [
  "valuable",
 ]
 
+[[package]]
+name = "tracing-log"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3"
+dependencies = [
+ "log",
+ "once_cell",
+ "tracing-core",
+]
+
 [[package]]
 name = "tracing-subscriber"
 version = "0.3.19"
@@ -2691,12 +2719,15 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
 dependencies = [
  "matchers",
+ "nu-ansi-term",
  "once_cell",
  "regex",
  "sharded-slab",
+ "smallvec",
  "thread_local",
  "tracing",
  "tracing-core",
+ "tracing-log",
 ]
 
 [[package]]
@@ -2925,6 +2956,22 @@ dependencies = [
  "wasm-bindgen",
 ]
 
+[[package]]
+name = "winapi"
+version = "0.3.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
+dependencies = [
+ "winapi-i686-pc-windows-gnu",
+ "winapi-x86_64-pc-windows-gnu",
+]
+
+[[package]]
+name = "winapi-i686-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
+
 [[package]]
 name = "winapi-util"
 version = "0.1.9"
@@ -2934,6 +2981,12 @@ dependencies = [
  "windows-sys 0.59.0",
 ]
 
+[[package]]
+name = "winapi-x86_64-pc-windows-gnu"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
+
 [[package]]
 name = "windows-core"
 version = "0.61.0"
diff --git a/Cargo.toml b/Cargo.toml
index f548ea8..8dd350d 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,6 +3,9 @@ name = "nixcp"
 version = "0.1.0"
 edition = "2024"
 
+[build]
+rustflags = ["--cfg", "tokio_unstable"]
+
 [profile.release]
 lto = true
 codegen-units = 1
@@ -29,6 +32,7 @@ tokio-util = { version = "0.7.15", features = ["io"] }
 bytes = "1.10.1"
 object_store = { version = "0.12.0", features = ["aws"] }
 ulid = "1.2.1"
+tracing-subscriber = "0.3.19"
 
 [build-dependencies]
 cxx-build = "1.0"
diff --git a/src/main.rs b/src/main.rs
index bca419b..0fefdf5 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,11 @@
 #![feature(let_chains)]
-#![feature(extend_one)]
 #![feature(exit_status_error)]
 
 use std::path::PathBuf;
 
 use anyhow::{Context, Result};
 use clap::{Args, Parser, Subcommand};
+use tracing_subscriber::{EnvFilter, prelude::*};
 
 use push::Push;
 use store::Store;
@@ -26,6 +26,10 @@ mod uploader;
 struct Cli {
     #[command(subcommand)]
     command: Commands,
+
+    /// Whether to enable tokio-console
+    #[arg(long)]
+    tokio_console: bool,
 }
 
 #[derive(Debug, Subcommand)]
@@ -70,9 +74,8 @@ pub struct PushArgs {
 
 #[tokio::main]
 async fn main() -> Result<()> {
-    console_subscriber::init();
-
     let cli = Cli::parse();
+    init_logging(cli.tokio_console);
 
     match &cli.command {
         Commands::Push(cli) => {
@@ -87,3 +90,23 @@ async fn main() -> Result<()> {
 
     Ok(())
 }
+
+fn init_logging(tokio_console: bool) {
+    let env_filter = EnvFilter::from_default_env();
+    let fmt_layer = tracing_subscriber::fmt::layer().with_filter(env_filter);
+
+    let console_layer = if tokio_console {
+        Some(console_subscriber::spawn())
+    } else {
+        None
+    };
+
+    tracing_subscriber::registry()
+        .with(fmt_layer)
+        .with(console_layer)
+        .init();
+
+    if tokio_console {
+        println!("tokio-console is enabled");
+    }
+}