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 {}