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#[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    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    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    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    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    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    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    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    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    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}