tower_http/trace/
service.rs

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

tower_http/trace/
service.rs
