async_lsp/
tracing.rs

1//! Attach [`tracing::Span`]s over underlying handlers.
2//!
3//! *Applies to both Language Servers and Language Clients.*
4//!
5//! This middleware attaches spans to logs in underlying implementations, with optional method
6//! strings of current processing requests/notifications.
7//! All of these methods are instrumented by the [`Default`] configuration:
8//! - [`Service::poll_ready`].
9//! - [`Future::poll`] of returned `Future` from [`Service::call`].
10//! - [`LspService::notify`].
11//! - [`LspService::emit`].
12use std::future::Future;
13use std::ops::ControlFlow;
14use std::pin::Pin;
15use std::task::{Context, Poll};
16
17use pin_project_lite::pin_project;
18use tower_layer::Layer;
19use tower_service::Service;
20use tracing::{info_span, Span};
21
22use crate::{AnyEvent, AnyNotification, AnyRequest, LspService, Result};
23
24/// The middleware attaching [`tracing::Span`]s over underlying handlers.
25///
26/// See [module level documentations](self) for details.
27#[derive(Default)]
28pub struct Tracing<S> {
29    service: S,
30    spans: TracingBuilder,
31}
32
33define_getters!(impl[S] Tracing<S>, service: S);
34
35impl<S: LspService> Service<AnyRequest> for Tracing<S> {
36    type Response = S::Response;
37    type Error = S::Error;
38    type Future = ResponseFuture<S::Future>;
39
40    fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
41        let _guard = self.spans.service_ready.map(|f| f().entered());
42        self.service.poll_ready(cx)
43    }
44
45    fn call(&mut self, req: AnyRequest) -> Self::Future {
46        ResponseFuture {
47            span: self.spans.request.map(|f| f(&req)),
48            fut: self.service.call(req),
49        }
50    }
51}
52
53pin_project! {
54    /// The [`Future`] type used by the [`Tracing`] middleware.
55    pub struct ResponseFuture<Fut> {
56        span: Option<Span>,
57        #[pin]
58        fut: Fut,
59    }
60}
61
62impl<Fut: Future> Future for ResponseFuture<Fut> {
63    type Output = Fut::Output;
64
65    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
66        let this = self.project();
67        let _guard = this.span.as_mut().map(|span| span.enter());
68        this.fut.poll(cx)
69    }
70}
71
72impl<S: LspService> LspService for Tracing<S> {
73    fn notify(&mut self, notif: AnyNotification) -> ControlFlow<Result<()>> {
74        let _guard = self.spans.notification.map(|f| f(&notif).entered());
75        self.service.notify(notif)
76    }
77
78    fn emit(&mut self, event: AnyEvent) -> ControlFlow<Result<()>> {
79        let _guard = self.spans.event.map(|f| f(&event).entered());
80        self.service.emit(event)
81    }
82}
83
84/// The builder of [`Tracing`] middleware.
85///
86/// See [module level documentations](self) for details.
87#[derive(Clone)]
88#[must_use]
89pub struct TracingBuilder {
90    service_ready: Option<fn() -> Span>,
91    request: Option<fn(&AnyRequest) -> Span>,
92    notification: Option<fn(&AnyNotification) -> Span>,
93    event: Option<fn(&AnyEvent) -> Span>,
94}
95
96impl Default for TracingBuilder {
97    fn default() -> Self {
98        Self {
99            service_ready: Some(|| info_span!("service_ready")),
100            request: Some(|req| info_span!("request", method = req.method)),
101            notification: Some(|notif| info_span!("notification", method = notif.method)),
102            event: Some(|event| info_span!("event", type_name = event.type_name())),
103        }
104    }
105}
106
107impl TracingBuilder {
108    /// Creating the builder with no spans configured.
109    ///
110    /// NB. This is **NOT** the same as [`TracingBuilder::default`] which configures ALL default
111    /// spans.
112    pub fn new() -> Self {
113        Self {
114            service_ready: None,
115            request: None,
116            notification: None,
117            event: None,
118        }
119    }
120
121    /// Set a [`tracing::Span`] builder to instrument [`Service::poll_ready`] method.
122    pub fn service_ready(mut self, f: fn() -> Span) -> Self {
123        self.service_ready = Some(f);
124        self
125    }
126
127    /// Set a [`tracing::Span`] builder to instrument [`Future::poll`] of the `Future` returned by
128    /// [`Service::call`].
129    pub fn request(mut self, f: fn(&AnyRequest) -> Span) -> Self {
130        self.request = Some(f);
131        self
132    }
133
134    /// Set a [`tracing::Span`] builder to instrument [`LspService::notify`].
135    pub fn notification(mut self, f: fn(&AnyNotification) -> Span) -> Self {
136        self.notification = Some(f);
137        self
138    }
139
140    /// Set a [`tracing::Span`] builder to instrument [`LspService::emit`].
141    pub fn event(mut self, f: fn(&AnyEvent) -> Span) -> Self {
142        self.event = Some(f);
143        self
144    }
145
146    /// Build the middleware with the current configuration.
147    pub fn build<S>(&self, service: S) -> Tracing<S> {
148        Tracing {
149            service,
150            spans: self.clone(),
151        }
152    }
153}
154
155/// A type alias of [`TracingLayer`] conforming to the naming convention of [`tower_layer`].
156pub type TracingLayer = TracingBuilder;
157
158impl<S> Layer<S> for TracingBuilder {
159    type Service = Tracing<S>;
160
161    fn layer(&self, inner: S) -> Self::Service {
162        self.build(inner)
163    }
164}
async_lsp/tracing.rs

async_lsp/
tracing.rs
