use crate::{
callsite, span,
subscriber::{self, NoSubscriber, Subscriber},
Event, LevelFilter, Metadata,
};
use crate::stdlib::{
any::Any,
fmt,
sync::{
atomic::{AtomicBool, AtomicUsize, Ordering},
Arc, Weak,
},
};
#[cfg(feature = "std")]
use crate::stdlib::{
cell::{Cell, Ref, RefCell},
error,
};
#[cfg(feature = "alloc")]
use alloc::sync::{Arc, Weak};
#[cfg(feature = "alloc")]
use core::ops::Deref;
#[derive(Clone)]
pub struct Dispatch {
subscriber: Kind<Arc<dyn Subscriber + Send + Sync>>,
}
#[derive(Clone)]
pub struct WeakDispatch {
subscriber: Kind<Weak<dyn Subscriber + Send + Sync>>,
}
#[derive(Clone)]
enum Kind<T> {
Global(&'static (dyn Subscriber + Send + Sync)),
Scoped(T),
}
#[cfg(feature = "std")]
thread_local! {
static CURRENT_STATE: State = State {
default: RefCell::new(None),
can_enter: Cell::new(true),
};
}
static EXISTS: AtomicBool = AtomicBool::new(false);
static GLOBAL_INIT: AtomicUsize = AtomicUsize::new(UNINITIALIZED);
#[cfg(feature = "std")]
static SCOPED_COUNT: AtomicUsize = AtomicUsize::new(0);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;
static mut GLOBAL_DISPATCH: Dispatch = Dispatch {
subscriber: Kind::Global(&NO_SUBSCRIBER),
};
static NONE: Dispatch = Dispatch {
subscriber: Kind::Global(&NO_SUBSCRIBER),
};
static NO_SUBSCRIBER: NoSubscriber = NoSubscriber::new();
#[cfg(feature = "std")]
struct State {
default: RefCell<Option<Dispatch>>,
can_enter: Cell<bool>,
}
#[cfg(feature = "std")]
struct Entered<'a>(&'a State);
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[derive(Debug)]
pub struct DefaultGuard(Option<Dispatch>);
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
pub fn with_default<T>(dispatcher: &Dispatch, f: impl FnOnce() -> T) -> T {
let _guard = set_default(dispatcher);
f()
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
#[must_use = "Dropping the guard unregisters the dispatcher."]
pub fn set_default(dispatcher: &Dispatch) -> DefaultGuard {
State::set_default(dispatcher.clone())
}
pub fn set_global_default(dispatcher: Dispatch) -> Result<(), SetGlobalDefaultError> {
if GLOBAL_INIT
.compare_exchange(
UNINITIALIZED,
INITIALIZING,
Ordering::SeqCst,
Ordering::SeqCst,
)
.is_ok()
{
let subscriber = {
let subscriber = match dispatcher.subscriber {
Kind::Global(s) => s,
Kind::Scoped(s) => unsafe {
&*Arc::into_raw(s)
},
};
Kind::Global(subscriber)
};
unsafe {
GLOBAL_DISPATCH = Dispatch { subscriber };
}
GLOBAL_INIT.store(INITIALIZED, Ordering::SeqCst);
EXISTS.store(true, Ordering::Release);
Ok(())
} else {
Err(SetGlobalDefaultError { _no_construct: () })
}
}
#[doc(hidden)]
#[inline(always)]
pub fn has_been_set() -> bool {
EXISTS.load(Ordering::Relaxed)
}
pub struct SetGlobalDefaultError {
_no_construct: (),
}
impl fmt::Debug for SetGlobalDefaultError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_tuple("SetGlobalDefaultError")
.field(&Self::MESSAGE)
.finish()
}
}
impl fmt::Display for SetGlobalDefaultError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.pad(Self::MESSAGE)
}
}
#[cfg(feature = "std")]
#[cfg_attr(docsrs, doc(cfg(feature = "std")))]
impl error::Error for SetGlobalDefaultError {}
impl SetGlobalDefaultError {
const MESSAGE: &'static str = "a global default trace dispatcher has already been set";
}
#[cfg(feature = "std")]
pub fn get_default<T, F>(mut f: F) -> T
where
F: FnMut(&Dispatch) -> T,
{
if SCOPED_COUNT.load(Ordering::Acquire) == 0 {
return f(get_global());
}
CURRENT_STATE
.try_with(|state| {
if let Some(entered) = state.enter() {
return f(&entered.current());
}
f(&NONE)
})
.unwrap_or_else(|_| f(&NONE))
}
#[cfg(feature = "std")]
#[doc(hidden)]
#[inline(never)]
pub fn get_current<T>(f: impl FnOnce(&Dispatch) -> T) -> Option<T> {
if SCOPED_COUNT.load(Ordering::Acquire) == 0 {
return Some(f(get_global()));
}
CURRENT_STATE
.try_with(|state| {
let entered = state.enter()?;
Some(f(&entered.current()))
})
.ok()?
}
#[cfg(not(feature = "std"))]
#[doc(hidden)]
pub fn get_current<T>(f: impl FnOnce(&Dispatch) -> T) -> Option<T> {
Some(f(get_global()))
}
#[cfg(not(feature = "std"))]
pub fn get_default<T, F>(mut f: F) -> T
where
F: FnMut(&Dispatch) -> T,
{
f(&get_global())
}
#[inline]
fn get_global() -> &'static Dispatch {
if GLOBAL_INIT.load(Ordering::SeqCst) != INITIALIZED {
return &NONE;
}
unsafe {
&GLOBAL_DISPATCH
}
}
#[cfg(feature = "std")]
pub(crate) struct Registrar(Kind<Weak<dyn Subscriber + Send + Sync>>);
impl Dispatch {
#[inline]
pub fn none() -> Self {
Dispatch {
subscriber: Kind::Global(&NO_SUBSCRIBER),
}
}
pub fn new<S>(subscriber: S) -> Self
where
S: Subscriber + Send + Sync + 'static,
{
let me = Dispatch {
subscriber: Kind::Scoped(Arc::new(subscriber)),
};
callsite::register_dispatch(&me);
me
}
#[cfg(feature = "std")]
pub(crate) fn registrar(&self) -> Registrar {
Registrar(self.subscriber.downgrade())
}
pub fn downgrade(&self) -> WeakDispatch {
WeakDispatch {
subscriber: self.subscriber.downgrade(),
}
}
#[inline(always)]
pub(crate) fn subscriber(&self) -> &(dyn Subscriber + Send + Sync) {
match self.subscriber {
Kind::Global(s) => s,
Kind::Scoped(ref s) => s.as_ref(),
}
}
#[inline]
pub fn register_callsite(&self, metadata: &'static Metadata<'static>) -> subscriber::Interest {
self.subscriber().register_callsite(metadata)
}
#[inline]
pub(crate) fn max_level_hint(&self) -> Option<LevelFilter> {
self.subscriber().max_level_hint()
}
#[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<'_>) {
let subscriber = self.subscriber();
if subscriber.event_enabled(event) {
subscriber.event(event);
}
}
pub fn enter(&self, span: &span::Id) {
self.subscriber().enter(span);
}
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);
}
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 {
<dyn Subscriber>::is::<T>(self.subscriber())
}
#[inline]
pub fn downcast_ref<T: Any>(&self) -> Option<&T> {
<dyn Subscriber>::downcast_ref(self.subscriber())
}
}
impl Default for Dispatch {
fn default() -> Self {
get_default(|default| default.clone())
}
}
impl fmt::Debug for Dispatch {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.subscriber {
Kind::Scoped(ref s) => f
.debug_tuple("Dispatch::Scoped")
.field(&format_args!("{:p}", s))
.finish(),
Kind::Global(s) => f
.debug_tuple("Dispatch::Global")
.field(&format_args!("{:p}", s))
.finish(),
}
}
}
impl<S> From<S> for Dispatch
where
S: Subscriber + Send + Sync + 'static,
{
#[inline]
fn from(subscriber: S) -> Self {
Dispatch::new(subscriber)
}
}
impl WeakDispatch {
pub fn upgrade(&self) -> Option<Dispatch> {
self.subscriber
.upgrade()
.map(|subscriber| Dispatch { subscriber })
}
}
impl fmt::Debug for WeakDispatch {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.subscriber {
Kind::Scoped(ref s) => f
.debug_tuple("WeakDispatch::Scoped")
.field(&format_args!("{:p}", s))
.finish(),
Kind::Global(s) => f
.debug_tuple("WeakDispatch::Global")
.field(&format_args!("{:p}", s))
.finish(),
}
}
}
#[cfg(feature = "std")]
impl Registrar {
pub(crate) fn upgrade(&self) -> Option<Dispatch> {
self.0.upgrade().map(|subscriber| Dispatch { subscriber })
}
}
impl Kind<Arc<dyn Subscriber + Send + Sync>> {
fn downgrade(&self) -> Kind<Weak<dyn Subscriber + Send + Sync>> {
match self {
Kind::Global(s) => Kind::Global(*s),
Kind::Scoped(ref s) => Kind::Scoped(Arc::downgrade(s)),
}
}
}
impl Kind<Weak<dyn Subscriber + Send + Sync>> {
fn upgrade(&self) -> Option<Kind<Arc<dyn Subscriber + Send + Sync>>> {
match self {
Kind::Global(s) => Some(Kind::Global(*s)),
Kind::Scoped(ref s) => Some(Kind::Scoped(s.upgrade()?)),
}
}
}
#[cfg(feature = "std")]
impl State {
#[inline]
fn set_default(new_dispatch: Dispatch) -> DefaultGuard {
let prior = CURRENT_STATE
.try_with(|state| {
state.can_enter.set(true);
state.default.replace(Some(new_dispatch))
})
.ok()
.flatten();
EXISTS.store(true, Ordering::Release);
SCOPED_COUNT.fetch_add(1, Ordering::Release);
DefaultGuard(prior)
}
#[inline]
fn enter(&self) -> Option<Entered<'_>> {
if self.can_enter.replace(false) {
Some(Entered(self))
} else {
None
}
}
}
#[cfg(feature = "std")]
impl<'a> Entered<'a> {
#[inline]
fn current(&self) -> Ref<'a, Dispatch> {
let default = self.0.default.borrow();
Ref::map(default, |default| match default {
Some(default) => default,
None => get_global(),
})
}
}
#[cfg(feature = "std")]
impl<'a> Drop for Entered<'a> {
#[inline]
fn drop(&mut self) {
self.0.can_enter.set(true);
}
}
#[cfg(feature = "std")]
impl Drop for DefaultGuard {
#[inline]
fn drop(&mut self) {
let prev = CURRENT_STATE.try_with(|state| state.default.replace(self.0.take()));
SCOPED_COUNT.fetch_sub(1, Ordering::Release);
drop(prev)
}
}
#[cfg(test)]
mod test {
use super::*;
#[cfg(feature = "std")]
use crate::stdlib::sync::atomic::{AtomicUsize, Ordering};
use crate::{
callsite::Callsite,
metadata::{Kind, Level, Metadata},
subscriber::Interest,
};
#[test]
fn dispatch_is() {
let dispatcher = Dispatch::new(NoSubscriber::default());
assert!(dispatcher.is::<NoSubscriber>());
}
#[test]
fn dispatch_downcasts() {
let dispatcher = Dispatch::new(NoSubscriber::default());
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]
#[cfg(feature = "std")]
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]
#[cfg(feature = "std")]
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 default_no_subscriber() {
let default_dispatcher = Dispatch::default();
assert!(default_dispatcher.is::<NoSubscriber>());
}
#[cfg(feature = "std")]
#[test]
fn default_dispatch() {
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<'_>) {}
fn enter(&self, _: &span::Id) {}
fn exit(&self, _: &span::Id) {}
}
let guard = set_default(&Dispatch::new(TestSubscriber));
let default_dispatcher = Dispatch::default();
assert!(default_dispatcher.is::<TestSubscriber>());
drop(guard);
let default_dispatcher = Dispatch::default();
assert!(default_dispatcher.is::<NoSubscriber>());
}
}