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}