athena_rs 3.26.1

//! Tokio console tracing helpers.
//!
//! Keeps the feature-specific subscriber logic out of the binary boot file.

use std::env;
use std::time::Duration;
use tracing_subscriber::filter::Targets;
use tracing_subscriber::layer::Layer;
use tracing_subscriber::registry::LookupSpan;

/// Returns whether the tokio-console subscriber should be enabled for this run.
pub fn env_opt_in() -> bool {
    env::var("ATHENA_BENCHMARK_CLIENT")
        .ok()
        .is_some_and(|v| v.trim().eq_ignore_ascii_case("true"))
}

/// Builds the tokio-console layer with repository-specific retention rules.
pub fn subscriber_layer<S>() -> impl Layer<S> + Send + Sync + 'static
where
    S: tracing::Subscriber + for<'a> LookupSpan<'a>,
{
    const MIN_RETENTION: Duration = Duration::from_secs(60 * 60);
    let retention = match env::var("TOKIO_CONSOLE_RETENTION") {
        Ok(s) => {
            let parsed = humantime::parse_duration(&s).unwrap_or_else(|e| {
                panic!("failed to parse a duration from `TOKIO_CONSOLE_RETENTION={s:?}`: {e}")
            });
            parsed.max(MIN_RETENTION)
        }
        Err(_) => console_subscriber::ConsoleLayer::DEFAULT_RETENTION,
    };

    console_subscriber::ConsoleLayer::builder()
        .with_default_env()
        .retention(retention)
        .spawn()
}

/// Builds the trace-target filter used when tokio-console is enabled.
pub fn trace_targets_filter() -> Targets {
    Targets::new()
        .with_target("tokio", tracing::Level::TRACE)
        .with_target("runtime", tracing::Level::TRACE)
}

/// Returns `true` when a target is internal tokio-console runtime telemetry that
/// should stay out of stdout and file logs.
pub fn should_suppress_output_target(target: &str) -> bool {
    target == "tokio"
        || target.starts_with("tokio::")
        || target == "runtime"
        || target.starts_with("runtime::")
        || target == "console_subscriber"
        || target.starts_with("console_subscriber::")
}

#[cfg(test)]
mod tests {
    use super::should_suppress_output_target;

    #[test]
    fn suppresses_tokio_console_runtime_targets() {
        assert!(should_suppress_output_target("tokio"));
        assert!(should_suppress_output_target("tokio::task::waker"));
        assert!(should_suppress_output_target("runtime"));
        assert!(should_suppress_output_target("runtime::resource::poll_op"));
        assert!(should_suppress_output_target("console_subscriber"));
        assert!(should_suppress_output_target("console_subscriber::builder"));
    }

    #[test]
    fn keeps_athena_and_dependency_targets_visible() {
        assert!(!should_suppress_output_target("athena_rs::bootstrap"));
        assert!(!should_suppress_output_target("sqlx_core::pool::inner"));
        assert!(!should_suppress_output_target("tokio_postgres"));
        assert!(!should_suppress_output_target("actix_http"));
    }
}