tracing: add Span::or_current to help with efficient propagation (#…

…1538)

This adds a new `Span::or_current` method that returns the span it's
called on, if that span is enabled, or the current span if that span
is not enabled.

This should provide a more efficient alternative to writing code like
```rust
tokio::spawn(
    future
        .instrument(some_span)
        .in_current_span()
);
```
since it doesn't wrap the future in two different spans which are
(potentially) both entered on polls. It's also a little more concise
to write, which (hopefully) will encourage people to use it instead
of the less efficient alternative.

`Span::or_current` may be useful in other cases as well.

Signed-off-by: Eliza Weisman <eliza@buoyant.io>

examples/examples/tower-load.rs

@@ -41,7 +41,9 @@ use std::{
 };
 use tokio::{time, try_join};
 use tower::{Service, ServiceBuilder, ServiceExt};
-use tracing::{self, debug, error, info, span, trace, warn, Instrument as _, Level, Span};
+use tracing::{
+    self, debug, error, info, info_span, span, trace, warn, Instrument as _, Level, Span,
+};
 use tracing_subscriber::{filter::EnvFilter, reload::Handle};
 use tracing_tower::{request_span, request_span::make};
 
@@ -368,7 +370,7 @@ async fn load_gen(addr: SocketAddr) -> Result<(), Err> {
             .instrument(span)
             .await
         }
-        .instrument(span!(target: "gen", Level::INFO, "generated_request", remote.addr=%addr));
+        .instrument(info_span!(target: "gen", "generated_request", remote.addr=%addr).or_current());
         tokio::spawn(f);
     }
 }

tracing/src/instrument.rs

@@ -34,7 +34,46 @@ pub trait Instrument: Sized {
     /// # }
     /// ```
     ///
-    /// [entered]: ../struct.Span.html#method.enter
+    /// The [`Span::or_current`] combinator can be used in combination with
+    /// `instrument` to ensure that the [current span] is attached to the
+    /// future if the span passed to `instrument` is [disabled]:
+    ///
+    /// ```W
+    /// use tracing::Instrument;
+    /// # mod tokio {
+    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
+    /// # }
+    ///
+    /// let my_future = async {
+    ///     // ...
+    /// };
+    ///
+    /// let outer_span = tracing::info_span!("outer").entered();
+    ///
+    /// // If the "my_future" span is enabled, then the spawned task will
+    /// // be within both "my_future" *and* "outer", since "outer" is
+    /// // "my_future"'s parent. However, if "my_future" is disabled,
+    /// // the spawned task will *not* be in any span.
+    /// tokio::spawn(
+    ///     my_future
+    ///         .instrument(tracing::debug_span!("my_future"))
+    /// );
+    ///
+    /// // Using `Span::or_current` ensures the spawned task is instrumented
+    /// // with the current span, if the new span passed to `instrument` is
+    /// // not enabled. This means that if the "my_future"  span is disabled,
+    /// // the spawned task will still be instrumented with the "outer" span:
+    /// # let my_future = async {};
+    /// tokio::spawn(
+    ///    my_future
+    ///         .instrument(tracing::debug_span!("my_future").or_current())
+    /// );
+    /// ```
+    ///
+    /// [entered]: super::Span::enter()
+    /// [`Span::or_current`]: super::Span::or_current()
+    /// [current span]: super::Span::current()
+    /// [disabled]: super::Span::is_disabled()
     fn instrument(self, span: Span) -> Instrumented<Self> {
         Instrumented { inner: self, span }
     }

tracing/src/span.rs

@@ -900,6 +900,135 @@ impl Span {
         }
     }
 
+    /// Returns this span, if it was [enabled] by the current [`Subscriber`], or
+    /// the [current span] (whose lexical distance may be further than expected),
+    ///  if this span [is disabled].
+    ///
+    /// This method can be useful when propagating spans to spawned threads or
+    /// [async tasks]. Consider the following:
+    ///
+    /// ```
+    /// let _parent_span = tracing::info_span!("parent").entered();
+    ///
+    /// // ...
+    ///
+    /// let child_span = tracing::debug_span!("child");
+    ///
+    /// std::thread::spawn(move || {
+    ///     let _entered = child_span.entered();
+    ///
+    ///     tracing::info!("spawned a thread!");
+    ///
+    ///     // ...
+    /// });
+    /// ```
+    ///
+    /// If the current [`Subscriber`] enables the [`DEBUG`] level, then both
+    /// the "parent" and "child" spans will be enabled. Thus, when the "spawaned
+    /// a thread!" event occurs, it will be inside of the "child" span. Because
+    /// "parent" is the parent of "child", the event will _also_ be inside of
+    /// "parent".
+    ///
+    /// However, if the [`Subscriber`] only enables the [`INFO`] level, the "child"
+    /// span will be disabled. When the thread is spawned, the
+    /// `child_span.entered()` call will do nothing, since "child" is not
+    /// enabled. In this case, the "spawned a thread!" event occurs outside of
+    /// *any* span, since the "child" span was responsible for propagating its
+    /// parent to the spawned thread.
+    ///
+    /// If this is not the desired behavior, `Span::or_current` can be used to
+    /// ensure that the "parent" span is propagated in both cases, either as a
+    /// parent of "child" _or_ directly. For example:
+    ///
+    /// ```
+    /// let _parent_span = tracing::info_span!("parent").entered();
+    ///
+    /// // ...
+    ///
+    /// // If DEBUG is enabled, then "child" will be enabled, and `or_current`
+    /// // returns "child". Otherwise, if DEBUG is not enabled, "child" will be
+    /// // disabled, and `or_current` returns "parent".
+    /// let child_span = tracing::debug_span!("child").or_current();
+    ///
+    /// std::thread::spawn(move || {
+    ///     let _entered = child_span.entered();
+    ///
+    ///     tracing::info!("spawned a thread!");
+    ///
+    ///     // ...
+    /// });
+    /// ```
+    ///
+    /// When spawning [asynchronous tasks][async tasks], `Span::or_current` can
+    /// be used similarly, in combination with [`instrument`]:
+    ///
+    /// ```
+    /// use tracing::Instrument;
+    /// # // lol
+    /// # mod tokio {
+    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
+    /// # }
+    ///
+    /// let _parent_span = tracing::info_span!("parent").entered();
+    ///
+    /// // ...
+    ///
+    /// let child_span = tracing::debug_span!("child");
+    ///
+    /// tokio::spawn(
+    ///     async {
+    ///         tracing::info!("spawned a task!");
+    ///
+    ///         // ...
+    ///
+    ///     }.instrument(child_span.or_current())
+    /// );
+    /// ```
+    ///
+    /// In general, `or_current` should be preferred over nesting an
+    /// [`instrument`]  call inside of an [`in_current_span`] call, as using
+    /// `or_current` will be more efficient.
+    ///
+    /// ```
+    /// use tracing::Instrument;
+    /// # // lol
+    /// # mod tokio {
+    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
+    /// # }
+    /// async fn my_async_fn() {
+    ///     // ...
+    /// }
+    ///
+    /// let _parent_span = tracing::info_span!("parent").entered();
+    ///
+    /// // Do this:
+    /// tokio::spawn(
+    ///     my_async_fn().instrument(tracing::debug_span!("child").or_current())
+    /// );
+    ///
+    /// // ...rather than this:
+    /// tokio::spawn(
+    ///     my_async_fn()
+    ///         .instrument(tracing::debug_span!("child"))
+    ///         .in_current_span()
+    /// );
+    /// ```
+    ///
+    /// [enabled]: crate::Subscriber::enabled
+    /// [`Subscriber`]: crate::Subscriber
+    /// [current span]: Span::current
+    /// [is disabled]: Span::is_disabled
+    /// [`INFO`]: crate::Level::INFO
+    /// [`DEBUG`]: crate::Level::DEBUG
+    /// [async tasks]: std::task
+    /// [`instrument`]: crate::instrument::Instrument
+    pub fn or_current(self) -> Self {
+        if self.is_disabled() {
+            return Self::current();
+        }
+        self
+    }
+
     #[inline]
     fn do_enter(&self) {
         if let Some(inner) = self.inner.as_ref() {