mock: correct contextual/explicit parent assertions (#3004)

## Motivation

When recording the parent of an event or span, the `MockCollector`
treats an explicit parent of `None` (i.e. an event or span that is an
explicit root) in the same way as if there is no explicit root. This
leads to it picking up the contextual parent or treating the event or
span as a contextual root.

## Solution

This change refactors the recording of the parent to use `is_contextual`
to distinguish whether or not an explicit parent has been specified. The
actual parent is also written into an `Ancestry` enum so that the
expected and actual values can be compared in a more explicit way.

Additionally, the `Ancestry` struct has been moved into its own module and
the check behavior has been fixed. The error message has also been
unified across all cases.

Another problem with the previous API is that the two methods
`with_contextual_parent` and `with_explicit_parent` are actually
mutually exclusive, a span or event cannot be both of them. It is also a
(small) mental leap for the user to go from `with_*_parent(None)` to
understanding that this means that a span or event is a root (either
contextual or explicit).

As such, the API has been reworked into a single method `with_ancestry`,
which takes an enum with the following four variants:
* `HasExplicitParent(String)` (parent span name)
* `IsExplicitRoot`
* `HasContextualParent(String)` (parent span name)
* `IsContextualRoot`

To make the interface as useable as possible, helper functions have been
defined in the `expect` module which can be used to create the enum
variants. Specifically, these take `Into<String>` parameter for the span
name.

Given the number of different cases involved in checking ancestry,
separate integration tests have been added to `tracing-mock`
specifically for testing all the positive and negative cases when
asserting on the ancestry of events and spans.

There were two tests in `tracing-attributes` which specified both an
explicit and a contextual parent. This behavior was never intended to
work as all events and spans are either contextual or not. The tests
have been corrected to only expect one of the two.

Fixes: #2440

Author: hds

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

@@ -71,7 +71,7 @@ fn span_on_drop() {
         .enter(expect::span().named("foo"))
         .event(
             expect::event()
-                .with_contextual_parent(Some("foo"))
+                .with_ancestry(expect::has_contextual_parent("foo"))
                 .at_level(Level::INFO),
         )
         .exit(expect::span().named("foo"))
@@ -81,7 +81,7 @@ fn span_on_drop() {
         .enter(expect::span().named("bar"))
         .event(
             expect::event()
-                .with_contextual_parent(Some("bar"))
+                .with_ancestry(expect::has_contextual_parent("bar"))
                 .at_level(Level::INFO),
         )
         .exit(expect::span().named("bar"))

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())
+    }
+}

tracing-mock/src/collector.rs

@@ -138,6 +138,7 @@
 //! [`Collect`]: trait@tracing::Collect
 //! [`MockCollector`]: struct@crate::collector::MockCollector
 use crate::{
+    ancestry::get_ancestry,
     event::ExpectedEvent,
     expect::Expect,
     field::ExpectedFields,
@@ -1039,16 +1040,20 @@ where
                         )
                     }
                 }
-                let get_parent_name = || {
-                    let stack = self.current.lock().unwrap();
-                    let spans = self.spans.lock().unwrap();
-                    event
-                        .parent()
-                        .and_then(|id| spans.get(id))
-                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
-                        .map(|s| s.name.to_string())
+                let event_get_ancestry = || {
+                    get_ancestry(
+                        event,
+                        || self.lookup_current(),
+                        |span_id| {
+                            self.spans
+                                .lock()
+                                .unwrap()
+                                .get(span_id)
+                                .map(|span| span.name)
+                        },
+                    )
                 };
-                expected.check(event, get_parent_name, &self.name);
+                expected.check(event, event_get_ancestry, &self.name);
             }
             Some(ex) => ex.bad(&self.name, format_args!("observed event {:#?}", event)),
         }
@@ -1108,14 +1113,18 @@ where
                 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()
-                        .and_then(|id| spans.get(id))
-                        .or_else(|| stack.last().and_then(|id| spans.get(id)))
-                        .map(|s| s.name.to_string())
-                };
-                expected.check(span, get_parent_name, &self.name);
+
+                expected.check(
+                    span,
+                    || {
+                        get_ancestry(
+                            span,
+                            || self.lookup_current(),
+                            |span_id| spans.get(span_id).map(|span| span.name),
+                        )
+                    },
+                    &self.name,
+                );
             }
         }
         spans.insert(
@@ -1265,6 +1274,16 @@ where
     }
 }
 
+impl<F> Running<F>
+where
+    F: Fn(&Metadata<'_>) -> bool,
+{
+    fn lookup_current(&self) -> Option<span::Id> {
+        let stack = self.current.lock().unwrap();
+        stack.last().cloned()
+    }
+}
+
 impl MockHandle {
     #[cfg(feature = "tracing-subscriber")]
     pub(crate) fn new(expected: Arc<Mutex<VecDeque<Expect>>>, name: String) -> Self {