From fbabcf0420bcc913971d529c11d58f6f888f1002 Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Thu, 16 Feb 2023 16:49:22 +0100
Subject: [PATCH] feat(tvix/store): add --json arg

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>
---
 tvix/Cargo.lock        | 13 +++++++++++++
 tvix/Cargo.nix         | 42 +++++++++++++++++++++++++++++++++++++++++-
 tvix/store/Cargo.toml  |  2 +-
 tvix/store/src/main.rs | 29 +++++++++++++++++++++++++++--
 4 files changed, 82 insertions(+), 4 deletions(-)

diff --git a/tvix/Cargo.lock b/tvix/Cargo.lock
index 106a1e38d..6c04176a1 100644
--- a/tvix/Cargo.lock
+++ b/tvix/Cargo.lock
@@ -2683,6 +2683,16 @@ dependencies = [
  "tracing-core",
 ]
 
+[[package]]
+name = "tracing-serde"
+version = "0.1.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bc6b213177105856957181934e4920de57730fc69bf42c37ee5bb664d406d9e1"
+dependencies = [
+ "serde",
+ "tracing-core",
+]
+
 [[package]]
 name = "tracing-subscriber"
 version = "0.3.16"
@@ -2690,11 +2700,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a6176eae26dd70d0c919749377897b54a9276bd7061339665dd68777926b5a70"
 dependencies = [
  "nu-ansi-term",
+ "serde",
+ "serde_json",
  "sharded-slab",
  "smallvec",
  "thread_local",
  "tracing-core",
  "tracing-log",
+ "tracing-serde",
 ]
 
 [[package]]
diff --git a/tvix/Cargo.nix b/tvix/Cargo.nix
index fc3f945df..4e9e42e47 100644
--- a/tvix/Cargo.nix
+++ b/tvix/Cargo.nix
@@ -7885,6 +7885,30 @@ rec {
         };
         resolvedDefaultFeatures = [ "log-tracer" "std" ];
       };
+      "tracing-serde" = rec {
+        crateName = "tracing-serde";
+        version = "0.1.3";
+        edition = "2018";
+        sha256 = "1qfr0va69djvxqvjrx4vqq7p6myy414lx4w1f6amcn0hfwqj2sxw";
+        authors = [
+          "Tokio Contributors <team@tokio.rs>"
+        ];
+        dependencies = [
+          {
+            name = "serde";
+            packageId = "serde";
+          }
+          {
+            name = "tracing-core";
+            packageId = "tracing-core";
+          }
+        ];
+        features = {
+          "valuable" = [ "valuable_crate" "valuable-serde" "tracing-core/valuable" ];
+          "valuable-serde" = [ "dep:valuable-serde" ];
+          "valuable_crate" = [ "dep:valuable_crate" ];
+        };
+      };
       "tracing-subscriber" = rec {
         crateName = "tracing-subscriber";
         version = "0.3.16";
@@ -7901,6 +7925,16 @@ rec {
             packageId = "nu-ansi-term";
             optional = true;
           }
+          {
+            name = "serde";
+            packageId = "serde";
+            optional = true;
+          }
+          {
+            name = "serde_json";
+            packageId = "serde_json";
+            optional = true;
+          }
           {
             name = "sharded-slab";
             packageId = "sharded-slab";
@@ -7928,6 +7962,11 @@ rec {
             usesDefaultFeatures = false;
             features = [ "log-tracer" "std" ];
           }
+          {
+            name = "tracing-serde";
+            packageId = "tracing-serde";
+            optional = true;
+          }
         ];
         devDependencies = [
           {
@@ -7962,7 +8001,7 @@ rec {
           "valuable-serde" = [ "dep:valuable-serde" ];
           "valuable_crate" = [ "dep:valuable_crate" ];
         };
-        resolvedDefaultFeatures = [ "alloc" "ansi" "default" "fmt" "nu-ansi-term" "registry" "sharded-slab" "smallvec" "std" "thread_local" "tracing-log" ];
+        resolvedDefaultFeatures = [ "alloc" "ansi" "default" "fmt" "json" "nu-ansi-term" "registry" "serde" "serde_json" "sharded-slab" "smallvec" "std" "thread_local" "tracing-log" "tracing-serde" ];
       };
       "try-lock" = rec {
         crateName = "try-lock";
@@ -8329,6 +8368,7 @@ rec {
           {
             name = "tracing-subscriber";
             packageId = "tracing-subscriber";
+            features = [ "json" ];
           }
         ];
         buildDependencies = [
diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index 196e986b8..766c44cb1 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 e71c1ef88..6c3b44eec 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();