rama_http/layer/trace/
make_span.rs

1use crate::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>: Send + Sync + 'static {
12    /// Make a span from a request.
13    fn make_span(&self, request: &Request<B>) -> Span;
14}
15
16impl<B> MakeSpan<B> for Span {
17    fn make_span(&self, _request: &Request<B>) -> Span {
18        self.clone()
19    }
20}
21
22impl<F, B> MakeSpan<B> for F
23where
24    F: Fn(&Request<B>) -> Span + Send + Sync + 'static,
25{
26    fn make_span(&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 const 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 const fn level(mut self, level: Level) -> Self {
56        self.level = level;
57        self
58    }
59
60    /// Set the [`Level`] used for the [tracing span].
61    ///
62    /// Defaults to [`Level::DEBUG`].
63    ///
64    /// [tracing span]: https://docs.rs/tracing/latest/tracing/#spans
65    pub fn set_level(&mut self, level: Level) -> &mut Self {
66        self.level = level;
67        self
68    }
69
70    /// Include request headers on the [`Span`].
71    ///
72    /// By default headers are not included.
73    ///
74    /// [`Span`]: tracing::Span
75    pub const fn include_headers(mut self, include_headers: bool) -> Self {
76        self.include_headers = include_headers;
77        self
78    }
79
80    /// Include request headers on the [`Span`].
81    ///
82    /// By default headers are not included.
83    ///
84    /// [`Span`]: tracing::Span
85    pub fn set_include_headers(&mut self, include_headers: bool) -> &mut Self {
86        self.include_headers = include_headers;
87        self
88    }
89}
90
91impl Default for DefaultMakeSpan {
92    fn default() -> Self {
93        Self::new()
94    }
95}
96
97impl<B> MakeSpan<B> for DefaultMakeSpan {
98    fn make_span(&self, request: &Request<B>) -> Span {
99        // This ugly macro is needed, unfortunately, because `tracing::span!`
100        // required the level argument to be static. Meaning we can't just pass
101        // `self.level`.
102        macro_rules! make_span {
103            ($level:expr) => {
104                if self.include_headers {
105                    tracing::span!(
106                        $level,
107                        "request",
108                        method = %request.method(),
109                        uri = %request.uri(),
110                        version = ?request.version(),
111                        headers = ?request.headers(),
112                    )
113                } else {
114                    tracing::span!(
115                        $level,
116                        "request",
117                        method = %request.method(),
118                        uri = %request.uri(),
119                        version = ?request.version(),
120                    )
121                }
122            }
123        }
124
125        match self.level {
126            Level::ERROR => make_span!(Level::ERROR),
127            Level::WARN => make_span!(Level::WARN),
128            Level::INFO => make_span!(Level::INFO),
129            Level::DEBUG => make_span!(Level::DEBUG),
130            Level::TRACE => make_span!(Level::TRACE),
131        }
132    }
133}
rama_http/layer/trace/make_span.rs

rama_http/layer/trace/
make_span.rs
