feat(tvix/store): add --json arg r/5924

This configures logging as JSON.

Change-Id: I22cdda84de215bfceda4e9d47bc8d487a5451a6e
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8130
Tested-by: BuildkiteCI
Reviewed-by: raitobezarius <tvl@lahfa.xyz>

2 files changed, 28 insertions, 3 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 196e986b89ab..766c44cb1775 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -23,7 +23,7 @@ tokio-stream = "0.1.11"
 tokio = { version = "1.23.0", features = ["rt-multi-thread"] }
 tonic = "0.8.2"
 tracing = "0.1.37"
-tracing-subscriber = "0.3.16"
+tracing-subscriber = { version = "0.3.16", features = ["json"] }
 
 [dependencies.tonic-reflection]
 optional = true
diff --git a/tvix/store/src/main.rs b/tvix/store/src/main.rs
index e71c1ef88614..6c3b44eec7f0 100644
--- a/tvix/store/src/main.rs
+++ b/tvix/store/src/main.rs
@@ -1,3 +1,4 @@
+use tracing_subscriber::prelude::*;
 use tvix_store::blobservice::SledBlobService;
 use tvix_store::chunkservice::SledChunkService;
 use tvix_store::directoryservice::SledDirectoryService;
@@ -23,6 +24,10 @@ struct Cli {
     #[clap(long, short = 'l')]
     listen_address: Option<String>,
 
+    /// Whether to log in JSON
+    #[clap(long)]
+    json: bool,
+
     #[clap(long)]
     log_level: Option<Level>,
 }
@@ -37,8 +42,28 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
         .unwrap();
 
     let level = cli.log_level.unwrap_or(Level::INFO);
-    let subscriber = tracing_subscriber::fmt().with_max_level(level).finish();
-    tracing::subscriber::set_global_default(subscriber).ok();
+
+    let subscriber = tracing_subscriber::registry()
+        .with(if cli.json {
+            Some(
+                tracing_subscriber::fmt::Layer::new()
+                    .with_writer(std::io::stdout.with_max_level(level))
+                    .json(),
+            )
+        } else {
+            None
+        })
+        .with(if !cli.json {
+            Some(
+                tracing_subscriber::fmt::Layer::new()
+                    .with_writer(std::io::stdout.with_max_level(level))
+                    .pretty(),
+            )
+        } else {
+            None
+        });
+
+    tracing::subscriber::set_global_default(subscriber).expect("Unable to set global subscriber");
 
     let mut server = Server::builder();