Line data Source code
1 : //! Crutch module to work around tracing infrastructure deficiencies
2 : //!
3 : //! We wish to collect granular request spans without impacting performance
4 : //! by much. Ideally, we should have zero overhead for a sampling rate of 0.
5 : //!
6 : //! The approach taken by the pageserver crate is to use a completely different
7 : //! span hierarchy for the performance spans. Spans are explicitly stored in
8 : //! the request context and use a different [`tracing::Subscriber`] in order
9 : //! to avoid expensive filtering.
10 : //!
11 : //! [`tracing::Span`] instances record their [`tracing::Dispatch`] and, implcitly,
12 : //! their [`tracing::Subscriber`] at creation time. However, upon exiting the span,
13 : //! the global default [`tracing::Dispatch`] is used. This is problematic if one
14 : //! wishes to juggle different subscribers.
15 : //!
16 : //! In order to work around this, this module provides a [`PerfSpan`] type which
17 : //! wraps a [`Span`] and sets the default subscriber when exiting the span. This
18 : //! achieves the correct routing.
19 : //!
20 : //! There's also a modified version of [`tracing::Instrument`] which works with
21 : //! [`PerfSpan`].
22 :
23 : use core::{
24 : future::Future,
25 : marker::Sized,
26 : mem::ManuallyDrop,
27 : pin::Pin,
28 : task::{Context, Poll},
29 : };
30 : use pin_project_lite::pin_project;
31 : use tracing::{Dispatch, span::Span};
32 :
33 : #[derive(Debug, Clone)]
34 : pub struct PerfSpan {
35 : inner: ManuallyDrop<Span>,
36 : dispatch: Dispatch,
37 : }
38 :
39 : #[must_use = "once a span has been entered, it should be exited"]
40 : pub struct PerfSpanEntered<'a> {
41 : span: &'a PerfSpan,
42 : }
43 :
44 : impl PerfSpan {
45 0 : pub fn new(span: Span, dispatch: Dispatch) -> Self {
46 0 : Self {
47 0 : inner: ManuallyDrop::new(span),
48 0 : dispatch,
49 0 : }
50 0 : }
51 :
52 0 : pub fn enter(&self) -> PerfSpanEntered {
53 0 : if let Some(ref id) = self.inner.id() {
54 0 : self.dispatch.enter(id);
55 0 : }
56 :
57 0 : PerfSpanEntered { span: self }
58 0 : }
59 :
60 0 : pub fn inner(&self) -> &Span {
61 0 : &self.inner
62 0 : }
63 : }
64 :
65 : impl Drop for PerfSpan {
66 0 : fn drop(&mut self) {
67 0 : // Bring the desired dispatch into scope before explicitly calling
68 0 : // the span destructor. This routes the span exit to the correct
69 0 : // [`tracing::Subscriber`].
70 0 : let _dispatch_guard = tracing::dispatcher::set_default(&self.dispatch);
71 0 : // SAFETY: ManuallyDrop in Drop implementation
72 0 : unsafe { ManuallyDrop::drop(&mut self.inner) }
73 0 : }
74 : }
75 :
76 : impl Drop for PerfSpanEntered<'_> {
77 0 : fn drop(&mut self) {
78 0 : assert!(self.span.inner.id().is_some());
79 :
80 0 : let _dispatch_guard = tracing::dispatcher::set_default(&self.span.dispatch);
81 0 : self.span.dispatch.exit(&self.span.inner.id().unwrap());
82 0 : }
83 : }
84 :
85 : pub trait PerfInstrument: Sized {
86 0 : fn instrument(self, span: PerfSpan) -> PerfInstrumented<Self> {
87 0 : PerfInstrumented {
88 0 : inner: ManuallyDrop::new(self),
89 0 : span,
90 0 : }
91 0 : }
92 : }
93 :
94 : pin_project! {
95 : #[project = PerfInstrumentedProj]
96 : #[derive(Debug, Clone)]
97 : #[must_use = "futures do nothing unless you `.await` or poll them"]
98 : pub struct PerfInstrumented<T> {
99 : // `ManuallyDrop` is used here to to enter instrument `Drop` by entering
100 : // `Span` and executing `ManuallyDrop::drop`.
101 : #[pin]
102 : inner: ManuallyDrop<T>,
103 : span: PerfSpan,
104 : }
105 :
106 : impl<T> PinnedDrop for PerfInstrumented<T> {
107 : fn drop(this: Pin<&mut Self>) {
108 : let this = this.project();
109 : let _enter = this.span.enter();
110 : // SAFETY: 1. `Pin::get_unchecked_mut()` is safe, because this isn't
111 : // different from wrapping `T` in `Option` and calling
112 : // `Pin::set(&mut this.inner, None)`, except avoiding
113 : // additional memory overhead.
114 : // 2. `ManuallyDrop::drop()` is safe, because
115 : // `PinnedDrop::drop()` is guaranteed to be called only
116 : // once.
117 : unsafe { ManuallyDrop::drop(this.inner.get_unchecked_mut()) }
118 : }
119 : }
120 : }
121 :
122 : impl<'a, T> PerfInstrumentedProj<'a, T> {
123 : /// Get a mutable reference to the [`Span`] a pinned mutable reference to
124 : /// the wrapped type.
125 0 : fn span_and_inner_pin_mut(self) -> (&'a mut PerfSpan, Pin<&'a mut T>) {
126 0 : // SAFETY: As long as `ManuallyDrop<T>` does not move, `T` won't move
127 0 : // and `inner` is valid, because `ManuallyDrop::drop` is called
128 0 : // only inside `Drop` of the `Instrumented`.
129 0 : let inner = unsafe { self.inner.map_unchecked_mut(|v| &mut **v) };
130 0 : (self.span, inner)
131 0 : }
132 : }
133 :
134 : impl<T: Future> Future for PerfInstrumented<T> {
135 : type Output = T::Output;
136 :
137 0 : fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
138 0 : let (span, inner) = self.project().span_and_inner_pin_mut();
139 0 : let _enter = span.enter();
140 0 : inner.poll(cx)
141 0 : }
142 : }
143 :
144 : impl<T: Sized> PerfInstrument for T {}
|
