msg_common/

span.rs

use std::{
    pin::Pin,
    task::{Context, Poll},
};

use derive_more::{Deref, DerefMut};

/// A container with a [`tracing::Span`] attached.
///
/// If `T` is a `Future`, the span will be entered when polling the future, making this struct
/// almost equivalent to [`tracing::Instrument`]. The main difference is that the return value will
/// contain the span as well so it can be re-entered later.
#[derive(Debug, Clone, Deref, DerefMut)]
pub struct WithSpan<T> {
    #[deref]
    #[deref_mut]
    pub inner: T,
    pub span: tracing::Span,
}

impl<T> WithSpan<T> {
    /// Create a spanned container using [`tracing::Span::current`] span.
    #[inline]
    pub fn current(inner: T) -> Self {
        Self { inner, span: tracing::Span::current() }
    }

    /// Create a container with a no-op [`tracing::Span::none`] span, to be eventually replaced
    /// with [`Self::with_span`]
    #[inline]
    pub fn new(inner: T) -> Self {
        Self { inner, span: tracing::Span::none() }
    }

    /// Replace the current [`tracing::Span`] with the provided one.
    #[inline]
    pub fn with_span(mut self, span: tracing::Span) -> Self {
        self.span = span;
        self
    }

    /// Break the spanned container into a tuple containing the inner object and the span.
    #[inline]
    pub fn into_parts(self) -> (T, tracing::Span) {
        (self.inner, self.span)
    }

    #[inline]
    pub fn into_inner(self) -> T {
        self.inner
    }
}

/// Projection type for [`WithSpan`], as generated by a library like `pin_project`.
struct WithSpanProjection<'a, T> {
    inner: Pin<&'a mut T>,
    // For the purpose of `poll`, we don't need a mutable reference.
    span: &'a tracing::Span,
}

/// The struct is [`Unpin`] only if its pinned fields are [`Unpin`].
impl<T: Unpin> Unpin for WithSpan<T> {}

impl<T> WithSpan<T> {
    /// Project the fields of the struct.
    #[inline]
    fn project(self: Pin<&mut Self>) -> WithSpanProjection<'_, T> {
        unsafe {
            // SAFETY: span is [`Unpin`], and self.inner is never moved while pinned.
            let this = self.get_unchecked_mut();
            WithSpanProjection { inner: Pin::new_unchecked(&mut this.inner), span: &mut this.span }
        }
    }
}

impl<T: Future> Future for WithSpan<T> {
    type Output = WithSpan<T::Output>;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let WithSpanProjection { inner, span } = self.project();

        let _g = span.enter();

        if let Poll::Ready(val) = inner.poll(cx) {
            return Poll::Ready(val.with_span(span));
        }

        Poll::Pending
    }
}

/// A container with a [`tracing::span::EnteredSpan`] attached.
#[derive(Debug, Deref, DerefMut)]
pub struct WithEntered<T> {
    #[deref]
    #[deref_mut]
    pub inner: T,
    pub span: tracing::span::EnteredSpan,
}

/// Trait to convert [`WithSpan<T>`] containers into [`WithEntered<T>`] containers, by
/// [`tracing::Span::enter`]ing the span.
///
/// The associated type [`EnterSpan::Output`] allows for conversions that change the outer type
/// which may wrap [`WithSpan<T>`], such as `Option` or `Poll`. For example, one implementation for
/// poll could allow to convert from `Poll<WithSpan<T>>` to `Poll<WithEntered<T>>`.
pub trait EnterSpan {
    type Output;

    fn enter(self) -> Self::Output;
}

impl<T> EnterSpan for WithSpan<T> {
    type Output = WithEntered<T>;

    #[inline]
    fn enter(self) -> Self::Output {
        let WithSpan { inner, span } = self;
        WithEntered { inner, span: span.entered() }
    }
}

impl<T> EnterSpan for WithEntered<T> {
    type Output = WithEntered<T>;

    #[inline]
    fn enter(self) -> Self::Output {
        self
    }
}

impl<T: EnterSpan> EnterSpan for Option<T> {
    type Output = Option<<T as EnterSpan>::Output>;

    #[inline]
    fn enter(self) -> Self::Output {
        self.map(|v| v.enter())
    }
}

impl<T: EnterSpan, E> EnterSpan for Result<T, E> {
    type Output = Result<<T as EnterSpan>::Output, E>;

    #[inline]
    fn enter(self) -> Self::Output {
        self.map(|v| v.enter())
    }
}

impl<T: EnterSpan> EnterSpan for Poll<T> {
    type Output = Poll<<T as EnterSpan>::Output>;

    #[inline]
    fn enter(self) -> Self::Output {
        match self {
            Poll::Ready(v) => Poll::Ready(v.enter()),
            Poll::Pending => Poll::Pending,
        }
    }
}

pub trait IntoSpanExt {
    fn into_span(self) -> tracing::Span;
}

impl IntoSpanExt for tracing::Span {
    #[inline]
    fn into_span(self) -> tracing::Span {
        self
    }
}

impl IntoSpanExt for &tracing::Span {
    #[inline]
    fn into_span(self) -> tracing::Span {
        self.clone()
    }
}

impl IntoSpanExt for tracing::span::EnteredSpan {
    #[inline]
    fn into_span(self) -> tracing::Span {
        self.clone()
    }
}

pub trait SpanExt<T> {
    fn with_span(self, span: impl IntoSpanExt) -> WithSpan<T>;

    fn with_current_span(self) -> WithSpan<T>;
}

impl<T> SpanExt<T> for T {
    #[inline]
    fn with_span(self, span: impl IntoSpanExt) -> WithSpan<T> {
        WithSpan { inner: self, span: span.into_span() }
    }

    fn with_current_span(self) -> WithSpan<T> {
        WithSpan::current(self)
    }
}