From db288b169e52021faa9aa44c3c6b5d2e3f78b917 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Wed, 26 Jun 2024 23:00:33 +0000
Subject: [PATCH] add config option for tokio_console runtime enablement

Signed-off-by: Jason Volk <jason@zemos.net>
---
 conduwuit-example.toml |  5 +++
 src/core/config/mod.rs |  4 +++
 src/main/tracing.rs    | 72 ++++++++++++++++++++++++++++++------------
 3 files changed, 61 insertions(+), 20 deletions(-)

diff --git a/conduwuit-example.toml b/conduwuit-example.toml
index 37aaf2fe3..f8f930707 100644
--- a/conduwuit-example.toml
+++ b/conduwuit-example.toml
@@ -371,6 +371,11 @@ allow_profile_lookup_federation_requests = true
 # If 'tracing_flame' is enabled, set the path to write the generated profile.
 # tracing_flame_output_path = "./tracing.folded"
 
+# Enable the tokio-console. This option is only relevant to developers.
+# See: docs/development.md#debugging-with-tokio-console for more information.
+#tokio_console = false
+
+
 ### Generic database options
 
 # Set this to any float value to multiply conduwuit's in-memory LRU caches with.
diff --git a/src/core/config/mod.rs b/src/core/config/mod.rs
index 4260c46e6..0b9cb2990 100644
--- a/src/core/config/mod.rs
+++ b/src/core/config/mod.rs
@@ -355,6 +355,9 @@ pub struct Config {
 	#[serde(default = "default_sentry_traces_sample_rate")]
 	pub sentry_traces_sample_rate: f32,
 
+	#[serde(default)]
+	pub tokio_console: bool,
+
 	#[serde(flatten)]
 	#[allow(clippy::zero_sized_map_values)] // this is a catchall, the map shouldn't be zero at runtime
 	catchall: BTreeMap<String, IgnoredAny>,
@@ -855,6 +858,7 @@ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 					.as_ref()
 					.map_or("", |url| url.as_str()),
 			),
+			("Enable the tokio-console", &self.tokio_console.to_string()),
 		];
 
 		let mut msg: String = "Active config values:\n\n".to_owned();
diff --git a/src/main/tracing.rs b/src/main/tracing.rs
index b5f7a7d16..bbfe4dc42 100644
--- a/src/main/tracing.rs
+++ b/src/main/tracing.rs
@@ -3,9 +3,10 @@
 use conduit::{
 	config,
 	config::Config,
+	debug_warn,
 	log::{capture, LogLevelReloadHandles, ReloadHandle},
 };
-use tracing_subscriber::{prelude::*, reload, EnvFilter, Registry};
+use tracing_subscriber::{layer::SubscriberExt, reload, EnvFilter, Layer, Registry};
 
 #[cfg(feature = "perf_measurements")]
 pub(crate) type TracingFlameGuard = Option<tracing_flame::FlushGuard<std::io::BufWriter<std::fs::File>>>;
@@ -14,7 +15,6 @@
 
 #[allow(clippy::redundant_clone)]
 pub(crate) fn init(config: &Config) -> (LogLevelReloadHandles, TracingFlameGuard, Arc<capture::State>) {
-	let registry = Registry::default();
 	let fmt_layer = tracing_subscriber::fmt::Layer::new();
 	let filter_layer = match EnvFilter::try_new(&config.log) {
 		Ok(s) => s,
@@ -25,22 +25,14 @@ pub(crate) fn init(config: &Config) -> (LogLevelReloadHandles, TracingFlameGuard
 	};
 
 	let mut reload_handles = Vec::<Box<dyn ReloadHandle<EnvFilter> + Send + Sync>>::new();
-	let subscriber = registry;
-
-	#[cfg(all(feature = "tokio_console", tokio_unstable))]
-	let subscriber = {
-		let console_layer = console_subscriber::spawn();
-		subscriber.with(console_layer)
-	};
-
 	let (fmt_reload_filter, fmt_reload_handle) = reload::Layer::new(filter_layer.clone());
 	reload_handles.push(Box::new(fmt_reload_handle));
 
+	let subscriber = Registry::default().with(fmt_layer.with_filter(fmt_reload_filter));
+
 	let cap_state = Arc::new(capture::State::new());
 	let cap_layer = capture::Layer::new(&cap_state);
-	let subscriber = subscriber
-		.with(fmt_layer.with_filter(fmt_reload_filter))
-		.with(cap_layer);
+	let subscriber = subscriber.with(cap_layer);
 
 	#[cfg(feature = "sentry_telemetry")]
 	let subscriber = {
@@ -95,13 +87,53 @@ pub(crate) fn init(config: &Config) -> (LogLevelReloadHandles, TracingFlameGuard
 	#[cfg_attr(not(feature = "perf_measurements"), allow(clippy::let_unit_value))]
 	let flame_guard = ();
 
-	tracing::subscriber::set_global_default(subscriber).expect("failed to set global tracing subscriber");
+	let ret = (LogLevelReloadHandles::new(reload_handles), flame_guard, cap_state);
+
+	// Enable the tokio console. This is slightly kludgy because we're judggling
+	// compile-time and runtime conditions to elide it, each of those changing the
+	// subscriber's type.
+	let (console_enabled, console_disabled_reason) = tokio_console_enabled(config);
+	#[cfg(all(feature = "tokio_console", tokio_unstable))]
+	if console_enabled {
+		let console_layer = console_subscriber::ConsoleLayer::builder()
+			.with_default_env()
+			.spawn();
+
+		set_global_default(subscriber.with(console_layer));
+		return ret;
+	}
 
-	#[cfg(all(feature = "tokio_console", feature = "release_max_log_level", tokio_unstable))]
-	tracing::error!(
-		"'tokio_console' feature and 'release_max_log_level' feature are incompatible, because console-subscriber \
-		 needs access to trace-level events. 'release_max_log_level' must be disabled to use tokio-console."
-	);
+	set_global_default(subscriber);
+
+	// If there's a reason the tokio console was disabled when it might be desired
+	// we output that here after initializing logging
+	if !console_enabled && !console_disabled_reason.is_empty() {
+		debug_warn!("{console_disabled_reason}");
+	}
+
+	ret
+}
+
+fn tokio_console_enabled(config: &Config) -> (bool, &'static str) {
+	if !cfg!(all(feature = "tokio_console", tokio_unstable)) {
+		return (false, "");
+	}
+
+	if cfg!(feature = "release_max_log_level") && !cfg!(debug_assertions) {
+		return (
+			false,
+			"'tokio_console' feature and 'release_max_log_level' feature are incompatible.",
+		);
+	}
+
+	if !config.tokio_console {
+		return (false, "tokio console is available but disabled by the configuration.");
+	}
+
+	(true, "")
+}
 
-	(LogLevelReloadHandles::new(reload_handles), flame_guard, cap_state)
+fn set_global_default<S: SubscriberExt + Send + Sync>(subscriber: S) {
+	tracing::subscriber::set_global_default(subscriber)
+		.expect("the global default tracing subscriber failed to be initialized");
 }
-- 
GitLab