metrics/recorder/
cell.rs

1use super::{Recorder, SetRecorderError};
2use std::{
3    cell::UnsafeCell,
4    sync::atomic::{AtomicUsize, Ordering},
5};
6
7/// The recorder is uninitialized.
8const UNINITIALIZED: usize = 0;
9
10/// The recorder is currently being initialized.
11const INITIALIZING: usize = 1;
12
13/// The recorder has been initialized successfully and can be read.
14const INITIALIZED: usize = 2;
15
16/// An specialized version of `OnceCell` for `Recorder`.
17pub struct RecorderOnceCell {
18    recorder: UnsafeCell<Option<&'static dyn Recorder>>,
19    state: AtomicUsize,
20}
21
22impl RecorderOnceCell {
23    /// Creates an uninitialized `RecorderOnceCell`.
24    pub const fn new() -> Self {
25        Self { recorder: UnsafeCell::new(None), state: AtomicUsize::new(UNINITIALIZED) }
26    }
27
28    pub fn set<R>(&self, recorder: R) -> Result<(), SetRecorderError<R>>
29    where
30        R: Recorder + 'static,
31    {
32        // Try and transition the cell from `UNINITIALIZED` to `INITIALIZING`, which would give
33        // us exclusive access to set the recorder.
34        match self.state.compare_exchange(
35            UNINITIALIZED,
36            INITIALIZING,
37            Ordering::Acquire,
38            Ordering::Relaxed,
39        ) {
40            Ok(UNINITIALIZED) => {
41                unsafe {
42                    // SAFETY: Access is unique because we can only be here if we won the race
43                    // to transition from `UNINITIALIZED` to `INITIALIZING` above.
44                    self.recorder.get().write(Some(Box::leak(Box::new(recorder))));
45                }
46
47                // Mark the recorder as initialized, which will make it visible to readers.
48                self.state.store(INITIALIZED, Ordering::Release);
49                Ok(())
50            }
51            _ => Err(SetRecorderError(recorder)),
52        }
53    }
54
55    pub fn try_load(&self) -> Option<&'static dyn Recorder> {
56        if self.state.load(Ordering::Acquire) != INITIALIZED {
57            None
58        } else {
59            // SAFETY: If the state is `INITIALIZED`, then we know that the recorder has been
60            // installed and is safe to read.
61            unsafe { self.recorder.get().read() }
62        }
63    }
64}
65
66// SAFETY: We can only mutate through `set`, which is protected by the `state` and unsafe
67// function where the caller has to guarantee synced-ness.
68unsafe impl Send for RecorderOnceCell {}
69unsafe impl Sync for RecorderOnceCell {}