feat(tvix/store): add opentelemetry r/7308

This is behind the otlp feature flag (currently enabled by default).

By default, this will try to push traces to a OTLP collector running at
https://localhost:4317.

You can make one available by running:

```
docker run -d --name jaeger \
  -e COLLECTOR_ZIPKIN_HOST_PORT=:9411 \
  -e COLLECTOR_OTLP_ENABLED=true \
  -p 6831:6831/udp \
  -p 6832:6832/udp \
  -p 5778:5778 \
  -p 16686:16686 \
  -p 4317:4317 \
  -p 4318:4318 \
  -p 14250:14250 \
  -p 14268:14268 \
  -p 14269:14269 \
  -p 9411:9411 --rm \
  jaegertracing/all-in-one:1.42
```

Started like that, jaeger brings a web interface at
http://localhost:16686/search

As documented in
https://docs.rs/opentelemetry-otlp/latest/opentelemetry_otlp/, you can
point this to another location by setting `OTEL_EXPORTER_OTLP_ENDPOINT`.

Change-Id: Id1dca367d70027b2ea98bb70bcf99a68363ec2be
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8194
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: aaqaishtyaq <aaqaishtyaq@gmail.com>

2 files changed, 43 insertions, 7 deletions

diff --git a/tvix/store/Cargo.toml b/tvix/store/Cargo.toml
index e6d17e0f6865..39bd0851bc54 100644
--- a/tvix/store/Cargo.toml
+++ b/tvix/store/Cargo.toml
@@ -16,6 +16,9 @@ lazy_static = "1.4.0"
 nix-compat = { path = "../nix-compat", features = ["async"] }
 pin-project-lite = "0.2.13"
 prost = "0.12.1"
+opentelemetry = { version = "0.21.0", optional = true}
+opentelemetry-otlp = { version = "0.14.0", optional = true }
+opentelemetry_sdk = { version = "0.21.0", features = ["rt-tokio"], optional = true}
 sha2 = "0.10.6"
 sled = { version = "0.34.7" }
 thiserror = "1.0.38"
@@ -26,6 +29,7 @@ tokio-util = { version = "0.7.9", features = ["io", "io-util", "compat"] }
 tonic = { version = "0.10.2", 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"] }
 tvix-castore = { path = "../castore" }
 url = "2.4.0"
@@ -48,7 +52,8 @@ tempfile = "3.3.0"
 tokio-retry = "0.3.0"
 
 [features]
-default = ["fuse", "tonic-reflection"]
+default = ["fuse", "otlp", "tonic-reflection"]
 fuse = ["tvix-castore/fuse"]
-virtiofs = ["tvix-castore/virtiofs"]
+otlp = ["dep:opentelemetry", "dep:opentelemetry-otlp", "dep:opentelemetry_sdk"]
 tonic-reflection = ["dep:tonic-reflection", "tvix-castore/tonic-reflection"]
+virtiofs = ["tvix-castore/virtiofs"]
diff --git a/tvix/store/src/bin/tvix-store.rs b/tvix/store/src/bin/tvix-store.rs
index a245af78d065..7c3d2ec0c797 100644
--- a/tvix/store/src/bin/tvix-store.rs
+++ b/tvix/store/src/bin/tvix-store.rs
@@ -1,6 +1,10 @@
+use clap::Parser;
 use clap::Subcommand;
 
 use futures::future::try_join_all;
+use tonic::transport::Server;
+use tracing::info;
+use tracing::Level;
 
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -9,6 +13,7 @@ use tokio_listener::SystemOptions;
 use tokio_listener::UserOptions;
 
 use tracing_subscriber::prelude::*;
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt};
 
 use tvix_castore::proto::blob_service_server::BlobServiceServer;
 use tvix_castore::proto::directory_service_server::DirectoryServiceServer;
@@ -32,10 +37,6 @@ use tvix_castore::proto::FILE_DESCRIPTOR_SET as CASTORE_FILE_DESCRIPTOR_SET;
 #[cfg(feature = "tonic-reflection")]
 use tvix_store::proto::FILE_DESCRIPTOR_SET;
 
-use clap::Parser;
-use tonic::{transport::Server, Result};
-use tracing::{info, Level};
-
 #[derive(Parser)]
 #[command(author, version, about, long_about = None)]
 struct Cli {
@@ -170,7 +171,37 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
             ),
         );
 
-    tracing::subscriber::set_global_default(subscriber).expect("Unable to set global subscriber");
+    // Add the otlp layer (when otlp is enabled), then init the registry.
+    // Or if the feature is disabled, just init without adding the layer.
+    // It's necessary to do this separately, as every with() call chains the
+    // layer into the type of the registry.
+    #[cfg(feature = "otlp")]
+    {
+        let opentelemetry_layer = {
+            let otlp_exporter = opentelemetry_otlp::new_exporter().tonic();
+            // TODO: re-add once https://github.com/open-telemetry/opentelemetry-rust/pull/1252 is solved.
+            // let mut metadata = tonic::metadata::MetadataMap::new();
+            // metadata.insert("service.name", "tvix.store".parse()?);
+            // otlp_exporter.with_metadata(metadata),
+
+            let tracer = opentelemetry_otlp::new_pipeline()
+                .tracing()
+                .with_exporter(otlp_exporter)
+                .install_batch(opentelemetry_sdk::runtime::Tokio)?;
+
+            // Create a tracing layer with the configured tracer
+            tracing_opentelemetry::layer().with_tracer(tracer)
+        };
+
+        let subscriber = subscriber.with(opentelemetry_layer);
+        subscriber.try_init()?;
+    }
+
+    // Init the registry (when otlp is not enabled)
+    #[cfg(not(feature = "otlp"))]
+    {
+        subscriber.try_init()?;
+    }
 
     match cli.command {
         Commands::Daemon {