feat(tvix/store): support RUST_LOG env var r/7723

This allows selectively increasing the log level for only parts of the
stack.

For example, the following RUST_LOG env var enables "tracing" level
logging for `tvix_store` and `tvix_castore`, while keeping it at "info"
for the rest of the stack:

export RUST_LOG='info,tvix_store=trace,tvix_castore=trace'

It only affects logs, not traces (if enabled).

Change-Id: Ib936bd132a405f216e75c843db83fbd71d20a18a
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11182
Autosubmit: flokli <flokli@flokli.de>
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>

1 files changed, 23 insertions, 5 deletions

diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index f3a5d01b605b..27a67b7c91c1 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -10,7 +10,8 @@ use tokio_listener::UserOptions;
 use tonic::transport::Server;
 use tracing::info;
 use tracing::Level;
-use tracing_subscriber::fmt::writer::MakeWriterExt;
+use tracing_subscriber::EnvFilter;
+use tracing_subscriber::Layer;
 use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 use tvix_castore::proto::blob_service_server::BlobServiceServer;
@@ -55,6 +56,10 @@ struct Cli {
     #[arg(long, default_missing_value = "true", default_value = "true", num_args(0..=1), require_equals(true), action(clap::ArgAction::Set))]
     otlp: bool,
 
+    /// A global log level to use when printing logs.
+    /// It's also possible to set `RUST_LOG` according to
+    /// `tracing_subscriber::filter::EnvFilter`, which will always have
+    /// priority.
     #[arg(long)]
     log_level: Option<Level>,
 
@@ -172,19 +177,32 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     // configure log settings
     let level = cli.log_level.unwrap_or(Level::INFO);
 
+    // Set up the tracing subscriber.
     let subscriber = tracing_subscriber::registry()
         .with(
             cli.json.then_some(
                 tracing_subscriber::fmt::Layer::new()
-                    .with_writer(std::io::stderr.with_max_level(level))
-                    .json(),
+                    .with_writer(std::io::stderr)
+                    .json()
+                    .with_filter(
+                        EnvFilter::builder()
+                            .with_default_directive(level.into())
+                            .from_env()
+                            .expect("invalid RUST_LOG"),
+                    ),
             ),
         )
         .with(
             (!cli.json).then_some(
                 tracing_subscriber::fmt::Layer::new()
-                    .with_writer(std::io::stderr.with_max_level(level))
-                    .pretty(),
+                    .with_writer(std::io::stderr)
+                    .pretty()
+                    .with_filter(
+                        EnvFilter::builder()
+                            .with_default_directive(level.into())
+                            .from_env()
+                            .expect("invalid RUST_LOG"),
+                    ),
             ),
         );