feat(tvix/tracing): add tracing-tracy support r/8275

This introduces another feature flag, "tracy" to the `tvix-tracing` crate.

If enabled (not enabled by default), it'll add an additional layer
emitting packets in a format that https://github.com/wolfpld/tracy can
display.

I had to be a bit tricky with the combinatorial complexity when adding
this, but the resulting code still seems manageable.

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

3 files changed, 28 insertions, 3 deletions

diff --git a/tvix/tracing/Cargo.toml b/tvix/tracing/Cargo.toml
index 41654d00c02f..490f086a27f3 100644
--- a/tvix/tracing/Cargo.toml
+++ b/tvix/tracing/Cargo.toml
@@ -16,6 +16,7 @@ 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 }
+tracing-tracy = { version = "0.11.0", features = ["flush-on-exit"], optional = true }
 
 [features]
 default = []
@@ -25,6 +26,9 @@ otlp = [
   "dep:opentelemetry-otlp",
   "dep:opentelemetry_sdk"
 ]
+tracy = [
+  "dep:tracing-tracy"
+]
 
 [lints]
 workspace = true
diff --git a/tvix/tracing/default.nix b/tvix/tracing/default.nix
index a4fe3a5d90c3..dd7dc200f2c6 100644
--- a/tvix/tracing/default.nix
+++ b/tvix/tracing/default.nix
@@ -6,6 +6,6 @@
   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" ];
+    features = [ "otlp" "tracy" ];
   };
 })
diff --git a/tvix/tracing/src/lib.rs b/tvix/tracing/src/lib.rs
index 36bd7cec10c0..08e100781073 100644
--- a/tvix/tracing/src/lib.rs
+++ b/tvix/tracing/src/lib.rs
@@ -13,6 +13,8 @@ use opentelemetry_sdk::{
     trace::BatchConfigBuilder,
     Resource,
 };
+#[cfg(feature = "tracy")]
+use tracing_tracy::TracyLayer;
 
 lazy_static! {
     pub static ref PB_PROGRESS_STYLE: ProgressStyle = ProgressStyle::with_template(
@@ -173,12 +175,31 @@ impl TracingBuilder {
                 let (tracer, tx) = gen_otlp_tracer(service_name.to_string());
                 // Create a tracing layer with the configured tracer
                 let layer = tracing_opentelemetry::layer().with_tracer(tracer);
-                subscriber.with(Some(layer)).try_init()?;
+
+                #[cfg(feature = "tracy")]
+                {
+                    subscriber
+                        .with(TracyLayer::default())
+                        .with(Some(layer))
+                        .try_init()?;
+                }
+
+                #[cfg(not(feature = "tracy"))]
+                {
+                    subscriber.with(Some(layer)).try_init()?;
+                }
                 return Ok(TracingHandle { tx: Some(tx) });
             }
         }
+        #[cfg(feature = "tracy")]
+        {
+            subscriber.with(TracyLayer::default()).try_init()?;
+        }
+        #[cfg(not(feature = "tracy"))]
+        {
+            subscriber.try_init()?;
+        }
 
-        subscriber.try_init()?;
         Ok(TracingHandle { tx: None })
     }
 }