use http::HeaderMap;
use crate::observability::{
current_span_id, current_trace_id, request_id_from_headers, span_id_from_traceparent,
trace_id_from_traceparent, traceparent_from_headers,
};
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CorrelationContext {
service: String,
transport: &'static str,
route: String,
method: String,
request_id: Option<String>,
traceparent: Option<String>,
trace_id: Option<String>,
span_id: Option<String>,
status: Option<String>,
}
impl CorrelationContext {
pub fn from_http_headers(
service: Option<&str>,
method: impl Into<String>,
route: Option<&str>,
headers: &HeaderMap,
) -> Self {
let traceparent = traceparent_from_headers(headers);
Self::new(
service.unwrap_or("unknown"),
"http",
route.unwrap_or("unknown"),
method,
)
.with_request_id(request_id_from_headers(headers))
.with_traceparent(traceparent)
.with_current_span_context()
}
#[cfg(feature = "rpc")]
pub fn from_rpc_metadata(
service: impl Into<String>,
method: impl Into<String>,
metadata: &tonic::metadata::MetadataMap,
) -> Self {
let method = method.into();
let traceparent = crate::observability::traceparent_from_metadata(metadata);
Self::new(service, "grpc", method.clone(), method)
.with_request_id(crate::observability::request_id_from_metadata(metadata))
.with_traceparent(traceparent)
.with_current_span_context()
}
pub fn from_rpc_parts(
service: impl Into<String>,
method: impl Into<String>,
request_id: Option<&str>,
traceparent: Option<&str>,
) -> Self {
let method = method.into();
Self::new(service, "grpc", method.clone(), method)
.with_request_id(request_id.map(ToOwned::to_owned))
.with_traceparent(traceparent.map(ToOwned::to_owned))
.with_current_span_context()
}
pub fn new(
service: impl Into<String>,
transport: &'static str,
route: impl Into<String>,
method: impl Into<String>,
) -> Self {
Self {
service: service.into(),
transport,
route: route.into(),
method: method.into(),
request_id: None,
traceparent: None,
trace_id: None,
span_id: None,
status: None,
}
}
pub fn with_status(mut self, status: impl Into<String>) -> Self {
self.status = Some(status.into());
self
}
pub fn service(&self) -> &str {
&self.service
}
pub fn transport(&self) -> &'static str {
self.transport
}
pub fn route(&self) -> &str {
&self.route
}
pub fn method(&self) -> &str {
&self.method
}
pub fn traceparent(&self) -> Option<&str> {
self.traceparent.as_deref()
}
pub fn request_id(&self) -> Option<&str> {
self.request_id.as_deref()
}
pub fn trace_id(&self) -> Option<&str> {
self.trace_id.as_deref()
}
pub fn span_id(&self) -> Option<&str> {
self.span_id.as_deref()
}
#[cfg(feature = "core")]
pub fn into_log_fields(self) -> crate::core::logging::LogFields {
use crate::core::logging::LogFields;
let mut fields = LogFields::new(self.service)
.with_transport(self.transport)
.with_route(self.route)
.with_method(self.method);
if let Some(request_id) = self.request_id {
fields = fields.with_request_id(request_id);
}
if let Some(trace_id) = self.trace_id {
fields = fields.with_trace_id(trace_id);
}
if let Some(span_id) = self.span_id {
fields = fields.with_span_id(span_id);
}
if let Some(status) = self.status {
fields = fields.with_status(status);
}
fields
}
pub fn as_pairs(&self) -> Vec<(String, String)> {
let mut pairs = vec![("service".to_string(), self.service.clone())];
push_optional(&mut pairs, "transport", Some(self.transport));
push_optional(&mut pairs, "route", Some(&self.route));
push_optional(&mut pairs, "method", Some(&self.method));
push_optional(&mut pairs, "request_id", self.request_id.as_deref());
push_optional(&mut pairs, "trace_id", self.trace_id.as_deref());
push_optional(&mut pairs, "span_id", self.span_id.as_deref());
push_optional(&mut pairs, "status", self.status.as_deref());
pairs
}
fn with_request_id(mut self, request_id: Option<String>) -> Self {
self.request_id = request_id;
self
}
fn with_traceparent(mut self, traceparent: Option<String>) -> Self {
if let Some(value) = traceparent {
self.trace_id = trace_id_from_traceparent(&value).map(ToOwned::to_owned);
self.span_id = span_id_from_traceparent(&value).map(ToOwned::to_owned);
self.traceparent = Some(value);
}
self
}
fn with_current_span_context(mut self) -> Self {
if let Some(trace_id) = current_trace_id() {
self.trace_id = Some(trace_id);
}
if let Some(span_id) = current_span_id() {
self.span_id = Some(span_id);
}
self
}
}
fn push_optional(pairs: &mut Vec<(String, String)>, key: &str, value: Option<&str>) {
if let Some(value) = value
&& !value.is_empty()
{
pairs.push((key.to_string(), value.to_string()));
}
}