From ba387ddc03f17eb91e70efe58ae007326ecb9cd4 Mon Sep 17 00:00:00 2001 From: Hayden Stainsby Date: Fri, 21 Jun 2024 11:11:29 +0200 Subject: [PATCH] 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` 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-mock/src/collector.rs | 9 ++ tracing-mock/src/expect.rs | 21 +++- tracing-mock/src/span.rs | 194 +++++++++++++++++++++++++++++++++- 3 files changed, 222 insertions(+), 2 deletions(-) diff --git a/tracing-mock/src/collector.rs b/tracing-mock/src/collector.rs index a91cdd332d..97a0148825 100644 --- a/tracing-mock/src/collector.rs +++ b/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, diff --git a/tracing-mock/src/expect.rs b/tracing-mock/src/expect.rs index 353bc52f5f..beb66d36f9 100644 --- a/tracing-mock/src/expect.rs +++ b/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, what: fmt::Arguments<'_>) { let name = name.as_ref(); diff --git a/tracing-mock/src/span.rs b/tracing-mock/src/span.rs index 42ed3a603e..b666cd7108 100644 --- a/tracing-mock/src/span.rs +++ b/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, 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, +} + 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 {}