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>

2 files changed, 24 insertions, 6 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 2dd155dbbc..366bd39216 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -31,7 +31,7 @@ tonic = { version = "0.11.0", features = ["tls", "tls-roots"] }
 tower = "0.4.13"
 tracing = "0.1.37"
 tracing-opentelemetry = "0.22.0"
-tracing-subscriber = { version = "0.3.16", features = ["json"] }
+tracing-subscriber = { version = "0.3.16", features = ["env-filter", "json"] }
 tvix-castore = { path = "../castore" }
 url = "2.4.0"
 walkdir = "2.4.0"
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index f3a5d01b60..27a67b7c91 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"),
+                    ),
             ),
         );