Propagate parent trace context

crates/cli/src/server.rs

@@ -19,21 +19,21 @@ use std::{
 
 use anyhow::Context;
 use clap::Clap;
-use hyper::{header, Server};
+use hyper::{header, Server, Version};
 use mas_config::RootConfig;
 use mas_core::{
     storage::MIGRATOR,
     tasks::{self, TaskQueue},
     templates::Templates,
 };
+use opentelemetry_http::HeaderExtractor;
 use tower::{make::Shared, ServiceBuilder};
 use tower_http::{
     compression::CompressionLayer,
     sensitive_headers::SetSensitiveHeadersLayer,
-    trace::{DefaultMakeSpan, DefaultOnResponse, TraceLayer},
-    LatencyUnit,
+    trace::{MakeSpan, OnResponse, TraceLayer},
 };
-use tracing::info;
+use tracing::{field, info};
 
 use super::RootCommand;
 
@@ -44,6 +44,72 @@ pub(super) struct ServerCommand {
     migrate: bool,
 }
 
+#[derive(Debug, Clone, Default)]
+struct OtelMakeSpan;
+
+impl<B> MakeSpan<B> for OtelMakeSpan {
+    fn make_span(&mut self, request: &hyper::Request<B>) -> tracing::Span {
+        // Extract the context from the headers
+        let headers = request.headers();
+        let extractor = HeaderExtractor(headers);
+
+        let cx = opentelemetry::global::get_text_map_propagator(|propagator| {
+            propagator.extract(&extractor)
+        });
+
+        // Attach the context so when the request span is created it gets properly
+        // parented
+        let _guard = cx.attach();
+
+        let version = match request.version() {
+            Version::HTTP_09 => "0.9",
+            Version::HTTP_10 => "1.0",
+            Version::HTTP_11 => "1.1",
+            Version::HTTP_2 => "2.0",
+            Version::HTTP_3 => "3.0",
+            _ => "",
+        };
+
+        let span = tracing::info_span!(
+            "request",
+            http.method = %request.method(),
+            http.target = %request.uri(),
+            http.flavor = version,
+            http.status_code = field::Empty,
+            http.user_agent = field::Empty,
+            otel.kind = "server",
+            otel.status_code = field::Empty,
+        );
+
+        if let Some(user_agent) = headers
+            .get(header::USER_AGENT)
+            .and_then(|s| s.to_str().ok())
+        {
+            span.record("http.user_agent", &user_agent);
+        }
+
+        span
+    }
+}
+
+#[derive(Debug, Clone, Default)]
+struct OtelOnResponse;
+
+impl<B> OnResponse<B> for OtelOnResponse {
+    fn on_response(self, response: &hyper::Response<B>, _latency: Duration, span: &tracing::Span) {
+        let s = response.status();
+        let status = if s.is_success() {
+            "ok"
+        } else if s.is_client_error() || s.is_server_error() {
+            "error"
+        } else {
+            "unset"
+        };
+        span.record("otel.status_code", &status);
+        span.record("http.status_code", &s.as_u16());
+    }
+}
+
 impl ServerCommand {
     pub async fn run(&self, root: &RootCommand) -> anyhow::Result<()> {
         let config: RootConfig = root.load_config()?;
@@ -80,12 +146,8 @@ impl ServerCommand {
             // Add high level tracing/logging to all requests
             .layer(
                 TraceLayer::new_for_http()
-                    .make_span_with(DefaultMakeSpan::new().include_headers(true))
-                    .on_response(
-                        DefaultOnResponse::new()
-                            .include_headers(true)
-                            .latency_unit(LatencyUnit::Micros),
-                    ),
+                    .make_span_with(OtelMakeSpan)
+                    .on_response(OtelOnResponse),
             )
             // Set a timeout
             .timeout(Duration::from_secs(10))