feat(tvix/tracing): introduce common tvix-tracing crate r/8242

Introduce a new common crate that contains tracing boilerplate which then
can be used in the cli, tvix-store and tvix-build crates.
It has otlp as an optional feature, which is currently only used by
tvix-store.

Change-Id: I41468ac4d9c65174515d721513b96fea463d6ed2
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11758
Tested-by: BuildkiteCI
Reviewed-by: flokli <flokli@flokli.de>
Autosubmit: Simon Hauser <simon.hauser@helsinki-systems.de>

3 files changed, 129 insertions, 0 deletions

diff --git a/tvix/tracing/Cargo.toml b/tvix/tracing/Cargo.toml
new file mode 100644
index 000000000000..eea6ca17711f
--- /dev/null
+++ b/tvix/tracing/Cargo.toml
@@ -0,0 +1,28 @@
+[package]
+name = "tvix-tracing"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+lazy_static = "1.4.0"
+tracing = { version = "0.1.40", features = ["max_level_trace", "release_max_level_info"] }
+tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
+indicatif = "0.17.8"
+tracing-indicatif = "0.3.6"
+
+tracing-opentelemetry = { version = "0.23.0", optional = true }
+opentelemetry = { version = "0.22.0", optional = true }
+opentelemetry-otlp = { version = "0.15.0", optional = true }
+opentelemetry_sdk = { version = "0.22.1", features = ["rt-tokio"], optional = true }
+
+[features]
+default = []
+otlp = [
+  "dep:tracing-opentelemetry",
+  "dep:opentelemetry",
+  "dep:opentelemetry-otlp",
+  "dep:opentelemetry_sdk"
+]
+
+[lints]
+workspace = true
diff --git a/tvix/tracing/default.nix b/tvix/tracing/default.nix
new file mode 100644
index 000000000000..a4fe3a5d90c3
--- /dev/null
+++ b/tvix/tracing/default.nix
@@ -0,0 +1,11 @@
+{ depot, lib, ... }:
+
+(depot.tvix.crates.workspaceMembers.tvix-tracing.build.override {
+  runTests = true;
+}).overrideAttrs (old: rec {
+  meta.ci.targets = lib.filter (x: lib.hasPrefix "with-features" x || x == "no-features") (lib.attrNames passthru);
+  passthru = depot.tvix.utils.mkFeaturePowerset {
+    inherit (old) crateName;
+    features = [ "otlp" ];
+  };
+})
diff --git a/tvix/tracing/src/lib.rs b/tvix/tracing/src/lib.rs
new file mode 100644
index 000000000000..8ad90835338e
--- /dev/null
+++ b/tvix/tracing/src/lib.rs
@@ -0,0 +1,90 @@
+use indicatif::ProgressStyle;
+use lazy_static::lazy_static;
+use tracing::Level;
+use tracing_indicatif::{filter::IndicatifFilter, IndicatifLayer};
+use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer};
+
+#[cfg(feature = "otlp")]
+use opentelemetry::KeyValue;
+#[cfg(feature = "otlp")]
+use opentelemetry_sdk::{
+    resource::{ResourceDetector, SdkProvidedResourceDetector},
+    trace::BatchConfig,
+    Resource,
+};
+
+lazy_static! {
+    pub static ref PB_PROGRESS_STYLE: ProgressStyle = ProgressStyle::with_template(
+        "{span_child_prefix}{bar:30} {wide_msg} [{elapsed_precise}]  {pos:>7}/{len:7}"
+    )
+    .expect("invalid progress template");
+    pub static ref PB_SPINNER_STYLE: ProgressStyle = ProgressStyle::with_template(
+        "{span_child_prefix}{spinner} {wide_msg} [{elapsed_precise}]  {pos:>7}/{len:7}"
+    )
+    .expect("invalid progress template");
+}
+
+// using a macro_rule here because of the complex return type
+macro_rules! init_base_subscriber {
+    ($level: expr) => {{
+        let indicatif_layer = IndicatifLayer::new().with_progress_style(PB_SPINNER_STYLE.clone());
+
+        // Set up the tracing subscriber.
+        tracing_subscriber::registry()
+            .with(
+                tracing_subscriber::fmt::Layer::new()
+                    .with_writer(indicatif_layer.get_stderr_writer())
+                    .compact()
+                    .with_filter(
+                        EnvFilter::builder()
+                            .with_default_directive($level.into())
+                            .from_env()
+                            .expect("invalid RUST_LOG"),
+                    ),
+            )
+            .with(indicatif_layer.with_filter(
+                // only show progress for spans with indicatif.pb_show field being set
+                IndicatifFilter::new(false),
+            ))
+    }};
+}
+
+pub fn init(level: Level) -> Result<(), tracing_subscriber::util::TryInitError> {
+    init_base_subscriber!(level).try_init()
+}
+
+#[cfg(feature = "otlp")]
+pub fn init_with_otlp(
+    level: Level,
+    service_name: &'static str,
+) -> Result<(), tracing_subscriber::util::TryInitError> {
+    let subscriber = init_base_subscriber!(level);
+
+    let tracer = opentelemetry_otlp::new_pipeline()
+        .tracing()
+        .with_exporter(opentelemetry_otlp::new_exporter().tonic())
+        .with_batch_config(BatchConfig::default())
+        .with_trace_config(opentelemetry_sdk::trace::config().with_resource({
+            // use SdkProvidedResourceDetector.detect to detect resources,
+            // but replace the default service name with our default.
+            // https://github.com/open-telemetry/opentelemetry-rust/issues/1298
+            let resources = SdkProvidedResourceDetector.detect(std::time::Duration::from_secs(0));
+            // SdkProvidedResourceDetector currently always sets
+            // `service.name`, but we don't like its default.
+            if resources.get("service.name".into()).unwrap() == "unknown_service".into() {
+                resources.merge(&Resource::new([KeyValue::new(
+                    "service.name",
+                    service_name,
+                )]))
+            } else {
+                resources
+            }
+        }))
+        .install_batch(opentelemetry_sdk::runtime::Tokio)
+        .expect("Failed to install tokio runtime");
+
+    // Create a tracing layer with the configured tracer
+    let layer = tracing_opentelemetry::layer().with_tracer(tracer);
+
+    subscriber.with(Some(layer)).try_init()
+}