添加qiming-rcoder模块

qiming-rcoder/crates/rcoder-telemetry/src/propagation.rs

@@ -0,0 +1,212 @@
+//! Trace Context 传播模块
+//!
+//! 提供跨服务的 trace context 传播功能，支持 gRPC 和 HTTP。
+
+use opentelemetry::propagation::{Extractor, Injector, TextMapPropagator};
+use opentelemetry::Context;
+use opentelemetry_sdk::propagation::TraceContextPropagator;
+use tonic::metadata::{MetadataKey, MetadataMap, MetadataValue};
+use tracing::debug;
+
+/// gRPC MetadataMap 的 Injector 实现
+struct MetadataMapInjector<'a>(&'a mut MetadataMap);
+
+impl Injector for MetadataMapInjector<'_> {
+    fn set(&mut self, key: &str, value: String) {
+        if let Ok(key) = MetadataKey::from_bytes(key.as_bytes()) {
+            if let Ok(value) = MetadataValue::try_from(&value) {
+                self.0.insert(key, value);
+            }
+        }
+    }
+}
+
+/// gRPC MetadataMap 的 Extractor 实现
+struct MetadataMapExtractor<'a>(&'a MetadataMap);
+
+impl Extractor for MetadataMapExtractor<'_> {
+    fn get(&self, key: &str) -> Option<&str> {
+        self.0.get(key).and_then(|value| value.to_str().ok())
+    }
+
+    fn keys(&self) -> Vec<&str> {
+        self.0
+            .keys()
+            .filter_map(|key| {
+                if let tonic::metadata::KeyRef::Ascii(k) = key {
+                    Some(k.as_str())
+                } else {
+                    None
+                }
+            })
+            .collect()
+    }
+}
+
+/// 注入 trace context 到 gRPC metadata
+///
+/// 将当前 span 的 trace context 注入到 gRPC metadata 中，
+/// 用于跨服务传播。
+///
+/// # Arguments
+///
+/// * `metadata` - gRPC metadata
+///
+/// # Example
+///
+/// ```no_run
+/// use tonic::metadata::MetadataMap;
+/// use rcoder_telemetry::propagation::inject_context;
+///
+/// let mut metadata = MetadataMap::new();
+/// inject_context(&mut metadata);
+/// // 现在 metadata 包含 traceparent 和 tracestate headers
+/// ```
+pub fn inject_context(metadata: &mut MetadataMap) {
+    let propagator = TraceContextPropagator::new();
+    let cx = Context::current();
+    let mut injector = MetadataMapInjector(metadata);
+    propagator.inject_context(&cx, &mut injector);
+
+    debug!("[Propagation] Trace context injected into gRPC metadata");
+}
+
+/// Extracting trace context from gRPC metadata
+///
+/// 从 gRPC metadata 中提取 trace context，
+/// 用于继续跨服务的 trace。
+///
+/// # Arguments
+///
+/// * `metadata` - gRPC metadata
+///
+/// # Returns
+///
+/// 返回提取的 `Context`，如果没有找到则返回当前 context。
+///
+/// # Example
+///
+/// ```no_run
+/// use tonic::metadata::MetadataMap;
+/// use rcoder_telemetry::propagation::extract_context;
+///
+/// let metadata = MetadataMap::new();
+/// let context = extract_context(&metadata);
+/// // 使用 context 创建新的 span
+/// ```
+pub fn extract_context(metadata: &MetadataMap) -> Context {
+    let propagator = TraceContextPropagator::new();
+    let extractor = MetadataMapExtractor(metadata);
+    let cx = propagator.extract(&extractor);
+
+    debug!("[Propagation] Extracting trace context from gRPC metadata");
+
+    cx
+}
+
+/// HTTP Headers 的 Injector 实现
+pub struct HttpHeaderInjector<'a>(pub &'a mut http::HeaderMap);
+
+impl Injector for HttpHeaderInjector<'_> {
+    fn set(&mut self, key: &str, value: String) {
+        if let Ok(name) = http::header::HeaderName::from_bytes(key.as_bytes()) {
+            if let Ok(value) = http::header::HeaderValue::from_str(&value) {
+                self.0.insert(name, value);
+            }
+        }
+    }
+}
+
+/// HTTP Headers 的 Extractor 实现
+pub struct HttpHeaderExtractor<'a>(pub &'a http::HeaderMap);
+
+impl Extractor for HttpHeaderExtractor<'_> {
+    fn get(&self, key: &str) -> Option<&str> {
+        self.0.get(key).and_then(|value| value.to_str().ok())
+    }
+
+    fn keys(&self) -> Vec<&str> {
+        self.0.keys().map(|key| key.as_str()).collect()
+    }
+}
+
+/// 注入 trace context 到 HTTP headers
+///
+/// # Arguments
+///
+/// * `headers` - HTTP headers
+pub fn inject_context_http(headers: &mut http::HeaderMap) {
+    let propagator = TraceContextPropagator::new();
+    let cx = Context::current();
+    let mut injector = HttpHeaderInjector(headers);
+    propagator.inject_context(&cx, &mut injector);
+
+    debug!("[Propagation] Trace context injected into HTTP headers");
+}
+
+/// Extracting trace context from HTTP headers
+///
+/// # Arguments
+///
+/// * `headers` - HTTP headers
+///
+/// # Returns
+///
+/// 返回提取的 `Context`
+pub fn extract_context_http(headers: &http::HeaderMap) -> Context {
+    let propagator = TraceContextPropagator::new();
+    let extractor = HttpHeaderExtractor(headers);
+    let cx = propagator.extract(&extractor);
+
+    debug!("[Propagation] Extracting trace context from HTTP headers");
+
+    cx
+}
+
+/// 设置全局 text map 传播器
+///
+/// 应该在应用启动时调用一次。
+pub fn set_global_propagator() {
+    opentelemetry::global::set_text_map_propagator(TraceContextPropagator::new());
+    debug!("[Propagation] Global TraceContextPropagator set");
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use opentelemetry::trace::TraceContextExt;
+
+    #[test]
+    fn test_metadata_injector_extractor() {
+        let mut metadata = MetadataMap::new();
+
+        // 手动设置一些 metadata
+        metadata.insert(
+            "traceparent",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+                .parse()
+                .unwrap(),
+        );
+
+        // 提取 context
+        let cx = extract_context(&metadata);
+        assert!(!cx.span().span_context().trace_id().to_string().is_empty());
+    }
+
+    #[test]
+    fn test_http_header_injector_extractor() {
+        let mut headers = http::HeaderMap::new();
+
+        // 手动设置 traceparent header
+        headers.insert(
+            "traceparent",
+            "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"
+                .parse()
+                .unwrap(),
+        );
+
+        // 提取 context
+        let cx = extract_context_http(&headers);
+        assert!(!cx.span().span_context().trace_id().to_string().is_empty());
+    }
+}