tracing_actix_web/
middleware.rs1use std::{
2 future::{ready, Future, Ready},
3 pin::Pin,
4 task::{Context, Poll},
5};
6
7use actix_web::{
8 body::{BodySize, MessageBody},
9 dev::{Service, ServiceRequest, ServiceResponse, Transform},
10 web::Bytes,
11 Error, HttpMessage,
12};
13use tracing::Span;
14
15use crate::{DefaultRootSpanBuilder, RequestId, RootSpan, RootSpanBuilder};
16
17pub struct TracingLogger<RootSpan: RootSpanBuilder> {
88 root_span_builder: std::marker::PhantomData<RootSpan>,
89}
90
91impl<RootSpan: RootSpanBuilder> Clone for TracingLogger<RootSpan> {
92 fn clone(&self) -> Self {
93 Self::new()
94 }
95}
96
97impl Default for TracingLogger<DefaultRootSpanBuilder> {
98 fn default() -> Self {
99 TracingLogger::new()
100 }
101}
102
103impl<RootSpan: RootSpanBuilder> TracingLogger<RootSpan> {
104 pub fn new() -> TracingLogger<RootSpan> {
105 TracingLogger {
106 root_span_builder: Default::default(),
107 }
108 }
109}
110
111impl<S, B, RootSpan> Transform<S, ServiceRequest> for TracingLogger<RootSpan>
112where
113 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
114 S::Future: 'static,
115 B: MessageBody + 'static,
116 RootSpan: RootSpanBuilder,
117{
118 type Response = ServiceResponse<StreamSpan<B>>;
119 type Error = Error;
120 type Transform = TracingLoggerMiddleware<S, RootSpan>;
121 type InitError = ();
122 type Future = Ready<Result<Self::Transform, Self::InitError>>;
123
124 fn new_transform(&self, service: S) -> Self::Future {
125 ready(Ok(TracingLoggerMiddleware {
126 service,
127 root_span_builder: std::marker::PhantomData,
128 }))
129 }
130}
131
132#[doc(hidden)]
133pub struct TracingLoggerMiddleware<S, RootSpanBuilder> {
134 service: S,
135 root_span_builder: std::marker::PhantomData<RootSpanBuilder>,
136}
137
138#[allow(clippy::type_complexity)]
139impl<S, B, RootSpanType> Service<ServiceRequest> for TracingLoggerMiddleware<S, RootSpanType>
140where
141 S: Service<ServiceRequest, Response = ServiceResponse<B>, Error = Error>,
142 S::Future: 'static,
143 B: MessageBody + 'static,
144 RootSpanType: RootSpanBuilder,
145{
146 type Response = ServiceResponse<StreamSpan<B>>;
147 type Error = Error;
148 type Future = TracingResponse<S::Future, RootSpanType>;
149
150 actix_web::dev::forward_ready!(service);
151
152 fn call(&self, req: ServiceRequest) -> Self::Future {
153 req.extensions_mut().insert(RequestId::generate());
154 let root_span = RootSpanType::on_request_start(&req);
155
156 let root_span_wrapper = RootSpan::new(root_span.clone());
157 req.extensions_mut().insert(root_span_wrapper);
158
159 let fut = root_span.in_scope(|| self.service.call(req));
160
161 TracingResponse {
162 fut,
163 span: root_span,
164 _root_span_type: std::marker::PhantomData,
165 }
166 }
167}
168
169#[doc(hidden)]
170#[pin_project::pin_project]
171pub struct TracingResponse<F, RootSpanType> {
172 #[pin]
173 fut: F,
174 span: Span,
175 _root_span_type: std::marker::PhantomData<RootSpanType>,
176}
177
178#[doc(hidden)]
179#[pin_project::pin_project]
180pub struct StreamSpan<B> {
181 #[pin]
182 body: B,
183 span: Span,
184}
185
186impl<F, B, RootSpanType> Future for TracingResponse<F, RootSpanType>
187where
188 F: Future<Output = Result<ServiceResponse<B>, Error>>,
189 B: MessageBody + 'static,
190 RootSpanType: RootSpanBuilder,
191{
192 type Output = Result<ServiceResponse<StreamSpan<B>>, Error>;
193
194 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
195 let this = self.project();
196
197 let fut = this.fut;
198 let span = this.span;
199
200 span.in_scope(|| match fut.poll(cx) {
201 Poll::Pending => Poll::Pending,
202 Poll::Ready(outcome) => {
203 RootSpanType::on_request_end(Span::current(), &outcome);
204
205 #[cfg(feature = "emit_event_on_error")]
206 {
207 emit_event_on_error(&outcome);
208 }
209
210 Poll::Ready(outcome.map(|service_response| {
211 service_response.map_body(|_, body| StreamSpan {
212 body,
213 span: span.clone(),
214 })
215 }))
216 }
217 })
218 }
219}
220
221impl<B> MessageBody for StreamSpan<B>
222where
223 B: MessageBody,
224{
225 type Error = B::Error;
226
227 fn size(&self) -> BodySize {
228 self.body.size()
229 }
230
231 fn poll_next(
232 self: Pin<&mut Self>,
233 cx: &mut Context<'_>,
234 ) -> Poll<Option<Result<Bytes, Self::Error>>> {
235 let this = self.project();
236
237 let body = this.body;
238 let span = this.span;
239 span.in_scope(|| body.poll_next(cx))
240 }
241}
242
243#[cfg(feature = "emit_event_on_error")]
244fn emit_event_on_error<B: 'static>(outcome: &Result<ServiceResponse<B>, actix_web::Error>) {
245 match outcome {
246 Ok(response) => {
247 if let Some(err) = response.response().error() {
248 emit_error_event(err.as_response_error(), response.status())
250 }
251 }
252 Err(error) => {
253 let response_error = error.as_response_error();
254 emit_error_event(response_error, response_error.status_code())
255 }
256 }
257}
258
259#[cfg(feature = "emit_event_on_error")]
260fn emit_error_event(
261 response_error: &dyn actix_web::ResponseError,
262 status_code: actix_web::http::StatusCode,
263) {
264 let error_msg_prefix = "Error encountered while processing the incoming HTTP request";
265 if status_code.is_client_error() {
266 tracing::warn!("{}: {:?}", error_msg_prefix, response_error);
267 } else {
268 tracing::error!("{}: {:?}", error_msg_prefix, response_error);
269 }
270}