use opentelemetry_http::HeaderExtractor;
use crate::{
http::headers::{self, HeaderMap, HeaderMapExt, HeaderName, HeaderValue},
otel::{
otel_trace::{FutureExt, Span, SpanKind, TraceContextExt, Tracer},
semantic_conventions::{attribute, trace},
*,
},
prelude::*,
};
pub struct TraceHoop {
tracer: ArcTracer,
}
impl Default for TraceHoop {
fn default() -> Self {
Self::new()
}
}
impl TraceHoop {
pub fn new() -> Self {
Self { tracer: arc_tracer() }
}
}
#[async_trait]
impl Handler for TraceHoop {
async fn handle(&self, req: &mut Request, depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl) {
let remote_addr = req.remote_addr().to_string();
let mut headers = HeaderMap::with_capacity(req.headers().len());
headers.extend(req.headers().into_iter().map(|(name, value)| {
let name = HeaderName::from_bytes(name.as_ref()).expect("Invalid header name");
let value = HeaderValue::from_bytes(value.as_ref()).expect("Invalid header value");
(name, value)
}));
let parent_cx = get_text_map_propagator(|propagator| propagator.extract(&HeaderExtractor(&headers)));
let mut attributes = Vec::new();
attributes.push(KeyValue::new(trace::HTTP_REQUEST_METHOD, req.method().to_string()));
attributes.push(KeyValue::new(trace::URL_FULL, req.uri().to_string()));
attributes.push(KeyValue::new(trace::CLIENT_ADDRESS, remote_addr));
attributes.push(KeyValue::new(
trace::NETWORK_PROTOCOL_VERSION,
format!("{:?}", req.version()),
));
let mut span = self
.tracer
.span_builder(format!("{} {}", req.method(), req.uri()))
.with_kind(SpanKind::Server)
.with_attributes(attributes)
.start_with_context(&self.tracer, &parent_cx);
span.add_event("request.started".to_string(), vec![]);
(async move {
ctrl.call_next(req, depot, res).await;
let cx = Context::current();
let span = cx.span();
let status = res.status_code.unwrap_or(StatusCode::OK);
let event = if status.is_client_error() || status.is_server_error() {
"request.failure"
} else {
"request.success"
};
span.add_event(event.to_string(), vec![]);
span.set_attribute(KeyValue::new(trace::HTTP_RESPONSE_STATUS_CODE, status.as_u16() as i64));
#[allow(clippy::cast_possible_truncation, clippy::cast_possible_wrap)]
if let Some(content_length) = res.headers().typed_get::<headers::ContentLength>() {
span.set_attribute(KeyValue::new(
attribute::HTTP_RESPONSE_BODY_SIZE,
content_length.0 as i64,
));
}
})
.with_context(Context::current_with_span(span))
.await
}
}