apollo_router/layers/
instrument.rs

1//! Instrumentation layer that allows services to be wrapped in a span.
2//!
3//! Using ServiceBuilderExt:
4//! ```rust
5//! # use tower::ServiceBuilder;
6//! # use tower_service::Service;
7//! # use tracing::info_span;
8//! # use apollo_router::layers::ServiceBuilderExt;
9//! # fn test<T>(service: impl Service<T>) {
10//! let instrumented = ServiceBuilder::new()
11//!             .instrument(|_request| info_span!("query_planning"))
12//!             .service(service);
13//! # }
14//! ```
15//! Now calls to the wrapped service will be wrapped in a span. You can attach attributes to the span from the request.
16//!
17//! See [`Layer`] and [`Service`] for more details.
18
19use std::marker::PhantomData;
20use std::task::Context;
21use std::task::Poll;
22
23use tower::Layer;
24use tower_service::Service;
25use tracing::Instrument;
26
27/// [`Layer`] for instrumentation. See [`ServiceBuilderExt::instrument()`](crate::layers::ServiceBuilderExt::instrument()).
28pub struct InstrumentLayer<F, Request>
29where
30    F: Fn(&Request) -> tracing::Span,
31{
32    span_fn: F,
33    phantom: PhantomData<Request>,
34}
35
36impl<F, Request> InstrumentLayer<F, Request>
37where
38    F: Fn(&Request) -> tracing::Span,
39{
40    #[allow(missing_docs)] // FIXME
41    pub fn new(span_fn: F) -> InstrumentLayer<F, Request> {
42        Self {
43            span_fn,
44            phantom: Default::default(),
45        }
46    }
47}
48
49impl<F, S, Request> Layer<S> for InstrumentLayer<F, Request>
50where
51    S: Service<Request>,
52    F: Fn(&Request) -> tracing::Span + Clone,
53{
54    type Service = InstrumentService<F, S, Request>;
55
56    fn layer(&self, inner: S) -> Self::Service {
57        InstrumentService {
58            inner,
59            span_fn: self.span_fn.clone(),
60            phantom: Default::default(),
61        }
62    }
63}
64
65/// [`Service`] for instrumentation. See [`ServiceBuilderExt::instrument()`](crate::layers::ServiceBuilderExt::instrument()).
66pub struct InstrumentService<F, S, Request>
67where
68    S: Service<Request>,
69    F: Fn(&Request) -> tracing::Span,
70{
71    inner: S,
72    span_fn: F,
73    phantom: PhantomData<Request>,
74}
75
76impl<F, S, Request> Service<Request> for InstrumentService<F, S, Request>
77where
78    F: Fn(&Request) -> tracing::Span,
79    S: Service<Request>,
80    <S as Service<Request>>::Future: Send + 'static,
81{
82    type Response = S::Response;
83    type Error = S::Error;
84    type Future = tracing::instrument::Instrumented<S::Future>;
85
86    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
87        self.inner.poll_ready(cx)
88    }
89
90    fn call(&mut self, req: Request) -> Self::Future {
91        let span = (self.span_fn)(&req);
92
93        // since other spans may be created while generating the future, enter this span right now so it will appear as parent
94        let _guard = span.enter();
95        let res = self.inner.call(req);
96        drop(_guard);
97
98        res.instrument(span)
99    }
100}
apollo_router/layers/instrument.rs

apollo_router/layers/
instrument.rs
