typeway_server/
request_id.rs1use std::convert::Infallible;
23use std::future::Future;
24use std::pin::Pin;
25use std::task::{Context, Poll};
26
27use crate::body::BoxBody;
28
29#[derive(Debug, Clone)]
31pub struct RequestId(pub String);
32
33#[derive(Clone)]
40pub struct RequestIdLayer;
41
42impl RequestIdLayer {
43 pub fn new() -> Self {
44 RequestIdLayer
45 }
46}
47
48impl Default for RequestIdLayer {
49 fn default() -> Self {
50 Self::new()
51 }
52}
53
54impl<S> tower_layer::Layer<S> for RequestIdLayer {
55 type Service = RequestIdService<S>;
56
57 fn layer(&self, inner: S) -> Self::Service {
58 RequestIdService { inner }
59 }
60}
61
62#[derive(Clone)]
64pub struct RequestIdService<S> {
65 inner: S,
66}
67
68impl<S, B> tower_service::Service<http::Request<B>> for RequestIdService<S>
69where
70 S: tower_service::Service<
71 http::Request<B>,
72 Response = http::Response<BoxBody>,
73 Error = Infallible,
74 > + Clone
75 + Send
76 + 'static,
77 S::Future: Send + 'static,
78 B: Send + 'static,
79{
80 type Response = http::Response<BoxBody>;
81 type Error = Infallible;
82 type Future = Pin<Box<dyn Future<Output = Result<Self::Response, Self::Error>> + Send>>;
83
84 fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
85 self.inner.poll_ready(cx)
86 }
87
88 fn call(&mut self, mut req: http::Request<B>) -> Self::Future {
89 let id = req
91 .headers()
92 .get("x-request-id")
93 .and_then(|v| v.to_str().ok())
94 .map(|s| s.to_string())
95 .unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
96
97 req.extensions_mut().insert(RequestId(id.clone()));
99
100 let mut inner = self.inner.clone();
101 Box::pin(async move {
102 let mut resp = inner.call(req).await?;
103 if let Ok(val) = http::HeaderValue::from_str(&id) {
105 resp.headers_mut().insert("x-request-id", val);
106 }
107 Ok(resp)
108 })
109 }
110}