From 7f2c6a10132a38cfc005e12f67a247a070baf17d Mon Sep 17 00:00:00 2001
From: Nikita Vilunov <nikita@vilunov.me>
Date: Sat, 4 May 2024 01:59:43 +0200
Subject: [PATCH] continue propagated traces in http request handlers

---
 src/http.rs | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/src/http.rs b/src/http.rs
index dfe23a2..1d24c14 100644
--- a/src/http.rs
+++ b/src/http.rs
@@ -74,12 +74,14 @@ async fn main_loop(
     Ok(())
 }
 
+#[tracing::instrument(skip_all)]
 async fn route(
     registry: &MetricsRegistry,
     core: &LavinaCore,
     storage: &Storage,
     request: Request<hyper::body::Incoming>,
 ) -> HttpResult<Response<Full<Bytes>>> {
+    propagade_span_from_headers(&request);
     let res = match (request.method(), request.uri().path()) {
         (&Method::GET, "/metrics") => endpoint_metrics(registry),
         (&Method::GET, "/rooms") => endpoint_rooms(&core.rooms).await,
@@ -161,6 +163,7 @@ async fn endpoint_set_password(
     Ok(empty_204_request())
 }
 
+#[tracing::instrument(skip_all)]
 async fn endpoint_send_room_message(
     request: Request<hyper::body::Incoming>,
     core: &LavinaCore,
@@ -183,6 +186,7 @@ async fn endpoint_send_room_message(
     }
 }
 
+#[tracing::instrument(skip_all)]
 async fn endpoint_set_room_topic(
     request: Request<hyper::body::Incoming>,
     core: &LavinaCore,
@@ -282,3 +286,24 @@ where
         Full::new(Bytes::from(buffer))
     }
 }
+
+fn propagade_span_from_headers<T>(req: &Request<T>) {
+    use opentelemetry::propagation::Extractor;
+    use tracing::Span;
+    use tracing_opentelemetry::OpenTelemetrySpanExt;
+    struct HttpReqExtractor<'a, T> {
+        req: &'a Request<T>,
+    }
+    impl<'a, T> Extractor for HttpReqExtractor<'a, T> {
+        fn get(&self, key: &str) -> Option<&str> {
+            self.req.headers().get(key).and_then(|v| v.to_str().ok())
+        }
+
+        fn keys(&self) -> Vec<&str> {
+            self.req.headers().keys().map(|k| k.as_str()).collect()
+        }
+    }
+
+    let ctx = opentelemetry::global::get_text_map_propagator(|pp| pp.extract(&HttpReqExtractor { req }));
+    Span::current().set_parent(ctx);
+}