1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166

use std::num::NonZeroU64;
use std::sync::atomic::{Ordering, AtomicU64, AtomicPtr};

use once_cell::sync::Lazy;

/// Store the id to be assigned to the next call site created.
static NEXT_CALL_SITE_ID: AtomicU64 = AtomicU64::new(1);
/// Store the global registry of call sites
static REGISTRY: Lazy<Registry> = Lazy::new(|| {
    Registry {
        head: AtomicPtr::new(std::ptr::null_mut()),
    }
});

/// Unique identifier of a [`CallSite`], attributed the first time the call site
/// is entered.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct CallSiteId(NonZeroU64);

impl CallSiteId {
    pub(crate) fn new(value: u64) -> CallSiteId {
        CallSiteId(std::num::NonZeroU64::new(value).expect("got a zero value for span id"))
    }
}

/// A [`CallSite`] identify uniquely a location in the source code, and record
/// multiple attributes associated with this location.
///
/// The only way to create a [`CallSite`] is with the [`macro@callsite`] macro,
/// which also takes care of registering the call site globally.
pub struct CallSite {
    /// Unique identifier of this call site
    id: CallSiteId,
    /// The name of the call site
    name: &'static str,
    /// The name of the Rust module where the call site occurred
    module_path: &'static str,
    /// The name of the source code file where the call site occurred
    file: &'static str,
    /// The line number in the source code file where the call site occurred
    line: u32,
    /// Call sites are registered using an atomic, append only intrusive linked
    /// list. If more than one call site are registered, this will be set to the
    /// last registered call site.
    next: AtomicPtr<CallSite>,
}

impl CallSite {
    /// Create a new `CallSite` with the given metadata. This function is
    /// private to this crate, and is only marked `pub` to be able to call it
    /// from inside macros.
    #[doc(hidden)]
    pub fn new(name: &'static str, module_path: &'static str, file: &'static str, line: u32) -> CallSite {
        let id = CallSiteId::new(NEXT_CALL_SITE_ID.fetch_add(1, Ordering::SeqCst));
        let next = AtomicPtr::new(std::ptr::null_mut());
        CallSite { id, name, module_path, file, line, next }
    }

    pub(crate) fn id(&self) -> CallSiteId {
        self.id
    }

    /// Get the user-provided name for this call site
    pub fn name(&self) -> &str {
        self.name
    }

    /// Get the rust module path to the source code location of this call site
    pub fn module_path(&self) -> &str {
        self.module_path
    }

    /// Get the path to the file containing this call site
    pub fn file(&self) -> &str {
        self.file
    }

    /// Get the line of the source file containing this call site
    pub fn line(&self) -> u32 {
        self.line
    }

    /// Get the full name of this call site, containing both the name and the
    /// module path.
    pub fn full_name(&self) -> String {
        let mut name = self.module_path.to_owned();
        name += "::";

        if self.name.contains(' ') {
            name += "{";
            name += self.name;
            name += "}";
        } else {
            name += self.name;
        }

        return name;
    }
}

/// Registry of CallSite, as the head pointer of an atomic, append-only linked
/// list.
struct Registry {
    head: AtomicPtr<CallSite>,
}

impl Registry {
    /// Register a new callsite within the list
    fn register(&self, callsite: &'static CallSite) {
        let mut head = self.head.load(Ordering::Acquire);

        loop {
            callsite.next.store(head, Ordering::Release);

            assert_ne!(
                callsite as *const _, head,
                "Attempted to register a `Callsite` that already exists! \
                This will cause an infinite loop when attempting to read from the \
                callsite registry."
            );

            match self.head.compare_exchange(
                head,
                callsite as *const _ as *mut _,
                Ordering::AcqRel,
                Ordering::Acquire,
            ) {
                Ok(_) => {
                    break;
                }
                Err(current) => head = current,
            }
        }
    }

    /// Execute the provided function on all elements of the list
    fn for_each(&self, mut f: impl FnMut(&'static CallSite)) {
        let mut head = self.head.load(Ordering::Acquire);

        while let Some(registered) = unsafe { head.as_ref() } {
            f(registered);
            head = registered.next.load(Ordering::Acquire);
        }
    }
}

/// Register a call site. This function is a private function of this crate. It
/// is only marked `pub` to be able to call it from inside macros.
#[doc(hidden)]
pub fn register_callsite(callsite: &'static CallSite) {
    REGISTRY.register(callsite);
}

/// Execute the given function on all call sites we know about.
///
/// # Examples
/// ```
/// # use time_graph::traverse_registered_callsite;
///
/// traverse_registered_callsite(|callsite| {
///     println!("got a callsite at {}:{}", callsite.file(), callsite.line());
/// })
/// ```
pub fn traverse_registered_callsite(function: impl FnMut(&'static CallSite)) {
    REGISTRY.for_each(function);
}