mock: add ExpectedId to link span expectations (#3007)

## Motivation

It currently isn't possible to differentiate spans with the same name,
target, and level when setting expectations on `enter`, `exit`, and
`drop_span`. This is not an issue for `tracing-mock`'s original (and
still primary) use case, which is to test `tracing` itself. However,
when testing the tracing instrumentation in library or application code,
this can be a limitation.

For example, when testing the instrumentation in tokio
(tokio-rs/tokio#6112), it isn't possible to set an expectation on which
task span is entered first, because the name, target, and level of those
spans are always identical - in fact, the spans have the same metadata
and only the field values are different.

## Solution

To make differentiating different spans possible, `ExpectId` has been
introduced. It is an opaque struct which represents a `span::Id` and can
be used to match spans from a `new_span` expectation (where a `NewSpan`
is accepted and all fields and values can be expected) through to
subsequent `enter`, `exit`, and `drop_span` expectations.

An `ExpectedId` is passed to an `ExpectedSpan` which then needs to be
expected with `MockCollector::new_span`. A clone of the `ExpectedId` (or
a clone of the `ExpectedSpan` with the `ExpectedId` already on it) will
then match the ID assigned to the span to the other span lifecycle
expectations.

The `ExpectedId` uses an `Arc<AtomicU64>` which has the ID for the new
span assigned to it, and then its clones will be matched against that
same ID.

In future changes it will also be possible to use this `ExpectedId` to
match parent spans, currently a parent is only matched by name.

Author: hds

tracing-mock/src/collector.rs

@@ -159,12 +159,17 @@ use tracing::{
 };
 
 pub(crate) struct SpanState {
+    id: u64,
     name: &'static str,
     refs: usize,
     meta: &'static Metadata<'static>,
 }
 
 impl SpanState {
+    pub(crate) fn id(&self) -> u64 {
+        self.id
+    }
+
     pub(crate) fn metadata(&self) -> &'static Metadata<'static> {
         self.meta
     }
@@ -1100,6 +1105,9 @@ where
         let mut spans = self.spans.lock().unwrap();
         if was_expected {
             if let Expect::NewSpan(mut expected) = expected.pop_front().unwrap() {
+                if let Some(expected_id) = &expected.span.id {
+                    expected_id.set(id.into_u64()).unwrap();
+                }
                 let get_parent_name = || {
                     let stack = self.current.lock().unwrap();
                     span.parent()
@@ -1113,6 +1121,7 @@ where
         spans.insert(
             id.clone(),
             SpanState {
+                id: id.into_u64(),
                 name: meta.name(),
                 refs: 1,
                 meta,

tracing-mock/src/expect.rs

@@ -3,7 +3,7 @@ use std::fmt;
 use crate::{
     event::ExpectedEvent,
     field::{ExpectedField, ExpectedFields, ExpectedValue},
-    span::{ExpectedSpan, NewSpan},
+    span::{ExpectedId, ExpectedSpan, NewSpan},
 };
 
 #[derive(Debug, Eq, PartialEq)]
@@ -51,6 +51,25 @@ pub fn span() -> ExpectedSpan {
     }
 }
 
+/// Returns a new, unset `ExpectedId`.
+///
+/// The `ExpectedId` needs to be attached to a [`NewSpan`] or an
+/// [`ExpectedSpan`] passed to [`MockCollector::new_span`] to
+/// ensure that it gets set. When the a clone of the same
+/// `ExpectedSpan` is attached to an [`ExpectedSpan`] and passed to
+/// any other method on [`MockCollector`] that accepts it, it will
+/// ensure that it is exactly the same span used across those
+/// distinct expectations.
+///
+/// For more details on how to use this struct, see the documentation
+/// on [`ExpectedSpan::with_id`].
+///
+/// [`MockCollector`]: struct@crate::collector::MockCollector
+/// [`MockCollector::new_span`]: fn@crate::collector::MockCollector::new_span
+pub fn id() -> ExpectedId {
+    ExpectedId::new_unset()
+}
+
 impl Expect {
     pub(crate) fn bad(&self, name: impl AsRef<str>, what: fmt::Arguments<'_>) {
         let name = name.as_ref();

tracing-mock/src/span.rs

@@ -94,7 +94,13 @@
 use crate::{
     collector::SpanState, expect, field::ExpectedFields, metadata::ExpectedMetadata, Parent,
 };
-use std::fmt;
+use std::{
+    error, fmt,
+    sync::{
+        atomic::{AtomicU64, Ordering},
+        Arc,
+    },
+};
 
 /// A mock span.
 ///
@@ -104,6 +110,7 @@ use std::fmt;
 /// [`collector`]: mod@crate::collector
 #[derive(Clone, Default, Eq, PartialEq)]
 pub struct ExpectedSpan {
+    pub(crate) id: Option<ExpectedId>,
     pub(crate) metadata: ExpectedMetadata,
 }
 
@@ -137,6 +144,24 @@ where
     expect::span().named(name)
 }
 
+/// A mock span ID.
+///
+/// This ID makes it possible to link together calls to different
+/// [`MockCollector`] span methods that take an [`ExpectedSpan`] in
+/// addition to those that take a [`NewSpan`].
+///
+/// Use [`expect::id`] to construct a new, unset `ExpectedId`.
+///
+/// For more details on how to use this struct, see the documentation
+/// on [`ExpectedSpan::with_id`].
+///
+/// [`expect::id`]: fn@crate::expect::id
+/// [`MockCollector`]: struct@crate::collector::MockCollector
+#[derive(Clone, Default)]
+pub struct ExpectedId {
+    inner: Arc<AtomicU64>,
+}
+
 impl ExpectedSpan {
     /// Sets a name to expect when matching a span.
     ///
@@ -188,6 +213,100 @@ impl ExpectedSpan {
                 name: Some(name.into()),
                 ..self.metadata
             },
+            ..self
+        }
+    }
+
+    /// Sets the `ID` to expect when matching a span.
+    ///
+    /// The [`ExpectedId`] can be used to differentiate spans that are
+    /// otherwise identical. An [`ExpectedId`] needs to be attached to
+    /// an `ExpectedSpan` or [`NewSpan`] which is passed to
+    /// [`MockCollector::new_span`]. The same [`ExpectedId`] can then
+    /// be used to match the exact same span when passed to
+    /// [`MockCollector::enter`], [`MockCollector::exit`], and
+    /// [`MockCollector::drop_span`].
+    ///
+    /// This is especially useful when `tracing-mock` is being used to
+    /// test the traces being generated within your own crate, in which
+    /// case you may need to distinguish between spans which have
+    /// identical metadata but different field values, which can
+    /// otherwise only be checked in [`MockCollector::new_span`].
+    ///
+    /// # Examples
+    ///
+    /// Here we expect that the span that is created first is entered
+    /// second:
+    ///
+    /// ```
+    /// use tracing_mock::{collector, expect};
+    /// let id1 = expect::id();
+    /// let span1 = expect::span().named("span").with_id(id1.clone());
+    /// let id2 = expect::id();
+    /// let span2 = expect::span().named("span").with_id(id2.clone());
+    ///
+    /// let (collector, handle) = collector::mock()
+    ///     .new_span(span1.clone())
+    ///     .new_span(span2.clone())
+    ///     .enter(span2)
+    ///     .enter(span1)
+    ///     .run_with_handle();
+    ///
+    /// tracing::collect::with_default(collector, || {
+    ///     fn create_span() -> tracing::Span {
+    ///         tracing::info_span!("span")
+    ///     }
+    ///
+    ///     let span1 = create_span();
+    ///     let span2 = create_span();
+    ///
+    ///     let _guard2 = span2.enter();
+    ///     let _guard1 = span1.enter();
+    /// });
+    ///
+    /// handle.assert_finished();
+    /// ```
+    ///
+    /// If the order that the spans are entered changes, the test will
+    /// fail:
+    ///
+    /// ```should_panic
+    /// use tracing_mock::{collector, expect};
+    /// let id1 = expect::id();
+    /// let span1 = expect::span().named("span").with_id(id1.clone());
+    /// let id2 = expect::id();
+    /// let span2 = expect::span().named("span").with_id(id2.clone());
+    ///
+    /// let (collector, handle) = collector::mock()
+    ///     .new_span(span1.clone())
+    ///     .new_span(span2.clone())
+    ///     .enter(span2)
+    ///     .enter(span1)
+    ///     .run_with_handle();
+    ///
+    /// tracing::collect::with_default(collector, || {
+    ///     fn create_span() -> tracing::Span {
+    ///         tracing::info_span!("span")
+    ///     }
+    ///
+    ///     let span1 = create_span();
+    ///     let span2 = create_span();
+    ///
+    ///     let _guard1 = span1.enter();
+    ///     let _guard2 = span2.enter();
+    /// });
+    ///
+    /// handle.assert_finished();
+    /// ```
+    ///
+    /// [`MockCollector::new_span`]: fn@crate::collector::MockCollector::new_span
+    /// [`MockCollector::enter`]: fn@crate::collector::MockCollector::enter
+    /// [`MockCollector::exit`]: fn@crate::collector::MockCollector::exit
+    /// [`MockCollector::drop_span`]: fn@crate::collector::MockCollector::drop_span
+    pub fn with_id(self, id: ExpectedId) -> Self {
+        Self {
+            id: Some(id),
+            ..self
         }
     }
 
@@ -241,6 +360,7 @@ impl ExpectedSpan {
                 level: Some(level),
                 ..self.metadata
             },
+            ..self
         }
     }
 
@@ -297,6 +417,7 @@ impl ExpectedSpan {
                 target: Some(target.into()),
                 ..self.metadata
             },
+            ..self
         }
     }
 
@@ -598,6 +719,11 @@ impl ExpectedSpan {
     pub(crate) fn check(&self, actual: &SpanState, collector_name: &str) {
         let meta = actual.metadata();
         let name = meta.name();
+
+        if let Some(expected_id) = &self.id {
+            expected_id.check(actual.id(), format_args!("span `{}`", name), collector_name);
+        }
+
         self.metadata
             .check(meta, format_args!("span `{}`", name), collector_name);
     }
@@ -760,3 +886,69 @@ impl fmt::Debug for NewSpan {
         s.finish()
     }
 }
+
+impl PartialEq for ExpectedId {
+    fn eq(&self, other: &Self) -> bool {
+        self.inner.load(Ordering::Relaxed) == other.inner.load(Ordering::Relaxed)
+    }
+}
+
+impl Eq for ExpectedId {}
+
+impl ExpectedId {
+    const UNSET: u64 = 0;
+
+    pub(crate) fn new_unset() -> Self {
+        Self {
+            inner: Arc::new(AtomicU64::from(Self::UNSET)),
+        }
+    }
+
+    pub(crate) fn set(&self, span_id: u64) -> Result<(), SetActualSpanIdError> {
+        self.inner
+            .compare_exchange(Self::UNSET, span_id, Ordering::Relaxed, Ordering::Relaxed)
+            .map_err(|current| SetActualSpanIdError {
+                previous_span_id: current,
+                new_span_id: span_id,
+            })?;
+        Ok(())
+    }
+
+    pub(crate) fn check(&self, actual: u64, ctx: fmt::Arguments<'_>, collector_name: &str) {
+        let id = self.inner.load(Ordering::Relaxed);
+
+        assert!(
+            id != Self::UNSET,
+            "\n[{}] expected {} to have expected ID set, but it hasn't been, \
+            perhaps this `ExpectedId` wasn't used in a call to `MockCollector::new_span()`?",
+            collector_name,
+            ctx,
+        );
+
+        assert_eq!(
+            id, actual,
+            "\n[{}] expected {} to have ID `{}`, but it has `{}` instead",
+            collector_name, ctx, id, actual,
+        );
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct SetActualSpanIdError {
+    previous_span_id: u64,
+    new_span_id: u64,
+}
+
+impl fmt::Display for SetActualSpanIdError {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        write!(
+            f,
+            "Could not set `ExpecedId` to {new}, \
+            it had already been set to {previous}",
+            new = self.new_span_id,
+            previous = self.previous_span_id
+        )
+    }
+}
+
+impl error::Error for SetActualSpanIdError {}