logo
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

//! Logging module
use std::time::Instant;

use tracing::{Instrument, Level};

use salvo_core::async_trait;
use salvo_core::http::{Request, Response, StatusCode};
use salvo_core::routing::FlowCtrl;
use salvo_core::{Depot, Handler};

/// LogHandler
#[derive(Default, Debug)]
pub struct LogHandler;

#[async_trait]
impl Handler for LogHandler {
    async fn handle(&self, req: &mut Request, depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl) {
        let span = tracing::span!(
            Level::INFO,
            "Request",
            remote_addr = %req.remote_addr().map(|addr|addr.to_string()).unwrap_or_else(|| "[Unknown]".into()),
            version = ?req.version(),
            method = %req.method(),
            path = %req.uri(),
        );

        async move {
            let now = Instant::now();
            ctrl.call_next(req, depot, res).await;
            let duration = now.elapsed();

            let status = match res.status_code() {
                Some(code) => code,
                None => {
                    if res.body().is_none() {
                        StatusCode::NOT_FOUND
                    } else {
                        StatusCode::OK
                    }
                }
            };
            tracing::info!(
                status = %status,
                duration = ?duration,
                "Response"
            );
        }
        .instrument(span)
        .await
    }
}

#[cfg(test)]
mod tests {
    use salvo_core::prelude::*;
    use tracing_test::traced_test;

    use super::*;

    #[tokio::test]
    #[traced_test]
    async fn test_log() {
        #[fn_handler]
        async fn hello() -> &'static str {
            "hello"
        }

        let router = Router::new()
            .hoop(LogHandler)
            .push(Router::with_path("hello").get(hello));
        let service = Service::new(router);

        let req: Request = hyper::Request::builder()
            .method("GET")
            .uri("http://127.0.0.1:7979/hello")
            .body(hyper::Body::empty())
            .unwrap()
            .into();
        service.handle(req).await.take_text().await.unwrap();
        assert!(logs_contain("duration"));
    }
}