rama_http/layer/trace/
on_response.rs

1use super::{DEFAULT_MESSAGE_LEVEL, Latency};
2use crate::Response;
3use rama_core::telemetry::tracing::{Level, Span, field};
4use rama_utils::latency::LatencyUnit;
5use std::time::Duration;
6
7/// Trait used to tell [`Trace`] what to do when a response has been produced.
8///
9/// See the [module docs](../trace/index.html#on_response) for details on exactly when the
10/// `on_response` callback is called.
11///
12/// [`Trace`]: super::Trace
13pub trait OnResponse<B>: Send + Sync + 'static {
14    /// Do the thing.
15    ///
16    /// `latency` is the duration since the request was received.
17    ///
18    /// `span` is the `tracing` [`Span`], corresponding to this request, produced by the closure
19    /// passed to [`TraceLayer::make_span_with`]. It can be used to [record field values][record]
20    /// that weren't known when the span was created.
21    ///
22    /// [`Span`]: https://docs.rs/tracing/latest/tracing/span/index.html
23    /// [record]: https://docs.rs/tracing/latest/tracing/span/struct.Span.html#method.record
24    /// [`TraceLayer::make_span_with`]: crate::layer::trace::TraceLayer::make_span_with
25    fn on_response(self, response: &Response<B>, latency: Duration, span: &Span);
26}
27
28impl<B> OnResponse<B> for () {
29    #[inline]
30    fn on_response(self, _: &Response<B>, _: Duration, _: &Span) {}
31}
32
33impl<B, F> OnResponse<B> for F
34where
35    F: Fn(&Response<B>, Duration, &Span) + Send + Sync + 'static,
36{
37    fn on_response(self, response: &Response<B>, latency: Duration, span: &Span) {
38        self(response, latency, span)
39    }
40}
41
42/// The default [`OnResponse`] implementation used by [`Trace`].
43///
44/// [`Trace`]: super::Trace
45#[derive(Clone, Debug)]
46pub struct DefaultOnResponse {
47    level: Level,
48    latency_unit: LatencyUnit,
49    include_headers: bool,
50}
51
52impl Default for DefaultOnResponse {
53    fn default() -> Self {
54        Self {
55            level: DEFAULT_MESSAGE_LEVEL,
56            latency_unit: LatencyUnit::Millis,
57            include_headers: false,
58        }
59    }
60}
61
62impl DefaultOnResponse {
63    /// Create a new `DefaultOnResponse`.
64    #[must_use]
65    pub fn new() -> Self {
66        Self::default()
67    }
68
69    rama_utils::macros::generate_set_and_with! {
70        /// Set the [`Level`] used for [tracing events].
71        ///
72        /// Please note that while this will set the level for the tracing events
73        /// themselves, it might cause them to lack expected information, like
74        /// request method or path. You can address this using
75        /// [`DefaultMakeSpan::level`].
76        ///
77        /// Defaults to [`Level::DEBUG`].
78        ///
79        /// [tracing events]: https://docs.rs/tracing/latest/tracing/#events
80        /// [`DefaultMakeSpan::level`]: crate::layer::trace::DefaultMakeSpan::level
81        pub fn level(mut self, level: Level) -> Self {
82            self.level = level;
83            self
84        }
85    }
86
87    rama_utils::macros::generate_set_and_with! {
88        /// Set the [`LatencyUnit`] latencies will be reported in.
89        ///
90        /// Defaults to [`LatencyUnit::Millis`].
91        pub fn latency_unit(mut self, latency_unit: LatencyUnit) -> Self {
92            self.latency_unit = latency_unit;
93            self
94        }
95    }
96
97    rama_utils::macros::generate_set_and_with! {
98        /// Include response headers on the [`Event`].
99        ///
100        /// By default headers are not included.
101        ///
102        /// [`Event`]: tracing::Event
103        pub fn include_headers(mut self, include_headers: bool) -> Self {
104            self.include_headers = include_headers;
105            self
106        }
107    }
108}
109
110impl<B> OnResponse<B> for DefaultOnResponse {
111    fn on_response(self, response: &Response<B>, latency: Duration, span: &Span) {
112        let latency = Latency {
113            unit: self.latency_unit,
114            duration: latency,
115        };
116        let response_headers = self
117            .include_headers
118            .then(|| field::debug(response.headers()));
119
120        event_dynamic_lvl!(
121            parent: span,
122            self.level,
123            %latency,
124            status = status(response),
125            response_headers,
126            "finished processing request"
127        );
128    }
129}
130
131fn status<B>(res: &Response<B>) -> Option<i32> {
132    use crate::layer::classify::grpc_errors_as_failures::ParsedGrpcStatus;
133
134    // gRPC-over-HTTP2 uses the "application/grpc[+format]" content type, and gRPC-Web uses
135    // "application/grpc-web[+format]" or "application/grpc-web-text[+format]", where "format" is
136    // the message format, e.g. +proto, +json.
137    //
138    // So, valid grpc content types include (but are not limited to):
139    //  - application/grpc
140    //  - application/grpc+proto
141    //  - application/grpc-web+proto
142    //  - application/grpc-web-text+proto
143    //
144    // For simplicity, we simply check that the content type starts with "application/grpc".
145    let is_grpc = res
146        .headers()
147        .get(rama_http_types::header::CONTENT_TYPE)
148        .is_some_and(|value| value.as_bytes().starts_with("application/grpc".as_bytes()));
149
150    if is_grpc {
151        match crate::layer::classify::grpc_errors_as_failures::classify_grpc_metadata(
152            res.headers(),
153            crate::layer::classify::GrpcCode::Ok.into_bitmask(),
154        ) {
155            ParsedGrpcStatus::Success | ParsedGrpcStatus::HeaderNotGrpcCode => Some(0),
156            ParsedGrpcStatus::NonSuccess(status) => Some(status.code_raw()),
157            // if `grpc-status` is missing then its a streaming response and there is no status
158            // _yet_, so its neither success nor error
159            ParsedGrpcStatus::GrpcStatusHeaderMissing => None,
160        }
161    } else {
162        Some(res.status().as_u16().into())
163    }
164}
rama_http/layer/trace/on_response.rs

rama_http/layer/trace/
on_response.rs
