wae_https/middleware/
tracing.rs

1//! 追踪中间件
2//!
3//! 为请求添加追踪日志。
4
5use http::{Request, Response};
6use pin_project_lite::pin_project;
7use std::{
8    future::Future,
9    pin::Pin,
10    task::{Context, Poll},
11};
12use tower::{Layer, Service};
13use tracing::{info, info_span};
14
15/// 追踪中间件层
16#[derive(Debug, Clone)]
17pub struct TracingLayer;
18
19impl<S> Layer<S> for TracingLayer {
20    type Service = TracingService<S>;
21
22    fn layer(&self, inner: S) -> Self::Service {
23        TracingService { inner }
24    }
25}
26
27/// 追踪服务
28#[derive(Debug, Clone)]
29pub struct TracingService<S> {
30    inner: S,
31}
32
33impl<S, ReqBody, ResBody> Service<Request<ReqBody>> for TracingService<S>
34where
35    S: Service<Request<ReqBody>, Response = Response<ResBody>>,
36    S::Future: Send + 'static,
37{
38    type Response = S::Response;
39    type Error = S::Error;
40    type Future = TracingFuture<S::Future>;
41
42    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
43        self.inner.poll_ready(cx)
44    }
45
46    fn call(&mut self, req: Request<ReqBody>) -> Self::Future {
47        let method = req.method().to_string();
48        let uri = req.uri().to_string();
49        let span = info_span!("http_request", method = %method, uri = %uri);
50
51        let future = self.inner.call(req);
52        TracingFuture { inner: future, span }
53    }
54}
55
56pin_project! {
57    /// 追踪未来
58    pub struct TracingFuture<F> {
59        #[pin]
60        inner: F,
61        span: tracing::Span,
62    }
63}
64
65impl<F, Res, E> Future for TracingFuture<F>
66where
67    F: Future<Output = Result<Res, E>>,
68{
69    type Output = Result<Res, E>;
70
71    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
72        let this = self.project();
73        let _enter = this.span.enter();
74        info!("Request started");
75
76        match this.inner.poll(cx) {
77            Poll::Ready(result) => {
78                info!("Request completed");
79                Poll::Ready(result)
80            }
81            Poll::Pending => Poll::Pending,
82        }
83    }
84}
wae_https/middleware/tracing.rs

wae_https/middleware/
tracing.rs
