tower_http/trace/
layer.rs

1use super::{
2    DefaultMakeSpan, DefaultOnBodyChunk, DefaultOnEos, DefaultOnFailure, DefaultOnRequest,
3    DefaultOnResponse, GrpcMakeClassifier, HttpMakeClassifier, Trace,
4};
5use crate::classify::{
6    GrpcErrorsAsFailures, MakeClassifier, ServerErrorsAsFailures, SharedClassifier,
7};
8use tower_layer::Layer;
9
10/// [`Layer`] that adds high level [tracing] to a [`Service`].
11///
12/// See the [module docs](crate::trace) for more details.
13///
14/// [`Layer`]: tower_layer::Layer
15/// [tracing]: https://crates.io/crates/tracing
16/// [`Service`]: tower_service::Service
17#[derive(Debug, Copy, 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`]: 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    pub fn new_for_http() -> Self {
183        Self {
184            make_classifier: SharedClassifier::new(ServerErrorsAsFailures::default()),
185            make_span: DefaultMakeSpan::new(),
186            on_response: DefaultOnResponse::default(),
187            on_request: DefaultOnRequest::default(),
188            on_body_chunk: DefaultOnBodyChunk::default(),
189            on_eos: DefaultOnEos::default(),
190            on_failure: DefaultOnFailure::default(),
191        }
192    }
193}
194
195impl TraceLayer<GrpcMakeClassifier> {
196    /// Create a new [`TraceLayer`] using [`GrpcErrorsAsFailures`] which supports classifying
197    /// gRPC responses and streams based on the `grpc-status` header.
198    pub fn new_for_grpc() -> Self {
199        Self {
200            make_classifier: SharedClassifier::new(GrpcErrorsAsFailures::default()),
201            make_span: DefaultMakeSpan::new(),
202            on_response: DefaultOnResponse::default(),
203            on_request: DefaultOnRequest::default(),
204            on_body_chunk: DefaultOnBodyChunk::default(),
205            on_eos: DefaultOnEos::default(),
206            on_failure: DefaultOnFailure::default(),
207        }
208    }
209}
210
211impl<S, M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure> Layer<S>
212    for TraceLayer<M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>
213where
214    M: Clone,
215    MakeSpan: Clone,
216    OnRequest: Clone,
217    OnResponse: Clone,
218    OnEos: Clone,
219    OnBodyChunk: Clone,
220    OnFailure: Clone,
221{
222    type Service = Trace<S, M, MakeSpan, OnRequest, OnResponse, OnBodyChunk, OnEos, OnFailure>;
223
224    fn layer(&self, inner: S) -> Self::Service {
225        Trace {
226            inner,
227            make_classifier: self.make_classifier.clone(),
228            make_span: self.make_span.clone(),
229            on_request: self.on_request.clone(),
230            on_eos: self.on_eos.clone(),
231            on_body_chunk: self.on_body_chunk.clone(),
232            on_response: self.on_response.clone(),
233            on_failure: self.on_failure.clone(),
234        }
235    }
236}