1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94

//! Instrumentation layer that allows services to be wrapped in a span.
//!
//! Using ServiceBuilderExt:
//! ```rust
//! # use tower::ServiceBuilder;
//! # use tower_service::Service;
//! # use tracing::info_span;
//! # use apollo_router::layers::ServiceBuilderExt;
//! # fn test<T>(service: impl Service<T>) {
//! let instrumented = ServiceBuilder::new()
//!             .instrument(|_request| info_span!("query_planning"))
//!             .service(service);
//! # }
//! ```
//! Now calls to the wrapped service will be wrapped in a span. You can attach attributes to the span from the request.
//!
//! See [`Layer`] and [`Service`] for more details.

use std::marker::PhantomData;
use std::task::Context;
use std::task::Poll;

use tower::Layer;
use tower_service::Service;
use tracing::Instrument;

/// [`Layer`] for instrumentation. See [`ServiceBuilderExt::instrument()`](crate::layers::ServiceBuilderExt::instrument()).
pub struct InstrumentLayer<F, Request>
where
    F: Fn(&Request) -> tracing::Span,
{
    span_fn: F,
    phantom: PhantomData<Request>,
}

impl<F, Request> InstrumentLayer<F, Request>
where
    F: Fn(&Request) -> tracing::Span,
{
    #[allow(missing_docs)] // FIXME
    pub fn new(span_fn: F) -> InstrumentLayer<F, Request> {
        Self {
            span_fn,
            phantom: Default::default(),
        }
    }
}

impl<F, S, Request> Layer<S> for InstrumentLayer<F, Request>
where
    S: Service<Request>,
    F: Fn(&Request) -> tracing::Span + Clone,
{
    type Service = InstrumentService<F, S, Request>;

    fn layer(&self, inner: S) -> Self::Service {
        InstrumentService {
            inner,
            span_fn: self.span_fn.clone(),
            phantom: Default::default(),
        }
    }
}

/// [`Service`] for instrumentation. See [`ServiceBuilderExt::instrument()`](crate::layers::ServiceBuilderExt::instrument()).
pub struct InstrumentService<F, S, Request>
where
    S: Service<Request>,
    F: Fn(&Request) -> tracing::Span,
{
    inner: S,
    span_fn: F,
    phantom: PhantomData<Request>,
}

impl<F, S, Request> Service<Request> for InstrumentService<F, S, Request>
where
    F: Fn(&Request) -> tracing::Span,
    S: Service<Request>,
    <S as Service<Request>>::Future: Send + 'static,
{
    type Response = S::Response;
    type Error = S::Error;
    type Future = tracing::instrument::Instrumented<S::Future>;

    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
        self.inner.poll_ready(cx)
    }

    fn call(&mut self, req: Request) -> Self::Future {
        let span = (self.span_fn)(&req);
        self.inner.call(req).instrument(span)
    }
}