mock: add ExpectedId to link span expectations

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::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
+    /// [`MockColletor::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 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::MockCollector::new_span
+    /// [`MockCollector::enter`]: fn@crate::MockCollector::enter
+    /// [`MockCollector::exit`]: fn@crate::MockCollector::exit
+    /// [`MockCollector::drop_span`]: fn@crate::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,72 @@ impl fmt::Debug for NewSpan {
         s.finish()
     }
 }
+
+impl PartialEq for ExpectedId {
+    fn eq(&self, other: &Self) -> bool {
+        let own_id = self.inner.load(Ordering::Relaxed);
+        let other_id = other.inner.load(Ordering::Relaxed);
+
+        own_id != Self::UNSET && own_id == other_id
+    }
+}
+
+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 {}