tracing_causality/
lib.rs

1use std::collections::HashMap;
2use std::fmt::Debug;
3use std::hash::{Hash, Hasher};
4use tracing_core::span::{Attributes, Id};
5use tracing_core::subscriber::Subscriber;
6use tracing_subscriber::layer::Context;
7use tracing_subscriber::registry::ExtensionsMut;
8use tracing_subscriber::registry::LookupSpan;
9use tracing_subscriber::registry::SpanData;
10
11mod channel;
12pub use channel::Updates;
13
14pub mod data;
15
16pub use data::Consequences;
17use data::Listeners;
18
19pub(crate) type Metadata = &'static tracing_core::Metadata<'static>;
20
21/// A causality graph, rooted at a given [`Id`].
22#[derive(Debug, Clone)]
23pub struct Trace<M = crate::Metadata>
24where
25    M: Clone + Debug,
26{
27    root: Span<M>,
28    adj: HashMap<Span<M>, Consequences<M>>,
29}
30
31/// A `Span`, represented by its [`Id`] and metadata.
32#[derive(Debug, Clone)]
33pub struct Span<M = crate::Metadata>
34where
35    M: Clone + Debug,
36{
37    pub id: Id,
38    pub metadata: M,
39}
40
41impl<M> Hash for Span<M>
42where
43    M: Clone + Debug,
44{
45    fn hash<H: Hasher>(&self, state: &mut H) {
46        self.id.hash(state);
47    }
48}
49
50impl<M> Eq for Span<M> where M: Clone + Debug {}
51
52impl<M, N> PartialEq<Span<N>> for Span<M>
53where
54    M: Clone + Debug,
55    N: Clone + Debug,
56{
57    fn eq(&self, other: &Span<N>) -> bool {
58        &self.id == &other.id
59    }
60}
61
62impl<M> Trace<M>
63where
64    M: Clone + Debug,
65{
66    pub fn with_root(root: Span<M>) -> Self {
67        Self {
68            root,
69            adj: Default::default(),
70        }
71    }
72
73    /// The `Id` of the root span of this trace.
74    pub fn root(&self) -> Span<M> {
75        self.root.to_owned()
76    }
77
78    /// The consequences of the `root` span of this trace.
79    pub fn root_consequences(&self) -> &Consequences<M> {
80        self.adj.get(&self.root).unwrap()
81    }
82
83    /// The consequences of the given `id`.
84    pub fn consequences(&self, span: &Span<M>) -> Option<&Consequences<M>> {
85        self.adj.get(span)
86    }
87
88    /// Update [`Trace`] with the given [`Update`].
89    pub fn apply(mut self, update: Update<M>) -> Option<Trace<M>> {
90        match update {
91            Update::OpenDirect { cause, consequence } => {
92                self.adj
93                    .entry(cause)
94                    .or_insert_with(Consequences::none)
95                    .add_direct(consequence);
96                Some(self)
97            }
98            Update::NewIndirect { cause, consequence } => {
99                self.adj
100                    .entry(cause)
101                    .or_insert_with(Consequences::none)
102                    .add_indirect(consequence);
103                Some(self)
104            }
105            Update::CloseDirect { span, direct_cause } => {
106                if let Some(direct_cause) = direct_cause {
107                    let _ = self.adj.remove(&span);
108                    if let Some(consequences) = self.adj.get_mut(&direct_cause) {
109                        consequences.remove_direct(&span);
110                    }
111                    Some(self)
112                } else {
113                    debug_assert_eq!(self.root, span);
114                    None
115                }
116            }
117            Update::CloseIndirect {
118                span: id,
119                indirect_causes,
120            } => {
121                let _ = self.adj.remove(&id);
122
123                for indirect_cause in indirect_causes {
124                    if let Some(consequences) = self.adj.get_mut(&indirect_cause) {
125                        consequences.remove_direct(&id);
126                    }
127                }
128
129                Some(self)
130            }
131            Update::CloseCyclic {
132                span: id,
133                direct_cause,
134                indirect_causes,
135            } => {
136                if let Some(direct_cause) = direct_cause {
137                    let _ = self.adj.remove(&id);
138                    if let Some(consequences) = self.adj.get_mut(&direct_cause) {
139                        consequences.remove_direct(&id);
140                    }
141                    for indirect_cause in indirect_causes {
142                        if let Some(consequences) = self.adj.get_mut(&indirect_cause) {
143                            consequences.remove_direct(&id);
144                        }
145                    }
146                    Some(self)
147                } else {
148                    debug_assert_eq!(self.root, id);
149                    None
150                }
151            }
152        }
153    }
154
155    /// A breadth-first traversal of [`Trace`].
156    pub fn iter(&self) -> impl Iterator<Item = (Span<M>, &Consequences<M>)> {
157        let mut queue = vec![(self.root.clone())];
158        std::iter::from_fn(move || {
159            let span = queue.pop()?;
160            let consequences = self.consequences(&span)?;
161            queue.extend(consequences.iter_direct());
162            Some((span, consequences))
163        })
164    }
165}
166
167fn get_or_init_with<'a, T, F>(extensions: &'a mut ExtensionsMut<'_>, f: F) -> &'a mut T
168where
169    T: 'static + Send + Sync,
170    F: FnOnce() -> T,
171{
172    if extensions.get_mut::<T>().is_none() {
173        extensions.insert::<T>(f());
174    }
175    extensions.get_mut::<T>().unwrap()
176}
177
178/// Produces the full causality graph for the span corresponding to a given
179/// `id`.
180///
181/// Returns both a [`Trace`] rooted at `id`, and a stream of updates
182/// that affect the produced trace, but occurred after the invocation of this
183/// method. If the span has already been closed, this function produces `None`.
184///
185/// ```
186/// use std::sync::Arc;
187/// use tracing_core::Subscriber;
188/// use tracing_causality as causality;
189/// use tracing_subscriber::{prelude::*, registry::Registry};
190///
191/// fn main() {
192///     let subscriber: Arc<dyn Subscriber + Send + Sync > =
193///         Arc::new(Registry::default().with(causality::Layer));
194///     subscriber.clone().init();
195///     let subscriber: Arc<dyn Subscriber> = subscriber;
196///     let subscriber = subscriber.downcast_ref::<Registry>().unwrap();
197///
198///     let a = tracing::trace_span!("a");
199///     let a_id_and_metadata = causality::Span {
200///         id: a.id().unwrap(),
201///         metadata: a.metadata().unwrap()
202///     };
203///
204///     let b = a.in_scope(|| tracing::trace_span!("b"));
205///     let b_id_and_metadata = causality::Span {
206///         id: b.id().unwrap(),
207///         metadata: b.metadata().unwrap()
208///     };
209///
210///     let (trace, updates) = causality::trace(subscriber, &a_id_and_metadata.id, 1).unwrap();
211///     assert!(trace.consequences(&a_id_and_metadata).unwrap().contains_direct(&b_id_and_metadata));
212///
213///     let c = b.in_scope(|| tracing::trace_span!("c"));
214///     let c_id_and_metadata = causality::Span {
215///         id: c.id().unwrap(),
216///         metadata: c.metadata().unwrap()
217///     };
218///     
219///     assert_eq!(
220///         updates.next(),
221///         Some(causality::Update::OpenDirect {
222///             cause: b_id_and_metadata,
223///             consequence: c_id_and_metadata,
224///         })
225///     );
226/// }
227/// ```
228pub fn trace<S>(s: &S, id: &Id, update_capacity: usize) -> Option<(Trace, Updates)>
229where
230    S: for<'span> LookupSpan<'span> + ?Sized,
231{
232    let (sender, updates) = channel::bounded(id.clone(), update_capacity);
233    let root = Span {
234        id: id.clone(),
235        metadata: s.span_data(id)?.metadata(),
236    };
237    let mut trace = Trace {
238        root: root.clone(),
239        adj: HashMap::default(),
240    };
241    let mut queue = vec![root.to_owned()];
242    while let Some(span) = queue.pop() {
243        if let Some(span_data) = s.span_data(&span.id) {
244            let mut extensions = span_data.extensions_mut();
245            // add the update listener
246            get_or_init_with::<Listeners, _>(&mut extensions, Listeners::new)
247                .insert(sender.clone());
248            if let Some(consequences) = extensions.get_mut::<Consequences>() {
249                let direct_consequences = consequences.clone();
250                // add any further consequences to the traversal queue
251                queue.extend(direct_consequences.direct.iter().cloned());
252                // and to the trace
253                trace.adj.insert(span, direct_consequences);
254            }
255        } else {
256            // the span has already been closed; do nothing
257        }
258    }
259    Some((trace, updates))
260}
261
262/// Produces the immediate consequences of the span corresponding to the given
263/// `id`, or `None` if that span has already been closed.
264///
265/// ```
266/// use std::sync::Arc;
267/// use tracing_core::Subscriber;
268/// use tracing_causality::{self as causality, Consequences};
269/// use tracing_subscriber::{prelude::*, registry::Registry};
270///
271/// fn main() {
272///     let subscriber: Arc<dyn Subscriber + Send + Sync > =
273///         Arc::new(Registry::default().with(causality::Layer));
274///     subscriber.clone().init();
275///     let subscriber: Arc<dyn Subscriber> = subscriber;
276///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
277///
278///     let a = tracing::trace_span!("a");
279///     let a_id_and_metadata = causality::Span {
280///         id: a.id().unwrap(),
281///         metadata: a.metadata().unwrap()
282///     };
283///
284///     assert_eq!(
285///         causality::consequences(registry, &a_id_and_metadata.id),
286///         Some(Consequences::none())
287///     );
288///
289///     let b = a.in_scope(|| tracing::trace_span!("b"));
290///     let b_id_and_metadata = causality::Span {
291///         id: b.id().unwrap(),
292///         metadata: b.metadata().unwrap()
293///     };
294///
295///     assert_eq!(
296///         causality::consequences(registry, &a_id_and_metadata.id),
297///         Some(Consequences::with_direct(b_id_and_metadata))
298///     );
299///
300///     drop(b);
301///
302///     assert_eq!(
303///         causality::consequences(registry, &a_id_and_metadata.id),
304///         Some(Consequences::none())
305///     );
306///
307///     drop(a);
308///
309///     assert_eq!(
310///         causality::consequences(registry, &a_id_and_metadata.id),
311///         None
312///     );
313/// }
314/// ```
315pub fn consequences<S>(subscriber: &S, span: &Id) -> Option<Consequences>
316where
317    S: for<'span> LookupSpan<'span> + ?Sized,
318{
319    Some(
320        subscriber
321            .span_data(span)?
322            .extensions()
323            .get::<Consequences>()
324            .cloned()
325            .unwrap_or_else(Consequences::default),
326    )
327}
328
329/// An update that should be applied to a [`Trace`].
330#[derive(Clone, Debug)]
331pub enum Update<M = crate::Metadata>
332where
333    M: Clone + Debug,
334{
335    /// Announces that `consequence` **directly** follows from `cause`.
336    ///
337    /// # Example
338    /// ```
339    /// use std::sync::Arc;
340    /// use tracing::Subscriber;
341    /// use tracing_causality::{self as causality, Update};
342    /// use tracing_subscriber::{prelude::*, registry::Registry};
343    ///
344    /// fn main() {
345    ///     let subscriber: Arc<dyn Subscriber + Send + Sync> =
346    ///         Arc::new(Registry::default().with(causality::Layer));
347    ///     let _guard = subscriber.clone().set_default();
348    ///     let subscriber: Arc<dyn Subscriber> = subscriber;
349    ///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
350    ///
351    ///     let cause = tracing::trace_span!("cause");
352    ///     let cause_id_and_metadata = causality::Span {
353    ///         id: cause.id().unwrap(),
354    ///         metadata: cause.metadata().unwrap()
355    ///     };
356    ///
357    ///     let (_trace, cause_updates) = causality::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
358    ///
359    ///     let consequence = cause.in_scope(|| tracing::trace_span!("consequence"));
360    ///     let consequence_id_and_metadata = causality::Span {
361    ///         id: consequence.id().unwrap(),
362    ///         metadata: consequence.metadata().unwrap()
363    ///     };
364    ///
365    ///     assert_eq!(
366    ///         cause_updates.drain().collect::<Vec<_>>(),
367    ///         vec![Update::OpenDirect {
368    ///             cause: cause_id_and_metadata,
369    ///             consequence: consequence_id_and_metadata,
370    ///         }],
371    ///         "The listeners on `cause` should be notified that it has a \
372    ///         direct `consequence`."
373    ///     );
374    /// }
375    /// ```
376    OpenDirect {
377        cause: Span<M>,
378        consequence: Span<M>,
379    },
380
381    /// Announces that `consequence` **indirectly** follows from `cause`.
382    ///
383    /// # Example
384    /// ```
385    /// use std::sync::Arc;
386    /// use tracing::Subscriber;
387    /// use tracing_causality::{self as causality, Update};
388    /// use tracing_subscriber::{prelude::*, registry::Registry};
389    ///
390    /// fn main() {
391    ///     let subscriber: Arc<dyn Subscriber + Send + Sync> =
392    ///         Arc::new(Registry::default().with(causality::Layer));
393    ///     let _guard = subscriber.clone().set_default();
394    ///     let subscriber: Arc<dyn Subscriber> = subscriber;
395    ///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
396    ///
397    ///     let cause = tracing::trace_span!("cause");
398    ///     let cause_id_and_metadata = causality::Span {
399    ///         id: cause.id().unwrap(),
400    ///         metadata: cause.metadata().unwrap()
401    ///     };
402    ///
403    ///     let (_trace, cause_updates) = causality::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
404    ///
405    ///     let consequence = tracing::trace_span!("consequence");
406    ///     let consequence_id_and_metadata = causality::Span {
407    ///         id: consequence.id().unwrap(),
408    ///         metadata: consequence.metadata().unwrap()
409    ///     };
410    ///     
411    ///     consequence.follows_from(&cause_id_and_metadata.id);
412    ///
413    ///     assert_eq!(
414    ///         cause_updates.drain().collect::<Vec<_>>(),
415    ///         vec![Update::NewIndirect {
416    ///             cause: cause_id_and_metadata.clone(),
417    ///             consequence: consequence_id_and_metadata.clone(),
418    ///         }],
419    ///         "The listeners on `cause` should be notified that it has a \
420    ///         indirect `consequence`."
421    ///     );
422    /// }
423    /// ```
424    NewIndirect {
425        cause: Span<M>,
426        consequence: Span<M>,
427    },
428
429    /// Announces that a direct consequence of a `Span` within [`Trace`] was
430    /// closed, and is thus no longer an extant consequence of `direct_cause`.
431    ///
432    /// # Example
433    /// ```
434    /// use std::sync::Arc;
435    /// use tracing::Subscriber;
436    /// use tracing_causality::{self as causality, Update};
437    /// use tracing_subscriber::{prelude::*, registry::Registry};
438    ///
439    /// fn main() {
440    ///     let subscriber: Arc<dyn Subscriber + Send + Sync> =
441    ///         Arc::new(Registry::default().with(causality::Layer));
442    ///     let _guard = subscriber.clone().set_default();
443    ///     let subscriber: Arc<dyn Subscriber> = subscriber;
444    ///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
445    ///
446    ///     let cause = tracing::trace_span!("cause");
447    ///     let cause_id_and_metadata = causality::Span {
448    ///         id: cause.id().unwrap(),
449    ///         metadata: cause.metadata().unwrap()
450    ///     };
451    ///
452    ///     let consequence = cause.in_scope(|| tracing::trace_span!("consequence"));
453    ///     let consequence_id_and_metadata = causality::Span {
454    ///         id: consequence.id().unwrap(),
455    ///         metadata: consequence.metadata().unwrap()
456    ///     };
457    ///
458    ///     let (_trace, cause_updates) = causality::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
459    ///
460    ///     drop(consequence);
461    ///
462    ///     assert_eq!(
463    ///         cause_updates.drain().collect::<Vec<_>>(),
464    ///         vec![Update::CloseDirect {
465    ///             span: consequence_id_and_metadata.clone(),
466    ///             direct_cause: Some(cause_id_and_metadata),
467    ///         }],
468    ///         "The listeners on `cause` should be notified that
469    ///         `consequence` was closed."
470    ///     );
471    /// }
472    /// ```
473    CloseDirect {
474        span: Span<M>,
475        direct_cause: Option<Span<M>>,
476    },
477
478    /// Announces that an indirect consequence of a `Span` within [`Trace`] was
479    /// closed, and is thus no longer an extant consequence of `indirect_causes`.
480    ///
481    /// # Example
482    /// ```
483    /// use std::sync::Arc;
484    /// use tracing::{Subscriber, trace_span};
485    /// use tracing_causality::{self as causality, Update};
486    /// use tracing_subscriber::{prelude::*, registry::Registry};
487    ///
488    /// fn main() {
489    ///     let subscriber: Arc<dyn Subscriber + Send + Sync> =
490    ///         Arc::new(Registry::default().with(causality::Layer));
491    ///     let _guard = subscriber.clone().set_default();
492    ///     let subscriber: Arc<dyn Subscriber> = subscriber;
493    ///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
494    ///
495    ///     let cause = tracing::trace_span!("cause");
496    ///     let cause_id_and_metadata = causality::Span {
497    ///         id: cause.id().unwrap(),
498    ///         metadata: cause.metadata().unwrap()
499    ///     };
500    ///
501    ///     let consequence = tracing::trace_span!("consequence");
502    ///     let consequence_id_and_metadata = causality::Span {
503    ///         id: consequence.id().unwrap(),
504    ///         metadata: consequence.metadata().unwrap()
505    ///     };
506    ///
507    ///     consequence.follows_from(&cause_id_and_metadata.id);
508    ///
509    ///     let (_trace, cause_updates) = causality::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
510    ///
511    ///     drop(consequence);
512    ///
513    ///     assert_eq!(
514    ///         cause_updates.drain().collect::<Vec<_>>(),
515    ///         vec![Update::CloseIndirect {
516    ///             span: consequence_id_and_metadata,
517    ///             indirect_causes: vec![cause_id_and_metadata],
518    ///         }],
519    ///         "The listeners on `cause` should be notified that
520    ///         `consequence` was closed."
521    ///     );
522    /// }
523    /// ```
524    CloseIndirect {
525        span: Span<M>,
526        indirect_causes: Vec<Span<M>>,
527    },
528
529    /// Announces that a self-cycling consequence of a `Span` within [`Trace`]
530    /// was closed, and is thus no longer an extant consequence of
531    /// `direct_cause` or `indirect_cause`.
532    ///
533    /// # Example
534    /// ```
535    /// use std::sync::Arc;
536    /// use tracing::{Subscriber, trace_span};
537    /// use tracing_causality::{self as causality, Update};
538    /// use tracing_subscriber::{prelude::*, registry::{LookupSpan, SpanData, Registry}};
539    ///
540    /// fn main() {
541    ///     let subscriber: Arc<dyn Subscriber + Send + Sync> =
542    ///         Arc::new(Registry::default().with(causality::Layer));
543    ///     let _guard = subscriber.clone().set_default();
544    ///     let subscriber: Arc<dyn Subscriber> = subscriber;
545    ///     let registry = subscriber.downcast_ref::<Registry>().unwrap();
546    ///
547    ///     let cause = tracing::trace_span!("cause");
548    ///     let cause_id_and_metadata = causality::Span {
549    ///         id: cause.id().unwrap(),
550    ///         metadata: cause.metadata().unwrap()
551    ///     };
552    ///
553    ///     let consequence = cause.clone();
554    ///     let consequence_id_and_metadata = causality::Span {
555    ///         id: consequence.id().unwrap(),
556    ///         metadata: consequence.metadata().unwrap()
557    ///     };
558    ///
559    ///     consequence.follows_from(&cause_id_and_metadata.id);
560    ///
561    ///     let (_trace, cause_updates) = causality::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
562    ///
563    ///     drop([cause, consequence]);
564    ///
565    ///     assert_eq!(
566    ///         cause_updates.into_iter().collect::<Vec<_>>(),
567    ///         vec![Update::CloseCyclic {
568    ///             span: consequence_id_and_metadata.clone(),
569    ///             direct_cause: None,
570    ///             indirect_causes: vec![cause_id_and_metadata.clone()],
571    ///         }],
572    ///         "The listeners on `cause` should be notified that
573    ///         `consequence` was closed."
574    ///     );
575    /// }
576    /// ```
577    CloseCyclic {
578        span: Span<M>,
579        direct_cause: Option<Span<M>>,
580        indirect_causes: Vec<Span<M>>,
581    },
582}
583
584impl<M> Eq for Update<M> where M: Clone + Debug + Eq {}
585
586impl<M> PartialEq<Update<M>> for Update<M>
587where
588    M: Clone + Debug + PartialEq,
589{
590    fn eq(&self, other: &Update<M>) -> bool {
591        use Update::*;
592        match (self, other) {
593            (
594                OpenDirect {
595                    cause: lhs_cause,
596                    consequence: lhs_consequence,
597                },
598                OpenDirect {
599                    cause: rhs_cause,
600                    consequence: rhs_consequence,
601                },
602            )
603            | (
604                NewIndirect {
605                    cause: lhs_cause,
606                    consequence: lhs_consequence,
607                },
608                NewIndirect {
609                    cause: rhs_cause,
610                    consequence: rhs_consequence,
611                },
612            ) => (lhs_cause == rhs_cause) && (lhs_consequence == rhs_consequence),
613            (
614                CloseDirect {
615                    span: lhs_span,
616                    direct_cause: lhs_direct_cause,
617                },
618                CloseDirect {
619                    span: rhs_span,
620                    direct_cause: rhs_direct_cause,
621                },
622            ) => (lhs_span == rhs_span) && (lhs_direct_cause == rhs_direct_cause),
623            (
624                CloseIndirect {
625                    span: lhs_span,
626                    indirect_causes: lhs_indirect_causes,
627                },
628                CloseIndirect {
629                    span: rhs_span,
630                    indirect_causes: rhs_indirect_causes,
631                },
632            ) => (lhs_span == rhs_span) && (lhs_indirect_causes == rhs_indirect_causes),
633            (
634                CloseCyclic {
635                    span: lhs_span,
636                    direct_cause: lhs_direct_cause,
637                    indirect_causes: lhs_indirect_causes,
638                },
639                CloseCyclic {
640                    span: rhs_span,
641                    direct_cause: rhs_direct_cause,
642                    indirect_causes: rhs_indirect_causes,
643                },
644            ) => {
645                (lhs_span == rhs_span)
646                    && (lhs_direct_cause == rhs_direct_cause)
647                    && (lhs_indirect_causes == rhs_indirect_causes)
648            }
649            _ => false,
650        }
651    }
652}
653
654/// A [tracing-subscriber layer] for monitoring the causal relationships between
655/// tracing spans.
656///
657/// [tracing-subscriber layer]: tracing_subscriber::layer::Layer
658pub struct Layer;
659
660impl Layer {
661    /// Record that a span `follows_from` itself.
662    pub(crate) fn on_follows_self<S>(&self, span_id: &Id, ctx: Context<'_, S>)
663    where
664        S: Subscriber + for<'span> LookupSpan<'span>,
665    {
666        use data::IndirectCauses;
667        let span = ctx.span(span_id).expect("Span not found, this is a bug");
668        let mut span_data = span.extensions_mut();
669        let id_and_metadata = Span {
670            id: span_id.to_owned(),
671            metadata: span.metadata(),
672        };
673
674        // 1. insert `consequence` as an indirect consequence of `cause`
675        if let Some(consequences) = span_data.get_mut::<Consequences>() {
676            consequences.indirect.insert(id_and_metadata.clone());
677        } else {
678            span_data.insert(Consequences::with_indirect(id_and_metadata.clone()));
679        }
680
681        // 2. insert `cause` as an indirect cause of `consequence`
682        if let Some(follows_from) = span_data.get_mut::<IndirectCauses>() {
683            follows_from.add_cause(id_and_metadata.clone());
684        } else {
685            span_data.insert(IndirectCauses::with_cause(id_and_metadata.clone()));
686        }
687
688        if let Some(listeners) = span_data.get_mut::<Listeners>() {
689            // 3. notify causes's listeners that it indirectly lead to `consequence`
690            channel::Sender::broadcast(
691                listeners,
692                Update::NewIndirect {
693                    cause: id_and_metadata.clone(),
694                    consequence: id_and_metadata.clone(),
695                },
696            );
697        }
698    }
699}
700
701impl<S> tracing_subscriber::layer::Layer<S> for Layer
702where
703    S: Subscriber + for<'span> LookupSpan<'span>,
704{
705    fn on_new_span(&self, _: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
706        let span = ctx.span(id).expect(
707            "The `id` provided to 
708            `tracing_causality::Layer::on_new_span` did not correspond to an
709            opened `Span` for the underlying subscriber. This span may have \
710            already been closed, or been assigned by a different subscriber, \
711            or there may be a bug in the subscriber underlying this layer.",
712        );
713        let consequence_id_and_metadata = Span {
714            id: id.to_owned(),
715            metadata: span.metadata(),
716        };
717
718        if let Some(direct_cause) = span.parent() {
719            let mut cause_extensions = direct_cause.extensions_mut();
720            let cause_id_and_metadata = Span {
721                id: direct_cause.id(),
722                metadata: direct_cause.metadata(),
723            };
724
725            if let Some(listeners) = cause_extensions.get_mut::<Listeners>() {
726                // 1. notify listeners, if any, that `id` is a consequence
727                channel::Sender::broadcast(
728                    listeners,
729                    Update::OpenDirect {
730                        cause: cause_id_and_metadata,
731                        consequence: consequence_id_and_metadata.clone(),
732                    },
733                );
734
735                // 2. copy cause's listeners onto `consequence`
736                crate::get_or_init_with::<Listeners, _>(&mut span.extensions_mut(), Listeners::new)
737                    .extend(listeners.iter().cloned());
738            }
739
740            // 3. register that `cause` directly lead to `consequence`.
741            if let Some(consequences) = cause_extensions.get_mut::<Consequences>() {
742                consequences.direct.insert(consequence_id_and_metadata);
743            } else {
744                cause_extensions.insert(Consequences::with_direct(consequence_id_and_metadata));
745            }
746        }
747    }
748
749    fn on_follows_from(
750        &self,
751        consequence_id_and_metadata: &Id,
752        cause_id_and_metadata: &Id,
753        ctx: Context<'_, S>,
754    ) {
755        use data::IndirectCauses;
756
757        if cause_id_and_metadata == consequence_id_and_metadata {
758            return self.on_follows_self(consequence_id_and_metadata, ctx);
759        }
760
761        let cause = ctx.span(cause_id_and_metadata).expect(
762            "The `cause_id_and_metadata` provided to 
763            `tracing_causality::Layer::on_follows_from` did not correspond to \
764            an opened `Span` for the underlying subscriber. This span may have \
765            already been closed, or been assigned by a different subscriber, \
766            or there may be a bug in the subscriber underlying this layer.",
767        );
768
769        let consequence = ctx.span(consequence_id_and_metadata).expect(
770            "The `consequence_id_and_metadata` provided to 
771            `tracing_causality::Layer::on_follows_from` did not correspond to \
772            an opened `Span` for the underlying subscriber. This span may have \
773            already been closed, or been assigned by a different subscriber, \
774            or there may be a bug in the subscriber underlying this layer.",
775        );
776
777        let cause_id_and_metadata = Span {
778            id: cause_id_and_metadata.to_owned(),
779            metadata: cause.metadata(),
780        };
781
782        let consequence_id_and_metadata = Span {
783            id: consequence_id_and_metadata.to_owned(),
784            metadata: consequence.metadata(),
785        };
786
787        let mut cause_data = cause.extensions_mut();
788        let mut consequence_data = consequence.extensions_mut();
789
790        // 1. insert `consequence` as an indirect consequence of `cause`
791        if let Some(consequences) = cause_data.get_mut::<Consequences>() {
792            consequences
793                .indirect
794                .insert(consequence_id_and_metadata.clone());
795        } else {
796            cause_data.insert(Consequences::with_indirect(
797                consequence_id_and_metadata.clone(),
798            ));
799        }
800
801        // 2. insert `cause` as an indirect cause of `consequence`
802        if let Some(follows_from) = consequence_data.get_mut::<IndirectCauses>() {
803            follows_from.add_cause(cause_id_and_metadata.clone());
804        } else {
805            consequence_data.insert(IndirectCauses::with_cause(cause_id_and_metadata.clone()));
806        }
807
808        if let Some(listeners) = cause_data.get_mut::<Listeners>() {
809            // 3. notify causes's listeners that it indirectly lead to `consequence`
810            channel::Sender::broadcast(
811                listeners,
812                Update::NewIndirect {
813                    cause: cause_id_and_metadata,
814                    consequence: consequence_id_and_metadata,
815                },
816            );
817        }
818    }
819
820    fn on_close(&self, id: Id, ctx: Context<'_, S>) {
821        use data::IndirectCauses;
822        let span = ctx.span(&id).expect(
823            "The `id` provided to 
824            `tracing_causality::Layer::close` did not correspond to an opened \
825            `Span` for the underlying subscriber. This span have \
826            already been closed, or been assigned by a different subscriber, \
827            or there may be a bug in the subscriber underlying this layer.",
828        );
829
830        let mut extensions = span.extensions_mut();
831
832        let closed_id_and_metadata = Span {
833            id: id.to_owned(),
834            metadata: span.metadata(),
835        };
836
837        // `None` if this span is not its own immediate indirect cause,
838        // otherwise `Some(indirect_causes)`.
839        let mut is_cyclic = None;
840
841        // 1. delete `id` as a consequence from each of its indirect causes
842        if let Some(follows_from) = extensions.remove::<IndirectCauses>() {
843            let indirect_causes: Vec<Span> = follows_from.causes.into_iter().collect();
844
845            let drop_update = Update::CloseIndirect {
846                span: closed_id_and_metadata.clone(),
847                indirect_causes: indirect_causes.clone(),
848            };
849
850            for cause in &indirect_causes {
851                if &cause.id == &id {
852                    is_cyclic = Some(indirect_causes.clone());
853                    continue;
854                } else if let Some(cause) = ctx.span(&cause.id) {
855                    let mut extensions = cause.extensions_mut();
856
857                    if let Some(consequences) = extensions.get_mut::<Consequences>() {
858                        consequences.remove_indirect(&closed_id_and_metadata);
859                    }
860
861                    if let Some(listeners) = extensions.get_mut::<Listeners>() {
862                        channel::Sender::broadcast(listeners, drop_update.clone());
863                    }
864                } else {
865                    // `cause_id_and_metadata` corresponds to a `Span` that has already been
866                    // closed. TODO: investigate this case by throwing a panic
867                    // here, and writing unit tests that trigger it.
868                }
869            }
870        }
871
872        let direct_cause = span.parent();
873
874        // 2. delete `id` as a direct consequence of its parent
875        if let Some(parent) = &direct_cause {
876            let mut parent_extensions = parent.extensions_mut();
877            if let Some(consequences) = parent_extensions.get_mut::<Consequences>() {
878                consequences.remove_direct(&closed_id_and_metadata);
879            }
880        }
881
882        // 3. notify listeners, if any, that `id` was dropped
883        if let Some(listeners) = extensions.get_mut::<Listeners>() {
884            let update = if let Some(indirect_causes) = is_cyclic {
885                Update::CloseCyclic {
886                    span: closed_id_and_metadata,
887                    direct_cause: direct_cause.map(|c| Span {
888                        id: c.id(),
889                        metadata: c.metadata(),
890                    }),
891                    indirect_causes,
892                }
893            } else {
894                Update::CloseDirect {
895                    span: closed_id_and_metadata,
896                    direct_cause: direct_cause.map(|c| Span {
897                        id: c.id(),
898                        metadata: c.metadata(),
899                    }),
900                }
901            };
902            channel::Sender::broadcast(listeners, update);
903        }
904    }
905}
906
907#[cfg(test)]
908mod test {
909    use crate::{self as causality, Consequences, Update};
910    use std::sync::Arc;
911    use tracing_core::Subscriber;
912    use tracing_subscriber::registry::{LookupSpan, SpanData};
913    use tracing_subscriber::{prelude::*, registry::Registry};
914
915    mod trace {
916        use crate::{self as causality};
917        use std::sync::Arc;
918        use tracing_core::Subscriber;
919        use tracing_subscriber::registry::{LookupSpan, SpanData};
920        use tracing_subscriber::{prelude::*, registry::Registry};
921
922        #[test]
923        fn should_install_listener() {
924            let subscriber: Arc<dyn Subscriber + Send + Sync> =
925                Arc::new(Registry::default().with(causality::Layer));
926            let _guard = subscriber.clone().set_default();
927            let subscriber: Arc<dyn Subscriber> = subscriber;
928            let registry = subscriber.downcast_ref::<Registry>().unwrap();
929
930            let a = tracing::trace_span!("a");
931            let a_id_and_metadata = causality::Span {
932                id: a.id().unwrap(),
933                metadata: a.metadata().unwrap(),
934            };
935
936            assert!(registry
937                .span_data(&a_id_and_metadata.id)
938                .unwrap()
939                .extensions()
940                .get::<crate::Listeners>()
941                .is_none());
942
943            let (_trace, _updates) =
944                causality::trace(registry, &a_id_and_metadata.id, 1024).unwrap();
945
946            // after `trace`, there should be 1 listener on `a`
947            assert_eq!(
948                registry
949                    .span_data(&a_id_and_metadata.id)
950                    .unwrap()
951                    .extensions()
952                    .get::<crate::Listeners>()
953                    .expect("there should be listeners on this span")
954                    .len(),
955                1
956            );
957        }
958
959        #[test]
960        fn should_copy_listener() {
961            let subscriber: Arc<dyn Subscriber + Send + Sync> =
962                Arc::new(Registry::default().with(causality::Layer));
963            let _guard = subscriber.clone().set_default();
964            let subscriber: Arc<dyn Subscriber> = subscriber;
965            let registry = subscriber.downcast_ref::<Registry>().unwrap();
966
967            let a = tracing::trace_span!("a");
968            let a_id_and_metadata = causality::Span {
969                id: a.id().unwrap(),
970                metadata: a.metadata().unwrap(),
971            };
972
973            let b = a.in_scope(|| tracing::trace_span!("b"));
974            let b_id_and_metadata = causality::Span {
975                id: b.id().unwrap(),
976                metadata: b.metadata().unwrap(),
977            };
978
979            assert!(registry
980                .span_data(&a_id_and_metadata.id)
981                .unwrap()
982                .extensions()
983                .get::<crate::Listeners>()
984                .is_none());
985
986            assert!(registry
987                .span_data(&b_id_and_metadata.id)
988                .unwrap()
989                .extensions()
990                .get::<crate::Listeners>()
991                .is_none());
992
993            let (_trace, _updates) =
994                causality::trace(registry, &a_id_and_metadata.id, 1024).unwrap();
995
996            // after `trace`, there should be 1 listener on `a`
997            assert_eq!(
998                registry
999                    .span_data(&a_id_and_metadata.id)
1000                    .unwrap()
1001                    .extensions()
1002                    .get::<crate::Listeners>()
1003                    .expect("there should be listeners on this span")
1004                    .len(),
1005                1
1006            );
1007
1008            // after `trace`, there should be 1 listener on `b`
1009            assert_eq!(
1010                registry
1011                    .span_data(&b_id_and_metadata.id)
1012                    .unwrap()
1013                    .extensions()
1014                    .get::<crate::Listeners>()
1015                    .expect("there should be listeners on this span")
1016                    .len(),
1017                1
1018            );
1019        }
1020
1021        #[test]
1022        fn should_not_overwrite_listeners() {
1023            let subscriber: Arc<dyn Subscriber + Send + Sync> =
1024                Arc::new(Registry::default().with(causality::Layer));
1025            let _guard = subscriber.clone().set_default();
1026            let subscriber: Arc<dyn Subscriber> = subscriber;
1027            let registry = subscriber.downcast_ref::<Registry>().unwrap();
1028
1029            let a = tracing::trace_span!("a");
1030            let a_id_and_metadata = causality::Span {
1031                id: a.id().unwrap(),
1032                metadata: a.metadata().unwrap(),
1033            };
1034
1035            let b = a.in_scope(|| tracing::trace_span!("b"));
1036            let b_id_and_metadata = causality::Span {
1037                id: b.id().unwrap(),
1038                metadata: b.metadata().unwrap(),
1039            };
1040
1041            assert!(registry
1042                .span_data(&a_id_and_metadata.id)
1043                .unwrap()
1044                .extensions()
1045                .get::<crate::Listeners>()
1046                .is_none());
1047
1048            assert!(registry
1049                .span_data(&b_id_and_metadata.id)
1050                .unwrap()
1051                .extensions()
1052                .get::<crate::Listeners>()
1053                .is_none());
1054
1055            // trace `b`
1056            let (_trace, _updates) =
1057                causality::trace(registry, &b_id_and_metadata.id, 1024).unwrap();
1058
1059            // after `trace`, there should be 0 listeners on `a`
1060            assert!(registry
1061                .span_data(&a_id_and_metadata.id)
1062                .unwrap()
1063                .extensions()
1064                .get::<crate::Listeners>()
1065                .is_none());
1066
1067            // after `trace`, there should be 1 listener on `b`
1068            assert_eq!(
1069                registry
1070                    .span_data(&b_id_and_metadata.id)
1071                    .unwrap()
1072                    .extensions()
1073                    .get::<crate::Listeners>()
1074                    .expect("there should be listeners on this span")
1075                    .len(),
1076                1
1077            );
1078
1079            // trace `a`
1080            let (_trace, _updates) =
1081                causality::trace(registry, &a_id_and_metadata.id, 1024).unwrap();
1082
1083            // after `trace`, there should be 1 listener on `a`
1084            assert_eq!(
1085                registry
1086                    .span_data(&a_id_and_metadata.id)
1087                    .unwrap()
1088                    .extensions()
1089                    .get::<crate::Listeners>()
1090                    .expect("there should be listeners on this span")
1091                    .len(),
1092                1
1093            );
1094
1095            // after `trace`, there should be 2 listeners on `b`
1096            assert_eq!(
1097                registry
1098                    .span_data(&b_id_and_metadata.id)
1099                    .unwrap()
1100                    .extensions()
1101                    .get::<crate::Listeners>()
1102                    .expect("there should be listeners on this span")
1103                    .len(),
1104                2
1105            );
1106        }
1107    }
1108
1109    mod layer {
1110        mod on_new_span {
1111            use crate::test::*;
1112
1113            /// If the span's parent has listeners, those listeners should be notified
1114            /// that the parent has a new consequence.
1115            #[test]
1116            fn should_notify_listeners() {
1117                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1118                    Arc::new(Registry::default().with(causality::Layer));
1119                let _guard = subscriber.clone().set_default();
1120                let subscriber: Arc<dyn Subscriber> = subscriber;
1121                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1122
1123                let a = tracing::trace_span!("a");
1124                let a_id_and_metadata = causality::Span {
1125                    id: a.id().unwrap(),
1126                    metadata: a.metadata().unwrap(),
1127                };
1128
1129                // prior to `trace`, there are no listeners on `a`
1130                assert!(registry
1131                    .span_data(&a_id_and_metadata.id)
1132                    .unwrap()
1133                    .extensions()
1134                    .get::<crate::Listeners>()
1135                    .is_none());
1136
1137                let (_trace, updates) =
1138                    causality::trace(registry, &a_id_and_metadata.id, 1024).unwrap();
1139
1140                // after `trace`, there should be 1 listener on `a`
1141                assert_eq!(
1142                    registry
1143                        .span_data(&a_id_and_metadata.id)
1144                        .unwrap()
1145                        .extensions()
1146                        .get::<crate::Listeners>()
1147                        .expect("there should be listeners on this span")
1148                        .len(),
1149                    1
1150                );
1151
1152                let b = a.in_scope(|| tracing::trace_span!("b"));
1153                let b_id_and_metadata = causality::Span {
1154                    id: b.id().unwrap(),
1155                    metadata: b.metadata().unwrap(),
1156                };
1157
1158                assert_eq!(
1159                    updates.next(),
1160                    Some(Update::OpenDirect {
1161                        cause: a_id_and_metadata,
1162                        consequence: b_id_and_metadata,
1163                    })
1164                );
1165
1166                assert!(updates.is_empty());
1167            }
1168
1169            /// If the parent of a new span has update listeners, those listeners should
1170            /// be copied to the new span.
1171            #[test]
1172            fn should_copy_listeners() {
1173                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1174                    Arc::new(Registry::default().with(causality::Layer));
1175                let _guard = subscriber.clone().set_default();
1176                let subscriber: Arc<dyn Subscriber> = subscriber;
1177                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1178
1179                let a = tracing::trace_span!("a");
1180                let a_id_and_metadata = causality::Span {
1181                    id: a.id().unwrap(),
1182                    metadata: a.metadata().unwrap(),
1183                };
1184
1185                let (_trace, _updates) =
1186                    causality::trace(registry, &a_id_and_metadata.id, 1024).unwrap();
1187
1188                let b = a.in_scope(|| tracing::trace_span!("b"));
1189                let b_id_and_metadata = causality::Span {
1190                    id: b.id().unwrap(),
1191                    metadata: b.metadata().unwrap(),
1192                };
1193
1194                let (_trace, _updates) =
1195                    causality::trace(registry, &b_id_and_metadata.id, 1024).unwrap();
1196
1197                let a_listeners = registry
1198                    .span_data(&a_id_and_metadata.id)
1199                    .unwrap()
1200                    .extensions()
1201                    .get::<crate::Listeners>()
1202                    .expect("there should be listeners on span `a`")
1203                    .clone();
1204
1205                let b_listeners = registry
1206                    .span_data(&b_id_and_metadata.id)
1207                    .unwrap()
1208                    .extensions()
1209                    .get::<crate::Listeners>()
1210                    .expect("there should be listeners on span `b`")
1211                    .clone();
1212
1213                dbg!(&a_listeners);
1214                dbg!(&b_listeners);
1215
1216                assert!(
1217                    b_listeners.is_superset(&a_listeners),
1218                    "the listeners on `a` should have been copied to `b`"
1219                );
1220
1221                assert_ne!(
1222                    b_listeners,
1223                    a_listeners,
1224                    "the listeners of `b` should not have been simply replaced the listeners on `a`"
1225                );
1226            }
1227
1228            /// If the new span as a parent, the new span should be recorded as a direct
1229            /// consequence of the parent.
1230            #[test]
1231            fn should_record_consequence() {
1232                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1233                    Arc::new(Registry::default().with(causality::Layer));
1234                let _guard = subscriber.clone().set_default();
1235                let subscriber: Arc<dyn Subscriber> = subscriber;
1236                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1237
1238                let a = tracing::trace_span!("a");
1239                let a_id_and_metadata = causality::Span {
1240                    id: a.id().unwrap(),
1241                    metadata: a.metadata().unwrap(),
1242                };
1243
1244                let a_consequences = causality::consequences(registry, &a_id_and_metadata.id)
1245                    .expect("span `a` should not have been closed yet");
1246
1247                assert_eq!(
1248                    a_consequences,
1249                    Consequences::default(),
1250                    "span `a` should not have any consequences"
1251                );
1252
1253                let b = a.in_scope(|| tracing::trace_span!("b"));
1254                let b_id_and_metadata = causality::Span {
1255                    id: b.id().unwrap(),
1256                    metadata: b.metadata().unwrap(),
1257                };
1258
1259                let a_consequences = causality::consequences(registry, &a_id_and_metadata.id)
1260                    .expect("span `a` should not have been closed yet");
1261
1262                assert_eq!(
1263                    a_consequences,
1264                    Consequences::with_direct(b_id_and_metadata),
1265                    "span `a` should only have the direct consequence `b`"
1266                );
1267            }
1268        }
1269
1270        mod on_follows_from {
1271            use crate::test::*;
1272
1273            /// Upon `consequence.follows_from(cause)`, the layer should record
1274            /// that `consequence` indirectly follows from `cause` in `causes`'s
1275            /// `Consequences`.
1276            #[test]
1277            fn should_record_consequence() {
1278                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1279                    Arc::new(Registry::default().with(causality::Layer));
1280                let _guard = subscriber.clone().set_default();
1281                let subscriber: Arc<dyn Subscriber> = subscriber;
1282                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1283
1284                let cause = tracing::trace_span!("cause");
1285                let cause_id_and_metadata = causality::Span {
1286                    id: cause.id().unwrap(),
1287                    metadata: cause.metadata().unwrap(),
1288                };
1289
1290                let consequence = tracing::trace_span!("consequence");
1291                let consequence_id_and_metadata = causality::Span {
1292                    id: consequence.id().unwrap(),
1293                    metadata: consequence.metadata().unwrap(),
1294                };
1295
1296                let consequences = causality::consequences(registry, &cause_id_and_metadata.id)
1297                    .expect("span `cause` should not have been closed yet");
1298
1299                assert_eq!(
1300                    consequences,
1301                    Consequences::default(),
1302                    "span `cause` should not have any consequences"
1303                );
1304
1305                consequence.follows_from(&cause_id_and_metadata.id);
1306
1307                let consequences = causality::consequences(registry, &cause_id_and_metadata.id)
1308                    .expect("span `cause` should not have been closed yet");
1309
1310                assert_eq!(
1311                    consequences,
1312                    Consequences::with_indirect(consequence_id_and_metadata),
1313                    "span `cause` should have an indirect `consequence`"
1314                );
1315            }
1316
1317            /// Upon `consequence.follows_from(cause)`, the layer should record
1318            /// that `cause` indirectly led to `consequence` in `consequence`'s
1319            /// `IndirectCauses`.
1320            #[test]
1321            fn should_record_cause() {
1322                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1323                    Arc::new(Registry::default().with(causality::Layer));
1324                let _guard = subscriber.clone().set_default();
1325                let subscriber: Arc<dyn Subscriber> = subscriber;
1326                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1327
1328                let cause = tracing::trace_span!("cause");
1329                let cause_id_and_metadata = causality::Span {
1330                    id: cause.id().unwrap(),
1331                    metadata: cause.metadata().unwrap(),
1332                };
1333
1334                let consequence = tracing::trace_span!("consequence");
1335                let consequence_id_and_metadata = causality::Span {
1336                    id: consequence.id().unwrap(),
1337                    metadata: consequence.metadata().unwrap(),
1338                };
1339
1340                assert!(
1341                    registry
1342                        .span_data(&consequence_id_and_metadata.id)
1343                        .expect("span `consequence` should not yet be closed.")
1344                        .extensions()
1345                        .get::<crate::data::IndirectCauses>()
1346                        .is_none(),
1347                    "span `consequence` should not yet have `IndirectCauses`"
1348                );
1349
1350                consequence.follows_from(&cause_id_and_metadata.id);
1351
1352                assert!(
1353                    registry
1354                        .span_data(&consequence_id_and_metadata.id)
1355                        .expect("span `consequence` should not yet be closed.")
1356                        .extensions()
1357                        .get::<crate::data::IndirectCauses>()
1358                        .expect("span `consequence` should have `IndirectCauses`")
1359                        .contains(&cause_id_and_metadata),
1360                    "`consequence`'s `IndirectCauses` should contain `cause_id_and_metadata`"
1361                );
1362            }
1363
1364            /// Upon `consequence.follows_from(cause)`, the layer should notify
1365            /// `cause`'s listeners that `cause` indirectly led to `consequence`.
1366            #[test]
1367            fn should_notify_listeners() {
1368                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1369                    Arc::new(Registry::default().with(causality::Layer));
1370                let _guard = subscriber.clone().set_default();
1371                let subscriber: Arc<dyn Subscriber> = subscriber;
1372                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1373
1374                let cause = tracing::trace_span!("cause");
1375                let cause_id_and_metadata = causality::Span {
1376                    id: cause.id().unwrap(),
1377                    metadata: cause.metadata().unwrap(),
1378                };
1379
1380                let consequence = tracing::trace_span!("consequence");
1381                let consequence_id_and_metadata = causality::Span {
1382                    id: consequence.id().unwrap(),
1383                    metadata: consequence.metadata().unwrap(),
1384                };
1385
1386                let (_trace, cause_updates) =
1387                    crate::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
1388                let (_trace, consequence_updates) =
1389                    crate::trace(registry, &consequence_id_and_metadata.id, 1024).unwrap();
1390
1391                assert!(consequence_updates.is_empty());
1392                assert!(cause_updates.is_empty());
1393
1394                consequence.follows_from(&cause_id_and_metadata.id);
1395
1396                assert!(
1397                    consequence_updates.is_empty(),
1398                    "The listeners on `consequence` should not have been \
1399                    notified of anything."
1400                );
1401                assert_eq!(
1402                    cause_updates.next(),
1403                    Some(Update::NewIndirect {
1404                        cause: cause_id_and_metadata.clone(),
1405                        consequence: consequence_id_and_metadata.clone(),
1406                    }),
1407                    "The listeners on `cause` should be notified that \
1408                    `consequence` indirectly follows from `cause`."
1409                );
1410                assert!(cause_updates.is_empty());
1411            }
1412        }
1413
1414        mod on_close {
1415            use crate::test::*;
1416
1417            /// Upon the closure of a span, the layer should, for each indirect
1418            /// cause of that span, erase the closed span as an indirect
1419            /// consequence.
1420            #[test]
1421            fn should_erase_consequence() {
1422                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1423                    Arc::new(Registry::default().with(causality::Layer));
1424                let _guard = subscriber.clone().set_default();
1425                let subscriber: Arc<dyn Subscriber> = subscriber;
1426                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1427
1428                let cause = tracing::trace_span!("cause");
1429                let cause_id_and_metadata = causality::Span {
1430                    id: cause.id().unwrap(),
1431                    metadata: cause.metadata().unwrap(),
1432                };
1433
1434                let consequence = tracing::trace_span!("consequence");
1435                let consequence_id_and_metadata = causality::Span {
1436                    id: consequence.id().unwrap(),
1437                    metadata: consequence.metadata().unwrap(),
1438                };
1439
1440                let (_trace, _updates) =
1441                    crate::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
1442                let (_trace, _updates) =
1443                    crate::trace(registry, &consequence_id_and_metadata.id, 1024).unwrap();
1444
1445                assert_eq!(
1446                    causality::consequences(registry, &cause_id_and_metadata.id)
1447                        .expect("span `cause` should not have been closed yet"),
1448                    Consequences::default(),
1449                    "span `cause` should not have any consequences"
1450                );
1451
1452                consequence.follows_from(&cause_id_and_metadata.id);
1453
1454                assert_eq!(
1455                    causality::consequences(registry, &cause_id_and_metadata.id)
1456                        .expect("span `cause` should not have been closed yet"),
1457                    Consequences::with_indirect(consequence_id_and_metadata.clone()),
1458                    "span `cause` should have one indirect consequence"
1459                );
1460
1461                drop(consequence);
1462
1463                assert_eq!(
1464                    causality::consequences(registry, &cause_id_and_metadata.id)
1465                        .expect("span `cause` should not have been closed yet"),
1466                    Consequences::default(),
1467                    "span `cause` should not have any consequences"
1468                );
1469            }
1470
1471            /// Upon the closure of a span, the layer should, for each indirect
1472            /// cause of that span, notify those cause's listeners that this
1473            /// indirect within their [`Trace`] has been closed.
1474            #[test]
1475            fn should_notify_causes_acyclic() {
1476                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1477                    Arc::new(Registry::default().with(causality::Layer));
1478                let _guard = subscriber.clone().set_default();
1479                let subscriber: Arc<dyn Subscriber> = subscriber;
1480                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1481
1482                let cause = tracing::trace_span!("cause");
1483                let cause_id_and_metadata = causality::Span {
1484                    id: cause.id().unwrap(),
1485                    metadata: cause.metadata().unwrap(),
1486                };
1487
1488                let consequence = tracing::trace_span!("consequence");
1489                let consequence_id_and_metadata = causality::Span {
1490                    id: consequence.id().unwrap(),
1491                    metadata: consequence.metadata().unwrap(),
1492                };
1493
1494                consequence.follows_from(&cause_id_and_metadata.id);
1495
1496                let (_trace, cause_updates) =
1497                    crate::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
1498                let (_trace, _consequence_updates) =
1499                    crate::trace(registry, &consequence_id_and_metadata.id, 1024).unwrap();
1500
1501                drop(consequence);
1502
1503                assert_eq!(
1504                    cause_updates.next(),
1505                    Some(Update::CloseIndirect {
1506                        span: consequence_id_and_metadata.clone(),
1507                        indirect_causes: vec![cause_id_and_metadata.clone()],
1508                    }),
1509                    "The listeners on `cause` should be notified that
1510                    `consequence` was closed."
1511                );
1512                assert!(cause_updates.is_empty());
1513            }
1514
1515            /// Upon the closure of a span, the layer should, for each indirect
1516            /// cause of that span, notify those cause's listeners that this
1517            /// indirect within their [`Trace`] has been closed.
1518            #[test]
1519            fn should_notify_causes_cyclic() {
1520                let subscriber: Arc<dyn Subscriber + Send + Sync> =
1521                    Arc::new(Registry::default().with(causality::Layer));
1522                let _guard = subscriber.clone().set_default();
1523                let subscriber: Arc<dyn Subscriber> = subscriber;
1524                let registry = subscriber.downcast_ref::<Registry>().unwrap();
1525
1526                let cause = tracing::trace_span!("cause");
1527                let cause_id_and_metadata = causality::Span {
1528                    id: cause.id().unwrap(),
1529                    metadata: cause.metadata().unwrap(),
1530                };
1531
1532                let consequence = cause.clone();
1533                let consequence_id_and_metadata = causality::Span {
1534                    id: consequence.id().unwrap(),
1535                    metadata: consequence.metadata().unwrap(),
1536                };
1537
1538                consequence.follows_from(&cause_id_and_metadata.id);
1539
1540                let (_trace, cause_updates) =
1541                    crate::trace(registry, &cause_id_and_metadata.id, 1024).unwrap();
1542                let (_trace, _consequence_updates) =
1543                    crate::trace(registry, &consequence_id_and_metadata.id, 1024).unwrap();
1544
1545                drop([cause, consequence]);
1546
1547                assert_eq!(
1548                    cause_updates.next(),
1549                    Some(Update::CloseCyclic {
1550                        span: consequence_id_and_metadata.clone(),
1551                        direct_cause: None,
1552                        indirect_causes: vec![cause_id_and_metadata.clone()],
1553                    }),
1554                    "The listeners on `cause` should be notified that
1555                    `consequence` was closed."
1556                );
1557
1558                assert!(cause_updates.is_empty());
1559            }
1560        }
1561    }
1562}
1563
1564#[cfg(test)]
1565mod test2 {
1566    use crate::{self as causality};
1567    use std::sync::Arc;
1568    use tracing_core::Subscriber;
1569    use tracing_subscriber::registry::{LookupSpan, SpanData};
1570    use tracing_subscriber::{prelude::*, registry::Registry};
1571
1572    #[test]
1573    fn should_update_transitively_1() {
1574        let subscriber: Arc<dyn Subscriber + Send + Sync> =
1575            Arc::new(Registry::default().with(crate::Layer));
1576        let _guard = subscriber.clone().set_default();
1577        let subscriber: Arc<dyn Subscriber> = subscriber;
1578        let subscriber = subscriber.downcast_ref::<Registry>().unwrap();
1579
1580        let a = tracing::trace_span!("a");
1581        let a_id_and_metadata = causality::Span {
1582            id: a.id().unwrap(),
1583            metadata: a.metadata().unwrap(),
1584        };
1585
1586        let b = a.in_scope(|| tracing::trace_span!("b"));
1587        let b_id_and_metadata = causality::Span {
1588            id: b.id().unwrap(),
1589            metadata: b.metadata().unwrap(),
1590        };
1591
1592        let (trace, updates) = crate::trace(subscriber, &a_id_and_metadata.id, 1).unwrap();
1593        assert!(trace
1594            .consequences(&a_id_and_metadata)
1595            .unwrap()
1596            .contains_direct(&b_id_and_metadata));
1597
1598        let c = b.in_scope(|| tracing::trace_span!("c"));
1599        let c_id_and_metadata = causality::Span {
1600            id: c.id().unwrap(),
1601            metadata: c.metadata().unwrap(),
1602        };
1603
1604        dbg!(subscriber
1605            .span_data(&b_id_and_metadata.id)
1606            .unwrap()
1607            .extensions()
1608            .get::<crate::Listeners>()
1609            .is_some());
1610
1611        assert_eq!(
1612            updates.next(),
1613            Some(crate::Update::OpenDirect {
1614                cause: b_id_and_metadata,
1615                consequence: c_id_and_metadata,
1616            })
1617        );
1618    }
1619
1620    #[test]
1621    fn should_update_transitively_2() {
1622        let subscriber: Arc<dyn Subscriber + Send + Sync> =
1623            Arc::new(Registry::default().with(crate::Layer));
1624        let _guard = subscriber.clone().set_default();
1625        let subscriber: Arc<dyn Subscriber> = subscriber;
1626        let subscriber = subscriber.downcast_ref::<Registry>().unwrap();
1627
1628        let a = tracing::trace_span!("a");
1629        let a_id_and_metadata = causality::Span {
1630            id: a.id().unwrap(),
1631            metadata: a.metadata().unwrap(),
1632        };
1633
1634        let (_trace, updates) = crate::trace(subscriber, &a_id_and_metadata.id, 1024).unwrap();
1635
1636        let b = a.in_scope(|| tracing::trace_span!("b"));
1637        let b_id_and_metadata = causality::Span {
1638            id: b.id().unwrap(),
1639            metadata: b.metadata().unwrap(),
1640        };
1641
1642        let c = b.in_scope(|| tracing::trace_span!("c"));
1643        let c_id_and_metadata = causality::Span {
1644            id: c.id().unwrap(),
1645            metadata: c.metadata().unwrap(),
1646        };
1647
1648        dbg!(subscriber
1649            .span_data(&b_id_and_metadata.id)
1650            .unwrap()
1651            .extensions()
1652            .get::<crate::Listeners>()
1653            .is_some());
1654
1655        dbg!(subscriber
1656            .span_data(&c_id_and_metadata.id)
1657            .unwrap()
1658            .extensions()
1659            .get::<crate::Listeners>()
1660            .is_some());
1661
1662        assert_eq!(
1663            updates.next(),
1664            Some(crate::Update::OpenDirect {
1665                cause: a_id_and_metadata,
1666                consequence: b_id_and_metadata.clone(),
1667            })
1668        );
1669
1670        assert_eq!(
1671            updates.next(),
1672            Some(crate::Update::OpenDirect {
1673                cause: b_id_and_metadata,
1674                consequence: c_id_and_metadata,
1675            })
1676        );
1677    }
1678}
tracing_causality/lib.rs

tracing_causality/
lib.rs
