postrust_server/
telemetry.rs1use axum::extract::MatchedPath;
38use http::{HeaderValue, Request, Response};
39use pin_project_lite::pin_project;
40use postrust_core::trace::{valid_request_id, TraceContext, TraceId};
41use std::cell::Cell;
42use std::future::Future;
43use std::pin::Pin;
44use std::task::{Context, Poll};
45use std::time::Instant;
46use tower::{Layer, Service};
47use tracing::field::Empty;
48use tracing::Span;
49
50#[cfg(feature = "otel")]
51use crate::otel::RequestMeter;
52
53#[cfg(not(feature = "otel"))]
56struct RequestMeter;
57
58#[cfg(not(feature = "otel"))]
59impl RequestMeter {
60 #[inline(always)]
61 fn start(_: &http::Method, _: Option<&str>) -> Self {
62 RequestMeter
63 }
64
65 #[inline(always)]
66 fn finish(&self, _: http::StatusCode, _: Instant) {}
67}
68
69pub const REQUEST_ID_HEADER: &str = "x-request-id";
71
72#[derive(Clone, Copy, Debug, Default)]
75pub struct RequestTraceLayer;
76
77impl<S> Layer<S> for RequestTraceLayer {
78 type Service = RequestTrace<S>;
79
80 fn layer(&self, inner: S) -> Self::Service {
81 RequestTrace { inner }
82 }
83}
84
85#[derive(Clone, Debug)]
87pub struct RequestTrace<S> {
88 inner: S,
89}
90
91impl<S, B, ResBody> Service<Request<B>> for RequestTrace<S>
92where
93 S: Service<Request<B>, Response = Response<ResBody>>,
94{
95 type Response = Response<ResBody>;
96 type Error = S::Error;
97 type Future = ResponseFuture<S::Future>;
98
99 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
100 self.inner.poll_ready(cx)
101 }
102
103 fn call(&mut self, mut request: Request<B>) -> Self::Future {
104 let start = Instant::now();
105 let request_id = request_id(request.headers());
106 let trace = TraceContext::from_headers(request.headers())
107 .unwrap_or_else(|| TraceContext::root(new_trace_id()));
108
109 let span = tracing::info_span!(
113 "request",
114 "http.request.method" = request.method().as_str(),
115 "url.path" = request.uri().path(),
116 "http.route" = Empty,
117 "http.response.status_code" = Empty,
118 "otel.kind" = "server",
119 "otel.status_code" = Empty,
120 "otel.name" = Empty,
123 request_id = request_id.to_str().unwrap_or_default(),
125 trace_id = %trace.trace_id,
126 );
127 let route = request.extensions().get::<MatchedPath>();
133 if let Some(route) = route {
134 span.record("http.route", route.as_str());
135 }
136 let meter = RequestMeter::start(request.method(), route.map(MatchedPath::as_str));
137
138 #[cfg(feature = "otel")]
144 {
145 crate::otel::name_request_span(&span, request.method(), route.map(MatchedPath::as_str));
146 crate::otel::join_trace(&span, &trace);
147 }
148
149 request.extensions_mut().insert(trace);
150
151 let inner = {
152 let _entered = span.enter();
153 self.inner.call(request)
154 };
155 ResponseFuture {
156 inner,
157 span,
158 request_id: Some(request_id),
159 start,
160 meter,
161 }
162 }
163}
164
165pin_project! {
166 pub struct ResponseFuture<F> {
168 #[pin]
169 inner: F,
170 span: Span,
171 request_id: Option<HeaderValue>,
172 start: Instant,
173 meter: RequestMeter,
176 }
177}
178
179impl<F, ResBody, E> Future for ResponseFuture<F>
180where
181 F: Future<Output = Result<Response<ResBody>, E>>,
182{
183 type Output = Result<Response<ResBody>, E>;
184
185 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
186 let this = self.project();
187 let _entered = this.span.enter();
188 let mut response = match this.inner.poll(cx) {
189 Poll::Pending => return Poll::Pending,
190 Poll::Ready(Ok(response)) => response,
191 Poll::Ready(Err(e)) => return Poll::Ready(Err(e)),
192 };
193
194 let status = response.status();
195 this.span
196 .record("http.response.status_code", status.as_u16());
197 if status.is_server_error() {
200 this.span.record("otel.status_code", "ERROR");
201 }
202 this.meter.finish(status, *this.start);
203 tracing::debug!(
204 status = status.as_u16(),
205 latency_ms = this.start.elapsed().as_secs_f64() * 1000.0,
206 "request finished"
207 );
208
209 if let Some(id) = this.request_id.take() {
213 response
214 .headers_mut()
215 .entry(REQUEST_ID_HEADER)
216 .or_insert(id);
217 }
218 Poll::Ready(Ok(response))
219 }
220}
221
222pub fn request_id(headers: &http::HeaderMap) -> HeaderValue {
229 let mut sent = headers.get_all(REQUEST_ID_HEADER).iter();
230 match (sent.next(), sent.next()) {
231 (Some(id), None) if valid_request_id(id.as_bytes()) => id.clone(),
232 _ => new_request_id(),
233 }
234}
235
236pub fn new_request_id() -> HeaderValue {
238 let mut bytes = random_u128().to_be_bytes();
239 bytes[6] = (bytes[6] & 0x0f) | 0x40;
241 bytes[8] = (bytes[8] & 0x3f) | 0x80;
242
243 const HEX: &[u8; 16] = b"0123456789abcdef";
244 let mut text = [0u8; 36];
245 let mut at = 0;
246 for (i, b) in bytes.iter().enumerate() {
247 if matches!(i, 4 | 6 | 8 | 10) {
248 text[at] = b'-';
249 at += 1;
250 }
251 text[at] = HEX[(b >> 4) as usize];
252 text[at + 1] = HEX[(b & 0x0f) as usize];
253 at += 2;
254 }
255 HeaderValue::from_bytes(&text).expect("hex digits and hyphens are a valid header value")
256}
257
258pub fn new_trace_id() -> TraceId {
260 loop {
261 if let Some(id) = TraceId::from_bytes(random_u128().to_be_bytes()) {
263 return id;
264 }
265 }
266}
267
268thread_local! {
269 static RNG_STATE: Cell<u64> = Cell::new(seed());
270}
271
272fn seed() -> u64 {
277 use std::hash::BuildHasher;
278 std::collections::hash_map::RandomState::new().hash_one(0x5eed_u64)
279}
280
281fn random_u128() -> u128 {
289 fn next(state: &Cell<u64>) -> u64 {
290 let s = state.get().wrapping_add(0x9e37_79b9_7f4a_7c15);
291 state.set(s);
292 let mut z = s;
293 z = (z ^ (z >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9);
294 z = (z ^ (z >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb);
295 z ^ (z >> 31)
296 }
297 RNG_STATE.with(|state| ((next(state) as u128) << 64) | next(state) as u128)
298}
299
300#[cfg(test)]
301mod tests {
302 use super::*;
303
304 #[test]
305 fn a_generated_request_id_is_a_v4_uuid() {
306 let id = new_request_id();
307 let text = id.to_str().unwrap();
308 assert_eq!(text.len(), 36);
309 let parts: Vec<&str> = text.split('-').collect();
310 assert_eq!(
311 parts.iter().map(|p| p.len()).collect::<Vec<_>>(),
312 [8, 4, 4, 4, 12]
313 );
314 assert!(parts[2].starts_with('4'));
315 assert!(matches!(parts[3].as_bytes()[0], b'8' | b'9' | b'a' | b'b'));
316 assert!(valid_request_id(id.as_bytes()));
317 assert_ne!(new_request_id(), id);
318 }
319
320 #[test]
321 fn a_usable_request_id_is_kept() {
322 let mut headers = http::HeaderMap::new();
323 headers.insert(REQUEST_ID_HEADER, "abc-123".parse().unwrap());
324 assert_eq!(request_id(&headers), "abc-123");
325 }
326
327 #[test]
328 fn an_unusable_request_id_is_replaced() {
329 for bad in ["has space", &"x".repeat(129)] {
330 let mut headers = http::HeaderMap::new();
331 headers.insert(REQUEST_ID_HEADER, bad.parse().unwrap());
332 let id = request_id(&headers);
333 assert_ne!(id, bad);
334 assert_eq!(id.len(), 36);
335 }
336 let mut headers = http::HeaderMap::new();
338 headers.append(REQUEST_ID_HEADER, "one".parse().unwrap());
339 headers.append(REQUEST_ID_HEADER, "two".parse().unwrap());
340 assert_eq!(request_id(&headers).len(), 36);
341 }
342
343 #[test]
344 fn trace_ids_differ() {
345 assert_ne!(new_trace_id(), new_trace_id());
346 }
347}