From a0c3c1ff17a47c5288b1aa93a28a16e95a33f953 Mon Sep 17 00:00:00 2001
From: Chris Bellew <cjbellew@gmail.com>
Date: Sat, 1 Jun 2024 10:22:32 +0800
Subject: [PATCH] Bring your own OTEL tracing backend

---
 apps/userspace-collector/src/main.rs    |  8 +++++++-
 apps/userspace-collector/src/tracing.rs | 14 ++++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/apps/userspace-collector/src/main.rs b/apps/userspace-collector/src/main.rs
index fe99251..ac961bb 100644
--- a/apps/userspace-collector/src/main.rs
+++ b/apps/userspace-collector/src/main.rs
@@ -4,6 +4,7 @@ use bpf::{attach_uprobes, init_bpf};
 use log::info;
 use metrics::init_metrics;
 use receive::listen_to_cpu;
+use std::env;
 use std::error::Error;
 use tokio::signal;
 use tracing::TraceEmitter;
@@ -26,7 +27,12 @@ async fn main() -> Result<(), Box<dyn Error>> {
 
     // Initialise the tracing system. We'll build traces from the
     // postgres events and send them to the tracing backend using OpenTelemetry.
-    let tracing = TraceEmitter::initialise()?;
+    // Use the provided bring-your-own backend endpoint if it's set, otherwise
+    // use the built-in backend.
+    let tracing_endpoint = env::var("OTEL_TRACING_ENDPOINT")
+        .ok()
+        .filter(|val| !val.is_empty());
+    let tracing = TraceEmitter::initialise(tracing_endpoint)?;
 
     // Initialise the metrics prometheus exporter. We'll collect metrics
     // about the postgres queries and this userspace collector, and expose
diff --git a/apps/userspace-collector/src/tracing.rs b/apps/userspace-collector/src/tracing.rs
index ade0967..e81c30e 100644
--- a/apps/userspace-collector/src/tracing.rs
+++ b/apps/userspace-collector/src/tracing.rs
@@ -22,13 +22,23 @@ type TraceIds = Arc<Mutex<HashMap<u32, u128>>>;
 type RemoteContexts = Arc<Mutex<HashMap<u32, Context>>>;
 
 impl TraceEmitter {
-    pub fn initialise() -> Result<Self, anyhow::Error> {
+    pub fn initialise(backend_endpoint: Option<String>) -> Result<Self, anyhow::Error> {
+        let endpoint = match backend_endpoint {
+            Some(endpoint) => {
+                println!("Sending traces to {}", endpoint);
+                endpoint
+            }
+            None => {
+                println!("Sending traces to built-in backend");
+                "http://localhost:4317".to_string()
+            }
+        };
         let tracer = opentelemetry_otlp::new_pipeline()
             .tracing()
             .with_exporter(
                 opentelemetry_otlp::new_exporter()
                     .tonic()
-                    .with_endpoint("http://localhost:4317"),
+                    .with_endpoint(endpoint),
             )
             .with_trace_config(
                 config()