rama_http/layer/trace/
layer.rs

1use super::{
2    DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnRequest,
3    DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, Trace,
4};
5use crate::layer::classify::{
6    GrpcErrorsAsFailures, MakeClassifier, ServerErrorsAsFailures, SharedClassifier,
7};
8use rama_core::Layer;
9use std::fmt;
10
11/// [`Layer`] that adds high level [tracing] to a [`Service`].
12///
13/// See the [module docs](crate::layer::trace) for more details.
14///
15/// [`Layer`]: rama_core::Layer
16/// [tracing]: https://crates.io/crates/tracing
17/// [`Service`]: rama_core::Service
18pub struct TraceLayer<
19    M,
20    MakeSpan = DefaultMakeSpan,
21    OnRequest = DefaultOnRequest,
22    OnResponse = DefaultOnResponse,
23    OnBodyChunk = DefaultOnBodyChunk,
24    OnEos = DefaultOnEos,
25    OnFailure = DefaultOnFailure,
26> {
27    pub(crate) make_classifier: M,
28    pub(crate) make_span: MakeSpan,
29    pub(crate) on_request: OnRequest,
30    pub(crate) on_response: OnResponse,
31    pub(crate) on_body_chunk: OnBodyChunk,
32    pub(crate) on_eos: OnEos,
33    pub(crate) on_failure: OnFailure,
34}
35
36impl<
37    M: fmt::Debug,
38    MakeSpan: fmt::Debug,
39    OnRequest: fmt::Debug,
40    OnResponse: fmt::Debug,
41    OnBodyChunk: fmt::Debug,
42    OnEos: fmt::Debug,
43    OnFailure: fmt::Debug,
44> fmt::Debug for TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
45{
46    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
47        f.debug_struct("TraceLayer")
48            .field("make_classifier", &self.make_classifier)
49            .field("make_span", &self.make_span)
50            .field("on_request", &self.on_request)
51            .field("on_response", &self.on_response)
52            .field("on_body_chunk", &self.on_body_chunk)
53            .field("on_eos", &self.on_eos)
54            .field("on_failure", &self.on_failure)
55            .finish()
56    }
57}
58
59impl<
60    M: Clone,
61    MakeSpan: Clone,
62    OnRequest: Clone,
63    OnResponse: Clone,
64    OnBodyChunk: Clone,
65    OnEos: Clone,
66    OnFailure: Clone,
67> Clone for TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
68{
69    fn clone(&self) -> Self {
70        Self {
71            make_classifier: self.make_classifier.clone(),
72            make_span: self.make_span.clone(),
73            on_request: self.on_request.clone(),
74            on_response: self.on_response.clone(),
75            on_body_chunk: self.on_body_chunk.clone(),
76            on_eos: self.on_eos.clone(),
77            on_failure: self.on_failure.clone(),
78        }
79    }
80}
81
82impl<M> TraceLayer<M> {
83    /// Create a new [`TraceLayer`] using the given [`MakeClassifier`].
84    pub fn new(make_classifier: M) -> Self
85    where
86        M: MakeClassifier,
87    {
88        Self {
89            make_classifier,
90            make_span: DefaultMakeSpan::new(),
91            on_failure: DefaultOnFailure::default(),
92            on_request: DefaultOnRequest::default(),
93            on_eos: DefaultOnEos::default(),
94            on_body_chunk: DefaultOnBodyChunk::default(),
95            on_response: DefaultOnResponse::default(),
96        }
97    }
98}
99
100impl<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
101    TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
102{
103    /// Customize what to do when a request is received.
104    ///
105    /// `NewOnRequest` is expected to implement [`OnRequest`].
106    ///
107    /// [`OnRequest`]: super::OnRequest
108    pub fn on_request<NewOnRequest>(
109        self,
110        new_on_request: NewOnRequest,
111    ) -> TraceLayer<M, MakeSpan, NewOnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure> {
112        TraceLayer {
113            on_request: new_on_request,
114            on_failure: self.on_failure,
115            on_eos: self.on_eos,
116            on_body_chunk: self.on_body_chunk,
117            make_span: self.make_span,
118            on_response: self.on_response,
119            make_classifier: self.make_classifier,
120        }
121    }
122
123    /// Customize what to do when a response has been produced.
124    ///
125    /// `NewOnResponse` is expected to implement [`OnResponse`].
126    ///
127    /// [`OnResponse`]: super::OnResponse
128    pub fn on_response<NewOnResponse>(
129        self,
130        new_on_response: NewOnResponse,
131    ) -> TraceLayer<M, MakeSpan, OnRequest, NewOnResponse, OnBodyChunk, OnEos, OnFailure> {
132        TraceLayer {
133            on_response: new_on_response,
134            on_request: self.on_request,
135            on_eos: self.on_eos,
136            on_body_chunk: self.on_body_chunk,
137            on_failure: self.on_failure,
138            make_span: self.make_span,
139            make_classifier: self.make_classifier,
140        }
141    }
142
143    /// Customize what to do when a body chunk has been sent.
144    ///
145    /// `NewOnBodyChunk` is expected to implement [`OnBodyChunk`].
146    ///
147    /// [`OnBodyChunk`]: super::OnBodyChunk
148    pub fn on_body_chunk<NewOnBodyChunk>(
149        self,
150        new_on_body_chunk: NewOnBodyChunk,
151    ) -> TraceLayer<M, MakeSpan, OnRequest, OnResponse, NewOnBodyChunk, OnEos, OnFailure> {
152        TraceLayer {
153            on_body_chunk: new_on_body_chunk,
154            on_eos: self.on_eos,
155            on_failure: self.on_failure,
156            on_request: self.on_request,
157            make_span: self.make_span,
158            on_response: self.on_response,
159            make_classifier: self.make_classifier,
160        }
161    }
162
163    /// Customize what to do when a streaming response has closed.
164    ///
165    /// `NewOnEos` is expected to implement [`OnEos`].
166    ///
167    /// [`OnEos`]: super::OnEos
168    pub fn on_eos<NewOnEos>(
169        self,
170        new_on_eos: NewOnEos,
171    ) -> TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, NewOnEos, OnFailure> {
172        TraceLayer {
173            on_eos: new_on_eos,
174            on_body_chunk: self.on_body_chunk,
175            on_failure: self.on_failure,
176            on_request: self.on_request,
177            make_span: self.make_span,
178            on_response: self.on_response,
179            make_classifier: self.make_classifier,
180        }
181    }
182
183    /// Customize what to do when a response has been classified as a failure.
184    ///
185    /// `NewOnFailure` is expected to implement [`OnFailure`].
186    ///
187    /// [`OnFailure`]: super::OnFailure
188    pub fn on_failure<NewOnFailure>(
189        self,
190        new_on_failure: NewOnFailure,
191    ) -> TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, NewOnFailure> {
192        TraceLayer {
193            on_failure: new_on_failure,
194            on_request: self.on_request,
195            on_eos: self.on_eos,
196            on_body_chunk: self.on_body_chunk,
197            make_span: self.make_span,
198            on_response: self.on_response,
199            make_classifier: self.make_classifier,
200        }
201    }
202
203    /// Customize how to make [`Span`]s that all request handling will be wrapped in.
204    ///
205    /// `NewMakeSpan` is expected to implement [`MakeSpan`].
206    ///
207    /// [`MakeSpan`]: super::MakeSpan
208    /// [`Span`]: tracing::Span
209    pub fn make_span_with<NewMakeSpan>(
210        self,
211        new_make_span: NewMakeSpan,
212    ) -> TraceLayer<M, NewMakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure> {
213        TraceLayer {
214            make_span: new_make_span,
215            on_request: self.on_request,
216            on_failure: self.on_failure,
217            on_body_chunk: self.on_body_chunk,
218            on_eos: self.on_eos,
219            on_response: self.on_response,
220            make_classifier: self.make_classifier,
221        }
222    }
223}
224
225impl TraceLayer<HttpMakeClassifier> {
226    /// Create a new [`TraceLayer`] using [`ServerErrorsAsFailures`] which supports classifying
227    /// regular HTTP responses based on the status code.
228    pub fn new_for_http() -> Self {
229        Self {
230            make_classifier: SharedClassifier::new(ServerErrorsAsFailures::default()),
231            make_span: DefaultMakeSpan::new(),
232            on_response: DefaultOnResponse::default(),
233            on_request: DefaultOnRequest::default(),
234            on_body_chunk: DefaultOnBodyChunk::default(),
235            on_eos: DefaultOnEos::default(),
236            on_failure: DefaultOnFailure::default(),
237        }
238    }
239}
240
241impl TraceLayer<GrpcMakeClassifier> {
242    /// Create a new [`TraceLayer`] using [`GrpcErrorsAsFailures`] which supports classifying
243    /// gRPC responses and streams based on the `grpc-status` header.
244    pub fn new_for_grpc() -> Self {
245        Self {
246            make_classifier: SharedClassifier::new(GrpcErrorsAsFailures::default()),
247            make_span: DefaultMakeSpan::new(),
248            on_response: DefaultOnResponse::default(),
249            on_request: DefaultOnRequest::default(),
250            on_body_chunk: DefaultOnBodyChunk::default(),
251            on_eos: DefaultOnEos::default(),
252            on_failure: DefaultOnFailure::default(),
253        }
254    }
255}
256
257impl<S, M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure> Layer<S>
258    for TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
259where
260    M: Clone,
261    MakeSpan: Clone,
262    OnRequest: Clone,
263    OnResponse: Clone,
264    OnEos: Clone,
265    OnBodyChunk: Clone,
266    OnFailure: Clone,
267{
268    type Service = Trace<S, M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>;
269
270    fn layer(&self, inner: S) -> Self::Service {
271        Trace {
272            inner,
273            make_classifier: self.make_classifier.clone(),
274            make_span: self.make_span.clone(),
275            on_request: self.on_request.clone(),
276            on_eos: self.on_eos.clone(),
277            on_body_chunk: self.on_body_chunk.clone(),
278            on_response: self.on_response.clone(),
279            on_failure: self.on_failure.clone(),
280        }
281    }
282
283    fn into_layer(self, inner: S) -> Self::Service {
284        Trace {
285            inner,
286            make_classifier: self.make_classifier,
287            make_span: self.make_span,
288            on_request: self.on_request,
289            on_eos: self.on_eos,
290            on_body_chunk: self.on_body_chunk,
291            on_response: self.on_response,
292            on_failure: self.on_failure,
293        }
294    }
295}
rama_http/layer/trace/layer.rs

rama_http/layer/trace/
layer.rs
