diff options
Diffstat (limited to '')
| -rw-r--r-- | third_party/rust/tracing/src/instrument.rs | 194 |
1 files changed, 194 insertions, 0 deletions
diff --git a/third_party/rust/tracing/src/instrument.rs b/third_party/rust/tracing/src/instrument.rs new file mode 100644 index 0000000000..d1070ba8e2 --- /dev/null +++ b/third_party/rust/tracing/src/instrument.rs @@ -0,0 +1,194 @@ +use crate::stdlib::pin::Pin; +use crate::stdlib::task::{Context, Poll}; +use crate::stdlib::{future::Future, marker::Sized}; +use crate::{dispatcher, span::Span, Dispatch}; +use pin_project_lite::pin_project; + +/// Attaches spans to a `std::future::Future`. +/// +/// Extension trait allowing futures to be +/// instrumented with a `tracing` [span]. +/// +/// [span]: ../struct.Span.html +pub trait Instrument: Sized { + /// Instruments this type with the provided `Span`, returning an + /// `Instrumented` wrapper. + /// + /// The attached `Span` will be [entered] every time the instrumented `Future` is polled. + /// + /// # Examples + /// + /// Instrumenting a future: + /// + /// ```rust + /// use tracing::Instrument; + /// + /// # async fn doc() { + /// let my_future = async { + /// // ... + /// }; + /// + /// my_future + /// .instrument(tracing::info_span!("my_future")) + /// .await + /// # } + /// ``` + /// + /// [entered]: ../struct.Span.html#method.enter + fn instrument(self, span: Span) -> Instrumented<Self> { + Instrumented { inner: self, span } + } + + /// Instruments this type with the [current] `Span`, returning an + /// `Instrumented` wrapper. + /// + /// If the instrumented type is a future, stream, or sink, the attached `Span` + /// will be [entered] every time it is polled. If the instrumented type + /// is a future executor, every future spawned on that executor will be + /// instrumented by the attached `Span`. + /// + /// This can be used to propagate the current span when spawning a new future. + /// + /// # Examples + /// + /// ```rust + /// use tracing::Instrument; + /// + /// # async fn doc() { + /// let span = tracing::info_span!("my_span"); + /// let _enter = span.enter(); + /// + /// // ... + /// + /// let future = async { + /// tracing::debug!("this event will occur inside `my_span`"); + /// // ... + /// }; + /// tokio::spawn(future.in_current_span()); + /// # } + /// ``` + /// + /// [current]: ../struct.Span.html#method.current + /// [entered]: ../struct.Span.html#method.enter + #[inline] + fn in_current_span(self) -> Instrumented<Self> { + self.instrument(Span::current()) + } +} + +/// Extension trait allowing futures to be instrumented with +/// a `tracing` [`Subscriber`]. +/// +/// [`Subscriber`]: ../trait.Subscriber.html +pub trait WithSubscriber: Sized { + /// Attaches the provided [`Subscriber`] to this type, returning a + /// `WithDispatch` wrapper. + /// + /// The attached subscriber will be set as the [default] when the returned `Future` is polled. + /// + /// [`Subscriber`]: ../trait.Subscriber.html + /// [default]: https://docs.rs/tracing/latest/tracing/dispatcher/index.html#setting-the-default-subscriber + fn with_subscriber<S>(self, subscriber: S) -> WithDispatch<Self> + where + S: Into<Dispatch>, + { + WithDispatch { + inner: self, + dispatch: subscriber.into(), + } + } + + /// Attaches the current [default] [`Subscriber`] to this type, returning a + /// `WithDispatch` wrapper. + /// + /// When the wrapped type is a future, stream, or sink, the attached + /// subscriber will be set as the [default] while it is being polled. + /// When the wrapped type is an executor, the subscriber will be set as the + /// default for any futures spawned on that executor. + /// + /// This can be used to propagate the current dispatcher context when + /// spawning a new future. + /// + /// [`Subscriber`]: ../trait.Subscriber.html + /// [default]: https://docs.rs/tracing/latest/tracing/dispatcher/index.html#setting-the-default-subscriber + #[inline] + fn with_current_subscriber(self) -> WithDispatch<Self> { + WithDispatch { + inner: self, + dispatch: dispatcher::get_default(|default| default.clone()), + } + } +} + +pin_project! { + /// A future that has been instrumented with a `tracing` subscriber. + #[derive(Clone, Debug)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct WithDispatch<T> { + #[pin] + inner: T, + dispatch: Dispatch, + } +} + +pin_project! { + /// A future that has been instrumented with a `tracing` span. + #[derive(Debug, Clone)] + #[must_use = "futures do nothing unless you `.await` or poll them"] + pub struct Instrumented<T> { + #[pin] + inner: T, + span: Span, + } +} + +impl<T: Future> Future for Instrumented<T> { + type Output = T::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { + let this = self.project(); + let _enter = this.span.enter(); + this.inner.poll(cx) + } +} + +impl<T: Sized> Instrument for T {} + +impl<T> Instrumented<T> { + /// Borrows the `Span` that this type is instrumented by. + pub fn span(&self) -> &Span { + &self.span + } + + /// Mutably borrows the `Span` that this type is instrumented by. + pub fn span_mut(&mut self) -> &mut Span { + &mut self.span + } + + /// Borrows the wrapped type. + pub fn inner(&self) -> &T { + &self.inner + } + + /// Mutably borrows the wrapped type. + pub fn inner_mut(&mut self) -> &mut T { + &mut self.inner + } + + /// Get a pinned reference to the wrapped type. + pub fn inner_pin_ref(self: Pin<&Self>) -> Pin<&T> { + self.project_ref().inner + } + + /// Get a pinned mutable reference to the wrapped type. + pub fn inner_pin_mut(self: Pin<&mut Self>) -> Pin<&mut T> { + self.project().inner + } + + /// Consumes the `Instrumented`, returning the wrapped type. + /// + /// Note that this drops the span. + pub fn into_inner(self) -> T { + self.inner + } +} |