tdm_server_rust/telemetry/
propagation.rs

1//! traceparent 入站提取与出站注入
2
3use axum::http::{HeaderMap, HeaderName, HeaderValue, Request};
4use opentelemetry::{
5    global,
6    propagation::{Extractor, Injector},
7    Context,
8};
9use opentelemetry_sdk::propagation::TraceContextPropagator;
10use std::str::FromStr;
11
12/// HTTP 请求头 trace 上下文提取器
13struct HeaderExtractor<'a> {
14    /// 请求头
15    headers: &'a HeaderMap,
16}
17
18impl Extractor for HeaderExtractor<'_> {
19    fn get(&self, key: &str) -> Option<&str> {
20        self.headers.get(key).and_then(|v| v.to_str().ok())
21    }
22
23    fn keys(&self) -> Vec<&str> {
24        self.headers.keys().map(|k| k.as_str()).collect()
25    }
26}
27
28/// HTTP 响应头 trace 上下文注入器
29struct HeaderInjector<'a> {
30    /// 响应头
31    headers: &'a mut HeaderMap,
32}
33
34impl Injector for HeaderInjector<'_> {
35    fn set(&mut self, key: &str, value: String) {
36        if let (Ok(name), Ok(val)) = (HeaderName::from_str(key), HeaderValue::from_str(&value)) {
37            self.headers.insert(name, val);
38        }
39    }
40}
41
42/// 从入站 HTTP 头提取父 trace 上下文
43pub fn extract_parent(headers: &HeaderMap) -> Context {
44    global::get_text_map_propagator(|prop| prop.extract(&HeaderExtractor { headers }))
45}
46
47/// 将当前 span 上下文写入响应头 `traceparent`
48pub fn inject_response_headers(headers: &mut HeaderMap) {
49    global::get_text_map_propagator(|prop| {
50        prop.inject_context(
51            &opentelemetry::Context::current(),
52            &mut HeaderInjector { headers },
53        );
54    });
55}
56
57/// 为 reqwest 出站请求注入 trace 上下文
58pub fn inject_reqwest(headers: &mut reqwest::header::HeaderMap) {
59    struct ReqwestInjector<'a>(&'a mut reqwest::header::HeaderMap);
60
61    impl Injector for ReqwestInjector<'_> {
62        fn set(&mut self, key: &str, value: String) {
63            if let (Ok(name), Ok(val)) = (
64                reqwest::header::HeaderName::from_str(key),
65                reqwest::header::HeaderValue::from_str(&value),
66            ) {
67                self.0.insert(name, val);
68            }
69        }
70    }
71
72    global::get_text_map_propagator(|prop| {
73        prop.inject_context(
74            &opentelemetry::Context::current(),
75            &mut ReqwestInjector(headers),
76        );
77    });
78}
79
80/// 安装 W3C TraceContext 传播器（进程启动时调用一次）
81pub fn install_propagator() {
82    global::set_text_map_propagator(TraceContextPropagator::new());
83}
84
85/// 从 axum Request 提取父上下文（便捷）
86pub fn extract_from_request<B>(req: &Request<B>) -> Context {
87    extract_parent(req.headers())
88}
tdm_server_rust/telemetry/propagation.rs

tdm_server_rust/telemetry/
propagation.rs
