122 lines
4.3 KiB
Rust
122 lines
4.3 KiB
Rust
use std::time::Instant;
|
|
|
|
use axum::{extract::Request, http::HeaderValue, middleware::Next, response::Response};
|
|
use crank_core::{CorrelationContext, RequestId, TraceContext};
|
|
use crank_metrics::ExemplarTraceId;
|
|
use crank_observability::{set_remote_trace_parent, with_request_correlation};
|
|
use tracing::{Instrument, info_span};
|
|
|
|
use crate::transport::HEADER_X_REQUEST_ID;
|
|
|
|
const HEADER_X_TRACE_ID: axum::http::HeaderName = axum::http::HeaderName::from_static("x-trace-id");
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(super) struct RequestContext {
|
|
pub(super) correlation: CorrelationContext,
|
|
started_at: Instant,
|
|
}
|
|
|
|
impl RequestContext {
|
|
pub(super) fn request_id(&self) -> &str {
|
|
self.correlation.request_id().as_str()
|
|
}
|
|
|
|
pub(super) fn started_at(&self) -> Instant {
|
|
self.started_at
|
|
}
|
|
|
|
pub(super) fn exemplar(&self) -> Option<ExemplarTraceId> {
|
|
self.correlation
|
|
.trace_context()
|
|
.is_sampled()
|
|
.then(|| ExemplarTraceId::parse(self.correlation.trace_id().as_str()))
|
|
.flatten()
|
|
}
|
|
}
|
|
|
|
pub(super) async fn apply_request_context(mut request: Request, next: Next) -> Response {
|
|
let (request_id, remote_parent) = resolve_correlation(request.headers());
|
|
let span = info_span!(
|
|
target: "crank::trace",
|
|
"mcp.request",
|
|
request_id = %request_id,
|
|
trace_id = tracing::field::Empty,
|
|
);
|
|
if let Some(remote_parent) = remote_parent.as_ref() {
|
|
set_canonical_parent(&span, remote_parent);
|
|
}
|
|
let trace_context = crank_trace::trace_context_for_span(&span).unwrap_or_else(|| {
|
|
remote_parent
|
|
.as_ref()
|
|
.map_or_else(TraceContext::generate, TraceContext::continue_local)
|
|
});
|
|
span.record("trace_id", trace_context.trace_id().as_str());
|
|
let context = RequestContext {
|
|
correlation: CorrelationContext::new(request_id, trace_context),
|
|
started_at: Instant::now(),
|
|
};
|
|
request.extensions_mut().insert(context.clone());
|
|
|
|
with_request_correlation(
|
|
context.correlation.request_id().to_string(),
|
|
context.correlation.trace_id().to_string(),
|
|
async move {
|
|
let mut response = next.run(request).instrument(span).await;
|
|
if let Ok(value) = HeaderValue::from_str(context.correlation.request_id().as_str()) {
|
|
response.headers_mut().insert(HEADER_X_REQUEST_ID, value);
|
|
}
|
|
if let Ok(value) = HeaderValue::from_str(context.correlation.trace_id().as_str()) {
|
|
response.headers_mut().insert(HEADER_X_TRACE_ID, value);
|
|
}
|
|
if let Some(exemplar) = context.exemplar() {
|
|
response.extensions_mut().insert(exemplar);
|
|
}
|
|
response
|
|
},
|
|
)
|
|
.await
|
|
}
|
|
|
|
fn resolve_correlation(headers: &axum::http::HeaderMap) -> (RequestId, Option<TraceContext>) {
|
|
let _tracestate_accepted = one_auxiliary_header_within_budget(
|
|
headers,
|
|
"tracestate",
|
|
TraceContext::tracestate_within_budget,
|
|
);
|
|
let _baggage_accepted =
|
|
one_auxiliary_header_within_budget(headers, "baggage", TraceContext::baggage_within_budget);
|
|
let mut request_ids = headers.get_all(HEADER_X_REQUEST_ID).iter();
|
|
let request_id = request_ids.next().and_then(|value| value.to_str().ok());
|
|
let request_id = if request_ids.next().is_some() {
|
|
RequestId::generate()
|
|
} else {
|
|
RequestId::resolve(request_id)
|
|
};
|
|
let mut traceparents = headers.get_all("traceparent").iter();
|
|
let traceparent = traceparents.next().and_then(|value| value.to_str().ok());
|
|
let remote_parent = if traceparents.next().is_some() {
|
|
None
|
|
} else {
|
|
traceparent.and_then(|value| TraceContext::parse(value).ok())
|
|
};
|
|
(request_id, remote_parent)
|
|
}
|
|
|
|
fn one_auxiliary_header_within_budget(
|
|
headers: &axum::http::HeaderMap,
|
|
name: &'static str,
|
|
validate: fn(&str) -> bool,
|
|
) -> bool {
|
|
let mut values = headers.get_all(name).iter();
|
|
let value = values.next().and_then(|value| value.to_str().ok());
|
|
values.next().is_none() && value.is_some_and(validate)
|
|
}
|
|
|
|
fn set_canonical_parent(span: &tracing::Span, context: &TraceContext) {
|
|
let mut headers = axum::http::HeaderMap::new();
|
|
if let Ok(value) = HeaderValue::from_str(context.traceparent()) {
|
|
headers.insert("traceparent", value);
|
|
set_remote_trace_parent(span, &headers);
|
|
}
|
|
}
|