use crate::{
callsite, span,
subscriber::{self, Subscriber},
Event, Metadata,
};
use std::{
any::Any,
cell::{Cell, RefCell},
error, fmt,
sync::{
atomic::{AtomicUsize, Ordering},
Arc, Weak,
},
};
#[derive(Clone)]
pub struct Dispatch {
subscriber: Arc<dyn Subscriber + Send + Sync>,
}
thread_local! {
static CURRENT_STATE: State = State {
default: RefCell::new(Dispatch::none()),
can_enter: Cell::new(true),
};
}
static GLOBAL_INIT: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;
static mut GLOBAL_DISPATCH: Option<Dispatch> = None;
struct State {
default: RefCell<Dispatch>,
can_enter: Cell<bool>,
}
struct ResetGuard(Option<Dispatch>);
pub fn with_default<T>(dispatcher: &Dispatch, f: impl FnOnce() -> T) -> T {
let _guard = State::set_default(dispatcher.clone());
f()
}
pub fn set_global_default(dispatcher: Dispatch) -> Result<(), SetGlobalDefaultError> {
if GLOBAL_INIT.compare_and_swap(UNINITIALIZED, INITIALIZING, Ordering::SeqCst) == UNINITIALIZED
{
unsafe {
GLOBAL_DISPATCH = Some(dispatcher.clone());
}
GLOBAL_INIT.store(INITIALIZED, Ordering::SeqCst);
Ok(())
} else {
Err(SetGlobalDefaultError { _no_construct: () })
}
}
#[derive(Debug)]
pub struct SetGlobalDefaultError {
_no_construct: (),
}
impl fmt::Display for SetGlobalDefaultError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("a global default trace dispatcher has already been set")
}
}
impl error::Error for SetGlobalDefaultError {}
pub fn get_default<T, F>(mut f: F) -> T
where
F: FnMut(&Dispatch) -> T,
{
struct Entered<'a>(&'a Cell<bool>);
impl<'a> Drop for Entered<'a> {
#[inline]
fn drop(&mut self) {
self.0.set(true);
}
}
CURRENT_STATE
.try_with(|state| {
if state.can_enter.replace(false) {
let _guard = Entered(&state.can_enter);
let mut default = state.default.borrow_mut();
if default.is::<NoSubscriber>() && GLOBAL_INIT.load(Ordering::SeqCst) == INITIALIZED
{
unsafe {
*default = GLOBAL_DISPATCH
.as_ref()
.expect("invariant violated: GLOBAL_DISPATCH must be initialized before GLOBAL_INIT is set")
.clone()
}
}
f(&*default)
} else {
f(&Dispatch::none())
}
})
.unwrap_or_else(|_| f(&Dispatch::none()))
}
pub(crate) struct Registrar(Weak<dyn Subscriber + Send + Sync>);
impl Dispatch {
#[inline]
pub fn none() -> Self {
Dispatch {
subscriber: Arc::new(NoSubscriber),
}
}
pub fn new<S>(subscriber: S) -> Self
where
S: Subscriber + Send + Sync + 'static,
{
let me = Dispatch {
subscriber: Arc::new(subscriber),
};
callsite::register_dispatch(&me);
me
}
pub(crate) fn registrar(&self) -> Registrar {
Registrar(Arc::downgrade(&self.subscriber))
}
#[inline]
pub fn register_callsite(&self, metadata: &'static Metadata<'static>) -> subscriber::Interest {
self.subscriber.register_callsite(metadata)
}
#[inline]
pub fn new_span(&self, span: &span::Attributes<'_>) -> span::Id {
self.subscriber.new_span(span)
}
#[inline]
pub fn record(&self, span: &span::Id, values: &span::Record<'_>) {
self.subscriber.record(span, values)
}
#[inline]
pub fn record_follows_from(&self, span: &span::Id, follows: &span::Id) {
self.subscriber.record_follows_from(span, follows)
}
#[inline]
pub fn enabled(&self, metadata: &Metadata<'_>) -> bool {
self.subscriber.enabled(metadata)
}
#[inline]
pub fn event(&self, event: &Event<'_>) {
self.subscriber.event(event)
}
#[inline]
pub fn enter(&self, span: &span::Id) {
self.subscriber.enter(span);
}
#[inline]
pub fn exit(&self, span: &span::Id) {
self.subscriber.exit(span);
}
#[inline]
pub fn clone_span(&self, id: &span::Id) -> span::Id {
self.subscriber.clone_span(&id)
}
#[inline]
#[deprecated(since = "0.1.2", note = "use `Dispatch::try_close` instead")]
pub fn drop_span(&self, id: span::Id) {
#[allow(deprecated)]
self.subscriber.drop_span(id);
}
#[inline]
pub fn try_close(&self, id: span::Id) -> bool {
self.subscriber.try_close(id)
}
#[inline]
pub fn current_span(&self) -> span::Current {
self.subscriber.current_span()
}
#[inline]
pub fn is<T: Any>(&self) -> bool {
Subscriber::is::<T>(&*self.subscriber)
}
#[inline]
pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
Subscriber::downcast_ref(&*self.subscriber)
}
}
impl fmt::Debug for Dispatch {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad("Dispatch(...)")
}
}
impl<S> From<S> for Dispatch
where
S: Subscriber + Send + Sync + 'static,
{
#[inline]
fn from(subscriber: S) -> Self {
Dispatch::new(subscriber)
}
}
struct NoSubscriber;
impl Subscriber for NoSubscriber {
#[inline]
fn register_callsite(&self, _: &'static Metadata<'static>) -> subscriber::Interest {
subscriber::Interest::never()
}
fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
span::Id::from_u64(0xDEAD)
}
fn event(&self, _event: &Event<'_>) {}
fn record(&self, _span: &span::Id, _values: &span::Record<'_>) {}
fn record_follows_from(&self, _span: &span::Id, _follows: &span::Id) {}
#[inline]
fn enabled(&self, _metadata: &Metadata<'_>) -> bool {
false
}
fn enter(&self, _span: &span::Id) {}
fn exit(&self, _span: &span::Id) {}
}
impl Registrar {
pub(crate) fn try_register(
&self,
metadata: &'static Metadata<'static>,
) -> Option<subscriber::Interest> {
self.0.upgrade().map(|s| s.register_callsite(metadata))
}
pub(crate) fn is_alive(&self) -> bool {
self.0.upgrade().is_some()
}
}
impl State {
#[inline]
fn set_default(new_dispatch: Dispatch) -> ResetGuard {
let prior = CURRENT_STATE
.try_with(|state| {
state.can_enter.set(true);
state.default.replace(new_dispatch)
})
.ok();
ResetGuard(prior)
}
}
impl Drop for ResetGuard {
#[inline]
fn drop(&mut self) {
if let Some(dispatch) = self.0.take() {
let _ = CURRENT_STATE.try_with(|state| {
*state.default.borrow_mut() = dispatch;
});
}
}
}
#[cfg(test)]
mod test {
use super::*;
use crate::{
callsite::Callsite,
metadata::{Kind, Level, Metadata},
span,
subscriber::{Interest, Subscriber},
Event,
};
use std::sync::atomic::{AtomicUsize, Ordering};
#[test]
fn dispatch_is() {
let dispatcher = Dispatch::new(NoSubscriber);
assert!(dispatcher.is::<NoSubscriber>());
}
#[test]
fn dispatch_downcasts() {
let dispatcher = Dispatch::new(NoSubscriber);
assert!(dispatcher.downcast_ref::<NoSubscriber>().is_some());
}
struct TestCallsite;
static TEST_CALLSITE: TestCallsite = TestCallsite;
static TEST_META: Metadata<'static> = metadata! {
name: "test",
target: module_path!(),
level: Level::DEBUG,
fields: &[],
callsite: &TEST_CALLSITE,
kind: Kind::EVENT
};
impl Callsite for TestCallsite {
fn set_interest(&self, _: Interest) {}
fn metadata(&self) -> &Metadata<'_> {
&TEST_META
}
}
#[test]
fn events_dont_infinite_loop() {
struct TestSubscriber;
impl Subscriber for TestSubscriber {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
span::Id::from_u64(0xAAAA)
}
fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
fn event(&self, _: &Event<'_>) {
static EVENTS: AtomicUsize = AtomicUsize::new(0);
assert_eq!(
EVENTS.fetch_add(1, Ordering::Relaxed),
0,
"event method called twice!"
);
Event::dispatch(&TEST_META, &TEST_META.fields().value_set(&[]))
}
fn enter(&self, _: &span::Id) {}
fn exit(&self, _: &span::Id) {}
}
with_default(&Dispatch::new(TestSubscriber), || {
Event::dispatch(&TEST_META, &TEST_META.fields().value_set(&[]))
})
}
#[test]
fn spans_dont_infinite_loop() {
fn mk_span() {
get_default(|current| {
current.new_span(&span::Attributes::new(
&TEST_META,
&TEST_META.fields().value_set(&[]),
))
});
}
struct TestSubscriber;
impl Subscriber for TestSubscriber {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
static NEW_SPANS: AtomicUsize = AtomicUsize::new(0);
assert_eq!(
NEW_SPANS.fetch_add(1, Ordering::Relaxed),
0,
"new_span method called twice!"
);
mk_span();
span::Id::from_u64(0xAAAA)
}
fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
fn event(&self, _: &Event<'_>) {}
fn enter(&self, _: &span::Id) {}
fn exit(&self, _: &span::Id) {}
}
with_default(&Dispatch::new(TestSubscriber), || mk_span())
}
#[test]
fn global_dispatch() {
struct TestSubscriberA;
impl Subscriber for TestSubscriberA {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
span::Id::from_u64(1)
}
fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
fn event(&self, _: &Event<'_>) {}
fn enter(&self, _: &span::Id) {}
fn exit(&self, _: &span::Id) {}
}
struct TestSubscriberB;
impl Subscriber for TestSubscriberB {
fn enabled(&self, _: &Metadata<'_>) -> bool {
true
}
fn new_span(&self, _: &span::Attributes<'_>) -> span::Id {
span::Id::from_u64(1)
}
fn record(&self, _: &span::Id, _: &span::Record<'_>) {}
fn record_follows_from(&self, _: &span::Id, _: &span::Id) {}
fn event(&self, _: &Event<'_>) {}
fn enter(&self, _: &span::Id) {}
fn exit(&self, _: &span::Id) {}
}
set_global_default(Dispatch::new(TestSubscriberA)).expect("global dispatch set failed");
get_default(|current| {
assert!(
current.is::<TestSubscriberA>(),
"global dispatch get failed"
)
});
with_default(&Dispatch::new(TestSubscriberB), || {
get_default(|current| {
assert!(
current.is::<TestSubscriberB>(),
"thread-local override of global dispatch failed"
)
});
});
get_default(|current| {
assert!(
current.is::<TestSubscriberA>(),
"reset to global override failed"
)
});
set_global_default(Dispatch::new(TestSubscriberA))
.expect_err("double global dispatch set succeeded");
}
}