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

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.

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.

tracing-attributes/tests/parents.rs

@@ -21,23 +21,16 @@ fn default_parent_test() {
         .new_span(
             contextual_parent
                 .clone()
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
-        )
-        .new_span(
-            child
-                .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(None),
+                .with_ancestry(expect::is_contextual_root()),
         )
+        .new_span(child.clone().with_ancestry(expect::is_contextual_root()))
         .enter(child.clone())
         .exit(child.clone())
         .enter(contextual_parent.clone())
         .new_span(
             child
                 .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(None),
+                .with_ancestry(expect::has_contextual_parent("contextual_parent")),
         )
         .enter(child.clone())
         .exit(child)
@@ -68,20 +61,14 @@ fn explicit_parent_test() {
         .new_span(
             contextual_parent
                 .clone()
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
-        )
-        .new_span(
-            explicit_parent
-                .with_contextual_parent(None)
-                .with_explicit_parent(None),
+                .with_ancestry(expect::is_contextual_root()),
         )
+        .new_span(explicit_parent.with_ancestry(expect::is_contextual_root()))
         .enter(contextual_parent.clone())
         .new_span(
             child
                 .clone()
-                .with_contextual_parent(Some("contextual_parent"))
-                .with_explicit_parent(Some("explicit_parent")),
+                .with_ancestry(expect::has_explicit_parent("explicit_parent")),
         )
         .enter(child.clone())
         .exit(child)

tracing-futures/tests/std_future.rs

@@ -69,13 +69,21 @@ fn span_on_drop() {
 
     let subscriber = subscriber::mock()
         .enter(expect::span().named("foo"))
-        .event(expect::event().at_level(Level::INFO))
+        .event(
+            expect::event()
+                .with_ancestry(expect::has_contextual_parent("foo"))
+                .at_level(Level::INFO),
+        )
         .exit(expect::span().named("foo"))
         .enter(expect::span().named("foo"))
         .exit(expect::span().named("foo"))
         .drop_span(expect::span().named("foo"))
         .enter(expect::span().named("bar"))
-        .event(expect::event().at_level(Level::INFO))
+        .event(
+            expect::event()
+                .with_ancestry(expect::has_contextual_parent("bar"))
+                .at_level(Level::INFO),
+        )
         .exit(expect::span().named("bar"))
         .drop_span(expect::span().named("bar"))
         .only()

tracing-mock/src/ancestry.rs

@@ -0,0 +1,148 @@
+//! Define the ancestry of an event or span.
+//!
+//! See the documentation on the [`Ancestry`] enum for further details.
+
+use tracing_core::{
+    span::{self, Attributes},
+    Event,
+};
+
+/// The ancestry of an event or span.
+///
+/// An event or span can have an explicitly assigned parent, or be an explicit root. Otherwise,
+/// an event or span may have a contextually assigned parent or in the final case will be a
+/// contextual root.
+#[derive(Debug, Eq, PartialEq)]
+pub enum Ancestry {
+    /// The event or span has an explicitly assigned parent (created with `parent: span_id`) with
+    /// the specified name.
+    HasExplicitParent(String),
+    /// The event or span is an explicitly defined root. It was created with `parent: None` and
+    /// has no parent.
+    IsExplicitRoot,
+    /// The event or span has a contextually assigned parent with the specified name. It has no
+    /// explicitly assigned parent, nor has it been explicitly defined as a root (it was created
+    /// without the `parent:` directive). There was a span in context when this event or span was
+    /// created.
+    HasContextualParent(String),
+    /// The event or span is a contextual root. It has no explicitly assigned parent, nor has it
+    /// been explicitly defined as a root (it was created without the `parent:` directive).
+    /// Additionally, no span was in context when this event or span was created.
+    IsContextualRoot,
+}
+
+impl Ancestry {
+    #[track_caller]
+    pub(crate) fn check(
+        &self,
+        actual_ancestry: &Ancestry,
+        ctx: impl std::fmt::Display,
+        collector_name: &str,
+    ) {
+        let expected_description = |ancestry: &Ancestry| match ancestry {
+            Self::IsExplicitRoot => "be an explicit root".to_string(),
+            Self::HasExplicitParent(name) => format!("have an explicit parent with name='{name}'"),
+            Self::IsContextualRoot => "be a contextual root".to_string(),
+            Self::HasContextualParent(name) => {
+                format!("have a contextual parent with name='{name}'")
+            }
+        };
+
+        let actual_description = |ancestry: &Ancestry| match ancestry {
+            Self::IsExplicitRoot => "was actually an explicit root".to_string(),
+            Self::HasExplicitParent(name) => {
+                format!("actually has an explicit parent with name='{name}'")
+            }
+            Self::IsContextualRoot => "was actually a contextual root".to_string(),
+            Self::HasContextualParent(name) => {
+                format!("actually has a contextual parent with name='{name}'")
+            }
+        };
+
+        assert_eq!(
+            self,
+            actual_ancestry,
+            "[{collector_name}] expected {ctx} to {expected_description}, but {actual_description}",
+            expected_description = expected_description(self),
+            actual_description = actual_description(actual_ancestry)
+        );
+    }
+}
+
+pub(crate) trait HasAncestry {
+    fn is_contextual(&self) -> bool;
+
+    fn is_root(&self) -> bool;
+
+    fn parent(&self) -> Option<&span::Id>;
+}
+
+impl HasAncestry for &Event<'_> {
+    fn is_contextual(&self) -> bool {
+        (self as &Event<'_>).is_contextual()
+    }
+
+    fn is_root(&self) -> bool {
+        (self as &Event<'_>).is_root()
+    }
+
+    fn parent(&self) -> Option<&span::Id> {
+        (self as &Event<'_>).parent()
+    }
+}
+
+impl HasAncestry for &Attributes<'_> {
+    fn is_contextual(&self) -> bool {
+        (self as &Attributes<'_>).is_contextual()
+    }
+
+    fn is_root(&self) -> bool {
+        (self as &Attributes<'_>).is_root()
+    }
+
+    fn parent(&self) -> Option<&span::Id> {
+        (self as &Attributes<'_>).parent()
+    }
+}
+
+/// Determines the ancestry of an actual span or event.
+///
+/// The rules for determining the ancestry are as follows:
+///
+/// +------------+--------------+-----------------+---------------------+
+/// | Contextual | Current Span | Explicit Parent | Ancestry            |
+/// +------------+--------------+-----------------+---------------------+
+/// | Yes        | Yes          | -               | HasContextualParent |
+/// | Yes        | No           | -               | IsContextualRoot    |
+/// | No         | -            | Yes             | HasExplicitParent   |
+/// | No         | -            | No              | IsExplicitRoot      |
+/// +------------+--------------+-----------------+---------------------+
+pub(crate) fn get_ancestry(
+    item: impl HasAncestry,
+    lookup_current: impl FnOnce() -> Option<span::Id>,
+    span_name: impl FnOnce(&span::Id) -> Option<&str>,
+) -> Ancestry {
+    if item.is_contextual() {
+        if let Some(parent_id) = lookup_current() {
+            let contextual_parent_name = span_name(&parent_id).expect(
+                "tracing-mock: contextual parent cannot \
+                            be looked up by ID. Was it recorded correctly?",
+            );
+            Ancestry::HasContextualParent(contextual_parent_name.to_string())
+        } else {
+            Ancestry::IsContextualRoot
+        }
+    } else if item.is_root() {
+        Ancestry::IsExplicitRoot
+    } else {
+        let parent_id = item.parent().expect(
+            "tracing-mock: is_contextual=false is_root=false \
+                        but no explicit parent found. This is a bug!",
+        );
+        let explicit_parent_name = span_name(parent_id).expect(
+            "tracing-mock: explicit parent cannot be looked \
+                        up by ID. Is the provided Span ID valid: {parent_id}",
+        );
+        Ancestry::HasExplicitParent(explicit_parent_name.to_string())
+    }
+}