merge: #3159

3159: feat(telemetry): extract otel context & associate with HTTP request spans r=fnichol a=fnichol

<img src="https://media2.giphy.com/media/gLiNeyEhLoLlshprKm/giphy.gif"/>

Co-authored-by: Fletcher Nichol <[email protected]>

lib/telemetry-http-rs/BUCK

@@ -4,9 +4,11 @@ rust_library(
     name = "telemetry-http",
     deps = [
         "//lib/telemetry-rs:telemetry",
+        "//third-party/rust:http",
         "//third-party/rust:hyper",
         "//third-party/rust:remain",
         "//third-party/rust:tower-http",
+        "//third-party/rust:tracing-opentelemetry",
     ],
     srcs = glob(["src/**/*.rs"]),
 )

lib/telemetry-http-rs/Cargo.toml

@@ -6,7 +6,9 @@ rust-version = "1.64"
 publish = false
 
 [dependencies]
+http = { workspace = true }
 hyper = { workspace = true }
 remain = { workspace = true }
 telemetry = { path = "../../lib/telemetry-rs" }
 tower-http = { workspace = true }
+tracing-opentelemetry = { workspace = true }

lib/telemetry-http-rs/src/lib.rs

@@ -13,6 +13,8 @@
 
 mod make_span;
 mod on_response;
+mod propagation;
 
 pub use make_span::{HttpMakeSpan, NetworkTransport};
 pub use on_response::HttpOnResponse;
+pub use propagation::{extract_opentelemetry_context, inject_opentelemetry_context};

lib/telemetry-http-rs/src/make_span.rs

@@ -1,6 +1,9 @@
 use hyper::header::USER_AGENT;
 use telemetry::prelude::*;
 use tower_http::trace::MakeSpan;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+use crate::extract_opentelemetry_context;
 
 /// An implementation of [`MakeSpan`] to generate [`Span`]s from incoming HTTP requests.
 #[derive(Clone, Debug)]
@@ -255,6 +258,10 @@ impl HttpMakeSpan {
             );
         }
 
+        // Extract OpenTelemetry parent span metadata from the request headers (if it exists) and
+        // associate it with this request span
+        span.set_parent(extract_opentelemetry_context(request.headers()));
+
         span
     }
 }

lib/telemetry-http-rs/src/propagation.rs

@@ -0,0 +1,51 @@
+use http::HeaderMap;
+use telemetry::opentelemetry::{global, Context};
+
+/// Extracts an OpenTelemetry [`Context`] from a [`HeaderMap`].
+pub fn extract_opentelemetry_context(headers: &HeaderMap) -> Context {
+    let extractor = self::headers::HeaderExtractor(headers);
+    global::get_text_map_propagator(|propagator| propagator.extract(&extractor))
+}
+
+/// Injects an OpenTelemetry [`Context`] into a [`HeaderMap`].
+pub fn inject_opentelemetry_context(ctx: &Context, headers: &mut HeaderMap) {
+    let mut injector = self::headers::HeaderInjector(headers);
+    global::get_text_map_propagator(|propagator| propagator.inject_context(ctx, &mut injector));
+}
+
+// Implementation vendored from `opentelemetry-http` crate, released under the Apache 2.0 license
+//
+// https://github.com/open-telemetry/opentelemetry-rust/blob/47881b20a2b8e94d8e1cdbd4877852dd74cc07de/opentelemetry-http/src/lib.rs#L13-L41
+mod headers {
+    use telemetry::opentelemetry::propagation::{Extractor, Injector};
+
+    pub struct HeaderInjector<'a>(pub &'a mut http::HeaderMap);
+
+    impl<'a> Injector for HeaderInjector<'a> {
+        /// Set a key and value in the HeaderMap.  Does nothing if the key or value are not valid inputs.
+        fn set(&mut self, key: &str, value: String) {
+            if let Ok(name) = http::header::HeaderName::from_bytes(key.as_bytes()) {
+                if let Ok(val) = http::header::HeaderValue::from_str(&value) {
+                    self.0.insert(name, val);
+                }
+            }
+        }
+    }
+
+    pub struct HeaderExtractor<'a>(pub &'a http::HeaderMap);
+
+    impl<'a> Extractor for HeaderExtractor<'a> {
+        /// Get a value for a key from the HeaderMap.  If the value is not valid ASCII, returns None.
+        fn get(&self, key: &str) -> Option<&str> {
+            self.0.get(key).and_then(|value| value.to_str().ok())
+        }
+
+        /// Collect all the keys from the HeaderMap.
+        fn keys(&self) -> Vec<&str> {
+            self.0
+                .keys()
+                .map(|value| value.as_str())
+                .collect::<Vec<_>>()
+        }
+    }
+}