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

rama_http/layer/trace/
layer.rs
