tower_http/trace/
make_span.rs

1use http::Request;
2use tracing::{Level, Span};
3
4use super::DEFAULT_MESSAGE_LEVEL;
5
6/// Trait used to generate [`Span`]s from requests. [`Trace`] wraps all request handling in this
7/// span.
8///
9/// [`Span`]: tracing::Span
10/// [`Trace`]: super::Trace
11pub trait MakeSpan<B> {
12    /// Make a span from a request.
13    fn make_span(&mut self, request: &Request<B>) -> Span;
14}
15
16impl<B> MakeSpan<B> for Span {
17    fn make_span(&mut self, _request: &Request<B>) -> Span {
18        self.clone()
19    }
20}
21
22impl<F, B> MakeSpan<B> for F
23where
24    F: FnMut(&Request<B>) -> Span,
25{
26    fn make_span(&mut self, request: &Request<B>) -> Span {
27        self(request)
28    }
29}
30
31/// The default way [`Span`]s will be created for [`Trace`].
32///
33/// [`Span`]: tracing::Span
34/// [`Trace`]: super::Trace
35#[derive(Debug, Clone)]
36pub struct DefaultMakeSpan {
37    level: Level,
38    include_headers: bool,
39}
40
41impl DefaultMakeSpan {
42    /// Create a new `DefaultMakeSpan`.
43    pub fn new() -> Self {
44        Self {
45            level: DEFAULT_MESSAGE_LEVEL,
46            include_headers: false,
47        }
48    }
49
50    /// Set the [`Level`] used for the [tracing span].
51    ///
52    /// Defaults to [`Level::DEBUG`].
53    ///
54    /// [tracing span]: https://docs.rs/tracing/latest/tracing/#spans
55    pub fn level(mut self, level: Level) -> Self {
56        self.level = level;
57        self
58    }
59
60    /// Include request headers on the [`Span`].
61    ///
62    /// By default headers are not included.
63    ///
64    /// [`Span`]: tracing::Span
65    pub fn include_headers(mut self, include_headers: bool) -> Self {
66        self.include_headers = include_headers;
67        self
68    }
69}
70
71impl Default for DefaultMakeSpan {
72    fn default() -> Self {
73        Self::new()
74    }
75}
76
77impl<B> MakeSpan<B> for DefaultMakeSpan {
78    fn make_span(&mut self, request: &Request<B>) -> Span {
79        // This ugly macro is needed, unfortunately, because `tracing::span!`
80        // required the level argument to be static. Meaning we can't just pass
81        // `self.level`.
82        macro_rules! make_span {
83            ($level:expr) => {
84                if self.include_headers {
85                    tracing::span!(
86                        $level,
87                        "request",
88                        method = %request.method(),
89                        uri = %request.uri(),
90                        version = ?request.version(),
91                        headers = ?request.headers(),
92                    )
93                } else {
94                    tracing::span!(
95                        $level,
96                        "request",
97                        method = %request.method(),
98                        uri = %request.uri(),
99                        version = ?request.version(),
100                    )
101                }
102            }
103        }
104
105        match self.level {
106            Level::ERROR => make_span!(Level::ERROR),
107            Level::WARN => make_span!(Level::WARN),
108            Level::INFO => make_span!(Level::INFO),
109            Level::DEBUG => make_span!(Level::DEBUG),
110            Level::TRACE => make_span!(Level::TRACE),
111        }
112    }
113}