refactor(logging): unify request-id propagation and fallback metrics (#2652)

crates/ecstore/Cargo.toml

@@ -55,6 +55,7 @@ flatbuffers.workspace = true
 futures.workspace = true
 futures-util.workspace = true
 tracing.workspace = true
+tracing-opentelemetry.workspace = true
 serde.workspace = true
 time.workspace = true
 bytesize.workspace = true
@@ -62,6 +63,7 @@ serde_json.workspace = true
 quick-xml = { workspace = true, features = ["serialize", "async-tokio"] }
 s3s.workspace = true
 http.workspace = true
+opentelemetry.workspace = true
 http-body = { workspace = true }
 http-body-util.workspace = true
 url.workspace = true
@@ -127,6 +129,7 @@ criterion = { workspace = true, features = ["html_reports"] }
 temp-env = { workspace = true }
 tracing-subscriber = { workspace = true }
 serial_test = { workspace = true }
+opentelemetry_sdk = { workspace = true }
 
 [build-dependencies]
 shadow-rs = { workspace = true, features = ["build", "metadata"] }

crates/ecstore/src/rpc/client.rs

@@ -20,6 +20,8 @@ use std::error::Error;
 use tonic::{service::interceptor::InterceptedService, transport::Channel};
 use tracing::debug;
 
+use super::context_propagation::{inject_request_id_into_metadata, inject_trace_context_into_metadata};
+
 /// 3. Subsequent calls will attempt fresh connections
 /// 4. If node is still down, connection will fail fast (3s timeout)
 pub async fn node_service_time_out_client(
@@ -55,6 +57,8 @@ impl tonic::service::Interceptor for TonicSignatureInterceptor {
     fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
         let headers = gen_signature_headers(TONIC_RPC_PREFIX, &Method::GET);
         req.metadata_mut().as_mut().extend(headers);
+        inject_trace_context_into_metadata(req.metadata_mut());
+        inject_request_id_into_metadata(req.metadata_mut());
         Ok(req)
     }
 }
@@ -84,3 +88,34 @@ impl tonic::service::Interceptor for TonicInterceptor {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tonic::service::Interceptor;
+
+    #[test]
+    fn test_signature_interceptor_keeps_auth_headers() {
+        let mut interceptor = TonicSignatureInterceptor;
+        let req = tonic::Request::new(());
+
+        let req = interceptor.call(req).expect("interceptor call should succeed");
+
+        assert!(req.metadata().contains_key("x-rustfs-signature"));
+        assert!(req.metadata().contains_key("x-rustfs-timestamp"));
+    }
+
+    #[test]
+    fn test_signature_interceptor_may_inject_request_id() {
+        let mut interceptor = TonicSignatureInterceptor;
+        let req = tonic::Request::new(());
+
+        let span = tracing::info_span!("grpc-rpc-test-span");
+        let _guard = span.enter();
+        let req = interceptor.call(req).expect("interceptor call should succeed");
+
+        if let Some(v) = req.metadata().get("x-request-id") {
+            assert!(!v.as_encoded_bytes().is_empty());
+        }
+    }
+}

crates/ecstore/src/rpc/context_propagation.rs

@@ -0,0 +1,223 @@
+// Copyright 2024 RustFS Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use http::{HeaderMap, HeaderValue};
+use opentelemetry::{global, propagation::Injector, trace::TraceContextExt};
+use tracing::Span;
+use tracing_opentelemetry::OpenTelemetrySpanExt;
+
+pub(crate) const REQUEST_ID_HEADER: &str = "x-request-id";
+
+struct HttpHeaderInjector<'a> {
+    headers: &'a mut HeaderMap,
+}
+
+impl Injector for HttpHeaderInjector<'_> {
+    fn set(&mut self, key: &str, value: String) {
+        let Ok(name) = http::header::HeaderName::from_bytes(key.as_bytes()) else {
+            return;
+        };
+        let Ok(val) = HeaderValue::from_str(&value) else {
+            return;
+        };
+        self.headers.insert(name, val);
+    }
+}
+
+struct MetadataInjector<'a> {
+    metadata: &'a mut tonic::metadata::MetadataMap,
+}
+
+impl Injector for MetadataInjector<'_> {
+    fn set(&mut self, key: &str, value: String) {
+        let Ok(meta_key) = tonic::metadata::MetadataKey::from_bytes(key.as_bytes()) else {
+            return;
+        };
+        let Ok(meta_value) = tonic::metadata::MetadataValue::try_from(value.as_str()) else {
+            return;
+        };
+        self.metadata.insert(meta_key, meta_value);
+    }
+}
+
+fn current_trace_id() -> Option<String> {
+    let current_context = Span::current().context();
+    let current_span = current_context.span();
+    let span_context = current_span.span_context();
+    if !span_context.is_valid() {
+        return None;
+    }
+    Some(span_context.trace_id().to_string())
+}
+
+fn fallback_request_id() -> String {
+    format!("req-{}", &uuid::Uuid::new_v4().to_string()[..8])
+}
+
+fn propagated_request_id() -> String {
+    current_trace_id()
+        .map(|trace_id| format!("trace-{trace_id}"))
+        .unwrap_or_else(fallback_request_id)
+}
+
+pub(crate) fn inject_trace_context_into_http_headers(headers: &mut HeaderMap) {
+    let current_context = Span::current().context();
+    global::get_text_map_propagator(|propagator| {
+        let mut injector = HttpHeaderInjector { headers };
+        propagator.inject_context(&current_context, &mut injector);
+    });
+}
+
+pub(crate) fn inject_request_id_into_http_headers(headers: &mut HeaderMap) {
+    if headers.contains_key(REQUEST_ID_HEADER) {
+        return;
+    }
+    let request_id = propagated_request_id();
+    if let Ok(value) = HeaderValue::from_str(&request_id) {
+        headers.insert(REQUEST_ID_HEADER, value);
+    }
+}
+
+pub(crate) fn inject_trace_context_into_metadata(metadata: &mut tonic::metadata::MetadataMap) {
+    let current_context = Span::current().context();
+    global::get_text_map_propagator(|propagator| {
+        let mut injector = MetadataInjector { metadata };
+        propagator.inject_context(&current_context, &mut injector);
+    });
+}
+
+pub(crate) fn inject_request_id_into_metadata(metadata: &mut tonic::metadata::MetadataMap) {
+    let request_id_key = tonic::metadata::MetadataKey::from_static(REQUEST_ID_HEADER);
+    if metadata.contains_key(&request_id_key) {
+        return;
+    }
+    let request_id = propagated_request_id();
+    let Ok(value) = tonic::metadata::MetadataValue::try_from(request_id.as_str()) else {
+        return;
+    };
+    metadata.insert(request_id_key, value);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use opentelemetry::trace::{SpanContext, TraceContextExt, TraceFlags, TraceId, TraceState, TracerProvider as _};
+    use opentelemetry_sdk::trace::SdkTracerProvider;
+    use tracing_opentelemetry::OpenTelemetrySpanExt;
+    use tracing_subscriber::{Registry, layer::SubscriberExt};
+
+    fn with_trace_parent<F>(trace_id_hex: &str, f: F)
+    where
+        F: FnOnce(),
+    {
+        let provider = SdkTracerProvider::builder().build();
+        let tracer = provider.tracer("context-propagation-tests");
+        let subscriber = Registry::default().with(tracing_opentelemetry::layer().with_tracer(tracer));
+
+        tracing::subscriber::with_default(subscriber, || {
+            let span = tracing::info_span!("context-propagation-test-span");
+
+            let trace_id = TraceId::from_hex(trace_id_hex).expect("trace id should be valid hex");
+            let span_id = opentelemetry::trace::SpanId::from_hex("0102030405060708").expect("span id should be valid hex");
+            let parent = SpanContext::new(trace_id, span_id, TraceFlags::SAMPLED, true, TraceState::default());
+            span.set_parent(opentelemetry::Context::new().with_remote_span_context(parent))
+                .expect("failed to set parent context");
+            let _guard = span.enter();
+
+            f();
+        });
+        let _ = provider.shutdown();
+    }
+
+    #[test]
+    fn test_inject_request_id_into_http_headers_preserves_existing_value() {
+        let mut headers = HeaderMap::new();
+        headers.insert(REQUEST_ID_HEADER, HeaderValue::from_static("req-upstream-123"));
+
+        with_trace_parent("0123456789abcdef0123456789abcdef", || {
+            inject_request_id_into_http_headers(&mut headers);
+        });
+
+        assert_eq!(headers.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()), Some("req-upstream-123"));
+    }
+
+    #[test]
+    fn test_inject_request_id_into_http_headers_uses_trace_id_when_missing() {
+        let trace_id = "abcdefabcdefabcdefabcdefabcdefab";
+        let mut headers = HeaderMap::new();
+
+        with_trace_parent(trace_id, || {
+            inject_request_id_into_http_headers(&mut headers);
+        });
+
+        assert_eq!(
+            headers.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()),
+            Some(format!("trace-{trace_id}").as_str())
+        );
+    }
+
+    #[test]
+    fn test_inject_request_id_into_metadata_preserves_existing_value() {
+        let mut metadata = tonic::metadata::MetadataMap::new();
+        metadata.insert(
+            tonic::metadata::MetadataKey::from_static(REQUEST_ID_HEADER),
+            tonic::metadata::MetadataValue::from_static("req-upstream-456"),
+        );
+
+        with_trace_parent("fedcba9876543210fedcba9876543210", || {
+            inject_request_id_into_metadata(&mut metadata);
+        });
+
+        assert_eq!(metadata.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()), Some("req-upstream-456"));
+    }
+
+    #[test]
+    fn test_inject_request_id_into_metadata_uses_trace_id_when_missing() {
+        let trace_id = "1234567890abcdef1234567890abcdef";
+        let mut metadata = tonic::metadata::MetadataMap::new();
+
+        with_trace_parent(trace_id, || {
+            inject_request_id_into_metadata(&mut metadata);
+        });
+
+        assert_eq!(
+            metadata.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()),
+            Some(format!("trace-{trace_id}").as_str())
+        );
+    }
+
+    #[test]
+    fn test_inject_request_id_into_http_headers_uses_req_fallback_when_trace_missing() {
+        let mut headers = HeaderMap::new();
+        inject_request_id_into_http_headers(&mut headers);
+
+        let request_id = headers
+            .get(REQUEST_ID_HEADER)
+            .and_then(|v| v.to_str().ok())
+            .expect("request id should be injected");
+        assert!(request_id.starts_with("req-"), "expected req- fallback, got: {request_id}");
+    }
+
+    #[test]
+    fn test_inject_request_id_into_metadata_uses_req_fallback_when_trace_missing() {
+        let mut metadata = tonic::metadata::MetadataMap::new();
+        inject_request_id_into_metadata(&mut metadata);
+
+        let request_id = metadata
+            .get(REQUEST_ID_HEADER)
+            .and_then(|v| v.to_str().ok())
+            .expect("request id should be injected");
+        assert!(request_id.starts_with("req-"), "expected req- fallback, got: {request_id}");
+    }
+}

crates/ecstore/src/rpc/http_auth.rs

@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use crate::rpc::context_propagation::{inject_request_id_into_http_headers, inject_trace_context_into_http_headers};
 use base64::Engine as _;
 use base64::engine::general_purpose;
 use hmac::{Hmac, KeyInit, Mac};
@@ -62,6 +63,8 @@ pub fn build_auth_headers(url: &str, method: &Method, headers: &mut HeaderMap) {
     let auth_headers = gen_signature_headers(url, method);
 
     headers.extend(auth_headers);
+    inject_trace_context_into_http_headers(headers);
+    inject_request_id_into_http_headers(headers);
 }
 
 pub fn gen_signature_headers(url: &str, method: &Method) -> HeaderMap {
@@ -132,6 +135,7 @@ pub fn verify_rpc_signature(url: &str, method: &Method, headers: &HeaderMap) ->
 #[cfg(test)]
 mod tests {
     use super::*;
+    use crate::rpc::context_propagation::REQUEST_ID_HEADER;
     use http::{HeaderMap, Method};
     use time::OffsetDateTime;
 
@@ -210,6 +214,33 @@ mod tests {
         assert!((current_time - timestamp).abs() <= 1, "Timestamp should be close to current time");
     }
 
+    #[test]
+    fn test_build_auth_headers_preserves_existing_request_id() {
+        let url = "http://example.com/api/test";
+        let method = Method::GET;
+        let mut headers = HeaderMap::new();
+        headers.insert(REQUEST_ID_HEADER, HeaderValue::from_static("req-upstream-123"));
+
+        build_auth_headers(url, &method, &mut headers);
+
+        assert_eq!(headers.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()), Some("req-upstream-123"));
+    }
+
+    #[test]
+    fn test_build_auth_headers_may_set_request_id_from_trace_id() {
+        let url = "http://example.com/api/test";
+        let method = Method::GET;
+        let mut headers = HeaderMap::new();
+
+        let span = tracing::info_span!("rpc-test-span");
+        let _guard = span.enter();
+        build_auth_headers(url, &method, &mut headers);
+
+        if let Some(value) = headers.get(REQUEST_ID_HEADER).and_then(|v| v.to_str().ok()) {
+            assert!(!value.is_empty(), "request id should not be empty");
+        }
+    }
+
     #[test]
     fn test_verify_rpc_signature_success() {
         let url = "http://example.com/api/test";

crates/ecstore/src/rpc/mod.rs

@@ -13,6 +13,7 @@
 // limitations under the License.
 
 mod client;
+mod context_propagation;
 mod http_auth;
 mod peer_rest_client;
 mod peer_s3_client;