From 3307791855fcce717c9265fab8868e3d8b5443ea Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Wed, 1 Mar 2023 00:37:08 +0100
Subject: feat(tvix/store): add opentelemetry

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>
---
 tvix/store/src/bin/tvix-store.rs | 41 +++++++++++++++++++++++++++++++++++-----
 1 file changed, 36 insertions(+), 5 deletions(-)

(limited to 'tvix/store/src/bin')

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 {
-- 
cgit 1.4.1