use sharded_slab::{pool::Ref, Clear, Pool};
use thread_local::ThreadLocal;
use super::stack::SpanStack;
use crate::{
registry::{
extensions::{Extensions, ExtensionsInner, ExtensionsMut},
LookupSpan, SpanData,
},
sync::RwLock,
};
use std::{
cell::{Cell, RefCell},
sync::atomic::{fence, AtomicUsize, Ordering},
};
use tracing_core::{
dispatcher::{self, Dispatch},
span::{self, Current, Id},
Event, Interest, Metadata, Subscriber,
};
#[cfg(feature = "registry")]
#[cfg_attr(docsrs, doc(cfg(feature = "registry")))]
#[derive(Debug)]
pub struct Registry {
spans: Pool<DataInner>,
current_spans: ThreadLocal<RefCell<SpanStack>>,
}
#[cfg(feature = "registry")]
#[cfg_attr(docsrs, doc(cfg(feature = "registry")))]
#[derive(Debug)]
pub struct Data<'a> {
inner: Ref<'a, DataInner>,
}
#[derive(Debug)]
struct DataInner {
metadata: &'static Metadata<'static>,
parent: Option<Id>,
ref_count: AtomicUsize,
pub(crate) extensions: RwLock<ExtensionsInner>,
}
impl Default for Registry {
fn default() -> Self {
Self {
spans: Pool::new(),
current_spans: ThreadLocal::new(),
}
}
}
#[inline]
fn idx_to_id(idx: usize) -> Id {
Id::from_u64(idx as u64 + 1)
}
#[inline]
fn id_to_idx(id: &Id) -> usize {
id.into_u64() as usize - 1
}
pub(crate) struct CloseGuard<'a> {
id: Id,
registry: &'a Registry,
is_closing: bool,
}
impl Registry {
fn get(&self, id: &Id) -> Option<Ref<'_, DataInner>> {
self.spans.get(id_to_idx(id))
}
pub(crate) fn start_close(&self, id: Id) -> CloseGuard<'_> {
CLOSE_COUNT.with(|count| {
let c = count.get();
count.set(c + 1);
});
CloseGuard {
id,
registry: &self,
is_closing: false,
}
}
}
thread_local! {
static CLOSE_COUNT: Cell<usize> = Cell::new(0);
}
impl Subscriber for Registry {
fn register_callsite(&self, _: &'static Metadata<'static>) -> Interest {
Interest::always()
}
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
#[inline]
fn new_span(&self, attrs: &span::Attributes<'_>) -> span::Id {
let parent = if attrs.is_root() {
None
} else if attrs.is_contextual() {
self.current_span().id().map(|id| self.clone_span(id))
} else {
attrs.parent().map(|id| self.clone_span(id))
};
let id = self
.spans
.create_with(|data| {
data.metadata = attrs.metadata();
data.parent = parent;
let refs = data.ref_count.get_mut();
debug_assert_eq!(*refs, 0);
*refs = 1;
})
.expect("Unable to allocate another span");
idx_to_id(id)
}
#[inline]
fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
fn record_follows_from(&self, _span: &span::Id, _follows: &span::Id) {}
fn event(&self, _: &Event<'_>) {}
fn enter(&self, id: &span::Id) {
if self
.current_spans
.get_or_default()
.borrow_mut()
.push(id.clone())
{
self.clone_span(id);
}
}
fn exit(&self, id: &span::Id) {
if let Some(spans) = self.current_spans.get() {
if spans.borrow_mut().pop(id) {
dispatcher::get_default(|dispatch| dispatch.try_close(id.clone()));
}
}
}
fn clone_span(&self, id: &span::Id) -> span::Id {
let span = self
.get(&id)
.unwrap_or_else(|| panic!("tried to clone {:?}, but no span exists with that ID", id));
let refs = span.ref_count.fetch_add(1, Ordering::Relaxed);
assert!(
refs != 0,
"tried to clone a span ({:?}) that already closed",
id
);
id.clone()
}
fn current_span(&self) -> Current {
self.current_spans
.get()
.and_then(|spans| {
let spans = spans.borrow();
let id = spans.current()?;
let span = self.get(id)?;
Some(Current::new(id.clone(), span.metadata))
})
.unwrap_or_else(Current::none)
}
fn try_close(&self, id: span::Id) -> bool {
let span = match self.get(&id) {
Some(span) => span,
None if std::thread::panicking() => return false,
None => panic!("tried to drop a ref to {:?}, but no such span exists!", id),
};
let refs = span.ref_count.fetch_sub(1, Ordering::Release);
if !std::thread::panicking() {
assert!(refs < std::usize::MAX, "reference count overflow!");
}
if refs > 1 {
return false;
}
fence(Ordering::Acquire);
true
}
}
impl<'a> LookupSpan<'a> for Registry {
type Data = Data<'a>;
fn span_data(&'a self, id: &Id) -> Option<Self::Data> {
let inner = self.get(id)?;
Some(Data { inner })
}
}
impl<'a> CloseGuard<'a> {
pub(crate) fn is_closing(&mut self) {
self.is_closing = true;
}
}
impl<'a> Drop for CloseGuard<'a> {
fn drop(&mut self) {
let _ = CLOSE_COUNT.try_with(|count| {
let c = count.get();
count.set(c - 1);
if c == 1 && self.is_closing {
self.registry.spans.clear(id_to_idx(&self.id));
}
});
}
}
impl<'a> SpanData<'a> for Data<'a> {
fn id(&self) -> Id {
idx_to_id(self.inner.key())
}
fn metadata(&self) -> &'static Metadata<'static> {
(*self).inner.metadata
}
fn parent(&self) -> Option<&Id> {
self.inner.parent.as_ref()
}
fn extensions(&self) -> Extensions<'_> {
Extensions::new(self.inner.extensions.read().expect("Mutex poisoned"))
}
fn extensions_mut(&self) -> ExtensionsMut<'_> {
ExtensionsMut::new(self.inner.extensions.write().expect("Mutex poisoned"))
}
}
impl Default for DataInner {
fn default() -> Self {
struct NullCallsite;
impl tracing_core::callsite::Callsite for NullCallsite {
fn set_interest(&self, _: Interest) {
unreachable!(
"/!\\ Tried to register the null callsite /!\\\n \
This should never have happened and is definitely a bug. \
A `tracing` bug report would be appreciated."
)
}
fn metadata(&self) -> &Metadata<'_> {
unreachable!(
"/!\\ Tried to access the null callsite's metadata /!\\\n \
This should never have happened and is definitely a bug. \
A `tracing` bug report would be appreciated."
)
}
}
static NULL_CALLSITE: NullCallsite = NullCallsite;
static NULL_METADATA: Metadata<'static> = tracing_core::metadata! {
name: "",
target: "",
level: tracing_core::Level::TRACE,
fields: &[],
callsite: &NULL_CALLSITE,
kind: tracing_core::metadata::Kind::SPAN,
};
Self {
metadata: &NULL_METADATA,
parent: None,
ref_count: AtomicUsize::new(0),
extensions: RwLock::new(ExtensionsInner::new()),
}
}
}
impl Clear for DataInner {
fn clear(&mut self) {
if self.parent.is_some() {
let subscriber = dispatcher::get_default(Dispatch::clone);
if let Some(parent) = self.parent.take() {
let _ = subscriber.try_close(parent);
}
}
self.extensions
.get_mut()
.unwrap_or_else(|l| {
l.into_inner()
})
.clear();
}
}
#[cfg(test)]
mod tests {
use super::Registry;
use crate::{layer::Context, registry::LookupSpan, Layer};
use std::{
collections::HashMap,
sync::{Arc, Mutex, Weak},
};
use tracing::{self, subscriber::with_default};
use tracing_core::{
dispatcher,
span::{Attributes, Id},
Subscriber,
};
struct AssertionLayer;
impl<S> Layer<S> for AssertionLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn on_close(&self, id: Id, ctx: Context<'_, S>) {
dbg!(format_args!("closing {:?}", id));
assert!(&ctx.span(&id).is_some());
}
}
#[test]
fn single_layer_can_access_closed_span() {
let subscriber = AssertionLayer.with_subscriber(Registry::default());
with_default(subscriber, || {
let span = tracing::debug_span!("span");
drop(span);
});
}
#[test]
fn multiple_layers_can_access_closed_span() {
let subscriber = AssertionLayer
.and_then(AssertionLayer)
.with_subscriber(Registry::default());
with_default(subscriber, || {
let span = tracing::debug_span!("span");
drop(span);
});
}
struct CloseLayer {
inner: Arc<Mutex<CloseState>>,
}
struct CloseHandle {
state: Arc<Mutex<CloseState>>,
}
#[derive(Default)]
struct CloseState {
open: HashMap<&'static str, Weak<()>>,
closed: Vec<(&'static str, Weak<()>)>,
}
struct SetRemoved(Arc<()>);
impl<S> Layer<S> for CloseLayer
where
S: Subscriber + for<'a> LookupSpan<'a>,
{
fn new_span(&self, _: &Attributes<'_>, id: &Id, ctx: Context<'_, S>) {
let span = ctx.span(id).expect("Missing span; this is a bug");
let mut lock = self.inner.lock().unwrap();
let is_removed = Arc::new(());
assert!(
lock.open
.insert(span.name(), Arc::downgrade(&is_removed))
.is_none(),
"test layer saw multiple spans with the same name, the test is probably messed up"
);
let mut extensions = span.extensions_mut();
extensions.insert(SetRemoved(is_removed));
}
fn on_close(&self, id: Id, ctx: Context<'_, S>) {
let span = if let Some(span) = ctx.span(&id) {
span
} else {
println!(
"span {:?} did not exist in `on_close`, are we panicking?",
id
);
return;
};
let name = span.name();
println!("close {} ({:?})", name, id);
if let Ok(mut lock) = self.inner.lock() {
if let Some(is_removed) = lock.open.remove(name) {
assert!(is_removed.upgrade().is_some());
lock.closed.push((name, is_removed));
}
}
}
}
impl CloseLayer {
fn new() -> (Self, CloseHandle) {
let state = Arc::new(Mutex::new(CloseState::default()));
(
Self {
inner: state.clone(),
},
CloseHandle { state },
)
}
}
impl CloseState {
fn is_open(&self, span: &str) -> bool {
self.open.contains_key(span)
}
fn is_closed(&self, span: &str) -> bool {
self.closed.iter().any(|(name, _)| name == &span)
}
}
impl CloseHandle {
fn assert_closed(&self, span: &str) {
let lock = self.state.lock().unwrap();
assert!(
lock.is_closed(span),
"expected {} to be closed{}",
span,
if lock.is_open(span) {
" (it was still open)"
} else {
", but it never existed (is there a problem with the test?)"
}
)
}
fn assert_open(&self, span: &str) {
let lock = self.state.lock().unwrap();
assert!(
lock.is_open(span),
"expected {} to be open{}",
span,
if lock.is_closed(span) {
" (it was still open)"
} else {
", but it never existed (is there a problem with the test?)"
}
)
}
fn assert_removed(&self, span: &str) {
let lock = self.state.lock().unwrap();
let is_removed = match lock.closed.iter().find(|(name, _)| name == &span) {
Some((_, is_removed)) => is_removed,
None => panic!(
"expected {} to be removed from the registry, but it was not closed {}",
span,
if lock.is_closed(span) {
" (it was still open)"
} else {
", but it never existed (is there a problem with the test?)"
}
),
};
assert!(
is_removed.upgrade().is_none(),
"expected {} to have been removed from the registry",
span
)
}
fn assert_not_removed(&self, span: &str) {
let lock = self.state.lock().unwrap();
let is_removed = match lock.closed.iter().find(|(name, _)| name == &span) {
Some((_, is_removed)) => is_removed,
None if lock.is_open(span) => return,
None => unreachable!(),
};
assert!(
is_removed.upgrade().is_some(),
"expected {} to have been removed from the registry",
span
)
}
#[allow(unused)]
fn assert_last_closed(&self, span: Option<&str>) {
let lock = self.state.lock().unwrap();
let last = lock.closed.last().map(|(span, _)| span);
assert_eq!(
last,
span.as_ref(),
"expected {:?} to have closed last",
span
);
}
fn assert_closed_in_order(&self, order: impl AsRef<[&'static str]>) {
let lock = self.state.lock().unwrap();
let order = order.as_ref();
for (i, name) in order.iter().enumerate() {
assert_eq!(
lock.closed.get(i).map(|(span, _)| span),
Some(name),
"expected close order: {:?}, actual: {:?}",
order,
lock.closed.iter().map(|(name, _)| name).collect::<Vec<_>>()
);
}
}
}
#[test]
fn spans_are_removed_from_registry() {
let (close_layer, state) = CloseLayer::new();
let subscriber = AssertionLayer
.and_then(close_layer)
.with_subscriber(Registry::default());
let dispatch = dispatcher::Dispatch::new(subscriber);
dispatcher::with_default(&dispatch, || {
let span = tracing::debug_span!("span1");
drop(span);
let span = tracing::info_span!("span2");
drop(span);
});
state.assert_removed("span1");
state.assert_removed("span2");
drop(dispatch);
}
#[test]
fn spans_are_only_closed_when_the_last_ref_drops() {
let (close_layer, state) = CloseLayer::new();
let subscriber = AssertionLayer
.and_then(close_layer)
.with_subscriber(Registry::default());
let dispatch = dispatcher::Dispatch::new(subscriber);
let span2 = dispatcher::with_default(&dispatch, || {
let span = tracing::debug_span!("span1");
drop(span);
let span2 = tracing::info_span!("span2");
let span2_clone = span2.clone();
drop(span2);
span2_clone
});
state.assert_removed("span1");
state.assert_not_removed("span2");
drop(span2);
state.assert_removed("span1");
drop(dispatch);
}
#[test]
fn span_enter_guards_are_dropped_out_of_order() {
let (close_layer, state) = CloseLayer::new();
let subscriber = AssertionLayer
.and_then(close_layer)
.with_subscriber(Registry::default());
let dispatch = dispatcher::Dispatch::new(subscriber);
dispatcher::with_default(&dispatch, || {
let span1 = tracing::debug_span!("span1");
let span2 = tracing::info_span!("span2");
let enter1 = span1.enter();
let enter2 = span2.enter();
drop(enter1);
drop(span1);
state.assert_removed("span1");
state.assert_not_removed("span2");
drop(enter2);
state.assert_not_removed("span2");
drop(span2);
state.assert_removed("span1");
state.assert_removed("span2");
});
}
#[test]
fn child_closes_parent() {
let (close_layer, state) = CloseLayer::new();
let subscriber = close_layer.with_subscriber(Registry::default());
let dispatch = dispatcher::Dispatch::new(subscriber);
dispatcher::with_default(&dispatch, || {
let span1 = tracing::info_span!("parent");
let span2 = tracing::info_span!(parent: &span1, "child");
state.assert_open("parent");
state.assert_open("child");
drop(span1);
state.assert_open("parent");
state.assert_open("child");
drop(span2);
state.assert_closed("parent");
state.assert_closed("child");
});
}
#[test]
fn child_closes_grandparent() {
let (close_layer, state) = CloseLayer::new();
let subscriber = close_layer.with_subscriber(Registry::default());
let dispatch = dispatcher::Dispatch::new(subscriber);
dispatcher::with_default(&dispatch, || {
let span1 = tracing::info_span!("grandparent");
let span2 = tracing::info_span!(parent: &span1, "parent");
let span3 = tracing::info_span!(parent: &span2, "child");
state.assert_open("grandparent");
state.assert_open("parent");
state.assert_open("child");
drop(span1);
drop(span2);
state.assert_open("grandparent");
state.assert_open("parent");
state.assert_open("child");
drop(span3);
state.assert_closed_in_order(&["child", "parent", "grandparent"]);
});
}
}