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| {
45        prop.extract(&HeaderExtractor { headers })
46    })
47}
48
49/// 将当前 span 上下文写入响应头 `traceparent`
50pub fn inject_response_headers(headers: &mut HeaderMap) {
51    global::get_text_map_propagator(|prop| {
52        prop.inject_context(
53            &opentelemetry::Context::current(),
54            &mut HeaderInjector { headers },
55        );
56    });
57}
58
59/// 为 reqwest 出站请求注入 trace 上下文
60pub fn inject_reqwest(headers: &mut reqwest::header::HeaderMap) {
61    struct ReqwestInjector<'a>(&'a mut reqwest::header::HeaderMap);
62
63    impl Injector for ReqwestInjector<'_> {
64        fn set(&mut self, key: &str, value: String) {
65            if let (Ok(name), Ok(val)) = (
66                reqwest::header::HeaderName::from_str(key),
67                reqwest::header::HeaderValue::from_str(&value),
68            ) {
69                self.0.insert(name, val);
70            }
71        }
72    }
73
74    global::get_text_map_propagator(|prop| {
75        prop.inject_context(
76            &opentelemetry::Context::current(),
77            &mut ReqwestInjector(headers),
78        );
79    });
80}
81
82/// 安装 W3C TraceContext 传播器（进程启动时调用一次）
83pub fn install_propagator() {
84    global::set_text_map_propagator(TraceContextPropagator::new());
85}
86
87/// 从 axum Request 提取父上下文（便捷）
88pub fn extract_from_request<B>(req: &Request<B>) -> Context {
89    extract_parent(req.headers())
90}
tdm_server_rust/telemetry/propagation.rs

tdm_server_rust/telemetry/
propagation.rs
