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, RefCell, RefMut},
error,
};
#[cfg(feature = "alloc")]
use alloc::sync::{Arc, Weak};
#[cfg(feature = "alloc")]
use core::ops::Deref;
#[derive(Clone)]
pub struct Dispatch {
subscriber: Arc<dyn Subscriber + Send + Sync>,
}
#[derive(Clone)]
pub struct WeakDispatch {
subscriber: Weak<dyn Subscriber + Send + Sync>,
}
#[cfg(feature = "alloc")]
#[derive(Clone)]
enum Kind<T> {
Global(&'static (dyn Collect + 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);
const UNINITIALIZED: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;
static mut GLOBAL_DISPATCH: Option<Dispatch> = None;
#[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()
{
unsafe {
GLOBAL_DISPATCH = Some(dispatcher);
}
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,
{
CURRENT_STATE
.try_with(|state| {
if let Some(entered) = state.enter() {
return f(&*entered.current());
}
f(&Dispatch::none())
})
.unwrap_or_else(|_| f(&Dispatch::none()))
}
#[cfg(feature = "std")]
#[doc(hidden)]
#[inline(never)]
pub fn get_current<T>(f: impl FnOnce(&Dispatch) -> T) -> Option<T> {
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> {
let dispatch = get_global()?;
Some(f(&dispatch))
}
#[cfg(not(feature = "std"))]
pub fn get_default<T, F>(mut f: F) -> T
where
F: FnMut(&Dispatch) -> T,
{
if let Some(d) = get_global() {
f(d)
} else {
f(&Dispatch::none())
}
}
fn get_global() -> Option<&'static Dispatch> {
if GLOBAL_INIT.load(Ordering::SeqCst) != INITIALIZED {
return None;
}
unsafe {
Some(GLOBAL_DISPATCH.as_ref().expect(
"invariant violated: GLOBAL_DISPATCH must be initialized before GLOBAL_INIT is set",
))
}
}
#[cfg(feature = "std")]
pub(crate) struct Registrar(Weak<dyn Subscriber + Send + Sync>);
impl Dispatch {
#[inline]
pub fn none() -> Self {
Dispatch {
subscriber: Arc::new(NoSubscriber::default()),
}
}
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
}
#[cfg(feature = "std")]
pub(crate) fn registrar(&self) -> Registrar {
Registrar(Arc::downgrade(&self.subscriber))
}
pub fn downgrade(&self) -> WeakDispatch {
WeakDispatch {
subscriber: Arc::downgrade(&self.subscriber),
}
}
#[inline(always)]
#[cfg(not(feature = "alloc"))]
pub(crate) fn subscriber(&self) -> &(dyn Subscriber + Send + Sync) {
&self.subscriber
}
#[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<'_>) {
if self.subscriber.event_enabled(event) {
self.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 {
f.debug_tuple("Dispatch")
.field(&format_args!("{:p}", self.subscriber))
.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 {
let mut tuple = f.debug_tuple("WeakDispatch");
match self.subscriber.upgrade() {
Some(subscriber) => tuple.field(&format_args!("Some({:p})", subscriber)),
None => tuple.field(&format_args!("None")),
};
tuple.finish()
}
}
#[cfg(feature = "std")]
impl Registrar {
pub(crate) fn upgrade(&self) -> Option<Dispatch> {
self.0.upgrade().map(|subscriber| Dispatch { subscriber })
}
}
#[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);
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) -> RefMut<'a, Dispatch> {
let default = self.0.default.borrow_mut();
RefMut::map(default, |default| {
default.get_or_insert_with(|| get_global().cloned().unwrap_or_else(Dispatch::none))
})
}
}
#[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()));
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>());
}
}