hwhkit 0.6.0-alpha.2

A batteries-included Rust web service toolkit (production-ready: health, metrics, OTel, graceful shutdown, JWT, rate-limit, idempotency, circuit-breaker, scheduler)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//! Request-ID middleware: read incoming `x-request-id` header (configurable)
//! or generate a UUID v7. Inject into the current tracing span and echo
//! back to the client in the response header.

use std::convert::Infallible;
use std::task::{Context, Poll};

use axum::http::header::HeaderName;
use axum::http::{HeaderValue, Request, Response};
use futures::future::BoxFuture;
use tower::{Layer, Service};
use uuid::Uuid;

const TRACING_FIELD: &str = "request_id";

#[derive(Clone)]
pub struct RequestIdLayer {
    header_name: HeaderName,
}

impl RequestIdLayer {
    pub fn new(header: &str) -> Self {
        let header_name = HeaderName::from_bytes(header.as_bytes())
            .unwrap_or_else(|_| HeaderName::from_static("x-request-id"));
        Self { header_name }
    }
}

impl Default for RequestIdLayer {
    fn default() -> Self {
        Self::new("x-request-id")
    }
}

impl<S> Layer<S> for RequestIdLayer {
    type Service = RequestIdService<S>;
    fn layer(&self, inner: S) -> Self::Service {
        RequestIdService {
            inner,
            header_name: self.header_name.clone(),
        }
    }
}

#[derive(Clone)]
pub struct RequestIdService<S> {
    inner: S,
    header_name: HeaderName,
}

impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for RequestIdService<S>
where
    S: Service<Request<ReqBody>, Response = Response<ResBody>, Error = Infallible>
        + Clone
        + Send
        + 'static,
    S::Future: Send + 'static,
    ReqBody: Send + 'static,
    ResBody: Send + 'static,
{
    type Response = Response<ResBody>;
    type Error = Infallible;
    type Future = BoxFuture<'static, Result<Self::Response, Self::Error>>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, mut req: Request<ReqBody>) -> Self::Future {
        let header_name = self.header_name.clone();
        let id = req
            .headers()
            .get(&header_name)
            .and_then(|h| h.to_str().ok())
            .map(|s| s.to_string())
            .unwrap_or_else(|| Uuid::now_v7().to_string());

        // Stamp the inbound request so handlers can read it via the same header.
        if let Ok(value) = HeaderValue::from_str(&id) {
            req.headers_mut().insert(header_name.clone(), value);
        }

        // Record on the current tracing span so log lines pick it up.
        tracing::Span::current().record(TRACING_FIELD, tracing::field::display(&id));

        let clone = self.inner.clone();
        let mut inner = std::mem::replace(&mut self.inner, clone);
        Box::pin(async move {
            let mut response = inner.call(req).await?;
            if let Ok(value) = HeaderValue::from_str(&id) {
                response.headers_mut().insert(header_name, value);
            }
            Ok(response)
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use axum::body::Body;
    use axum::http::{Request, StatusCode};
    use axum::routing::get;
    use axum::Router;
    use tower::{Service, ServiceExt};

    /// Regression for F9: the layered service must remain ready across
    /// multiple invocations. Three sequential calls is enough to catch
    /// a poll_ready-after-call mistake under the typical tower runtime.
    #[tokio::test]
    async fn request_id_layer_handles_repeated_calls() {
        let app: Router = Router::new()
            .route("/x", get(|| async { "ok" }))
            .layer(RequestIdLayer::default());

        for _ in 0..3 {
            let mut svc = app.clone();
            let response = ServiceExt::<Request<Body>>::ready(&mut svc)
                .await
                .expect("ready")
                .call(Request::builder().uri("/x").body(Body::empty()).unwrap())
                .await
                .expect("call");
            assert_eq!(response.status(), StatusCode::OK);
            assert!(response.headers().get("x-request-id").is_some());
        }
    }
}