pointer.rs

/// Atomic smart pointer that supports replacing itself.
///
/// ## Metric handles and expiration
///
/// In metric handles, we have to deal with the potential that a registry expires a metric, which
/// could otherwise leave a handle attached to a metric that is reported nowhere. In order to handle
/// this, we need the ability to be able to reset/update the internal field that holds the smart
/// pointer reference to the inner data, which requires a mechanism for interior mutability given
/// the immutable references by which handles are accessed through.
///
/// ## Existing smart pointer types
///
/// In nearly all cases, the existing `Arc<T>` from the standard library works for atomically
/// sharing immutable access to some piece of wrapped data. It also additionally supports the trick
/// of being able to be transparently coerced to `Arc<dyn Trait>` from `Arc<T>` if `T: Trait`, which
/// we utilize to allow erasing the concrete type so that the handle types themselves are not
/// required to carry any generic type.
///
/// All of this culiminates in a few requirements:
///
/// - we want to be able to allow callers to use `Arc<T>` where possible because it's already
///   commonly used
/// - we need to be able to support dynamically-sized types for the handle's inner smart pointer
/// - we need a way to expose interior mutability without affecting performance in the normal path
///   (metric not expired)
///
/// `Pointer<T>` provides a mechanism to do this that interoperates with `Arc<T>`.
///
/// ## Design
///
/// Internally, `Pointer<T>` simply manages an atomic pointer to `Arc<T>` (or `Weak<T>`, more on
/// that later) through a newtype wrapper. This is important, as it's a key part of the high
/// performance design. Normally, keeping an atomic pointer directly to `Arc<T>`, by converting it
/// to a raw pointer via `Arc::into_raw`, would only be possible if both of these constraints could
/// be met:
///
/// - ability to destructure a "fat" pointer into its constituent data pointer and virtual table
///   pointer
/// - support for 128-bit atomics (for storing the data pointer + virtual table pointer together)
///
/// While support for 128-bit atomics is generally available, it is not current possible on stable
/// Rust to destructure a fat pointer to even be able to atomically store it. We store `Arc<T>`
/// and `Weak<T>`, even when `T` is unsized, in a newtype wrapper that then provides us a concrete
/// type to reference when creating and storing the raw pointer to those allocations.
///
/// We are acheiving this via double indirection -- first the newtype wrapper, and then
/// `Arc<T>`/`Weak<T>` itself -- which isn't optimal, but crucially, we can load the pointer to the
/// wrapper allocation in a lock-free fashion.
///
/// ## Replacement of weak pointers
///
/// Additionally, we support updating weak pointers in order to reattach them to a live strong reference.
///
/// Like the paradigm of `Arc<T>` and `Weak<T>` being "strong" and "weak" references, `Pointer<T>`
/// allows for the same paradigm. However, as mentioned prior, we have to contend with the potential
/// of a registry expiring a metric, which is the equivalent of all strong references dropping.
///
/// In order to deal with this, all dereference calls (through `with_ref`) must pass a closure which
/// can provide a new `Pointer<T>`, which is meant to act as a replacement. If, upon attempting to
/// dereference a weak pointer, we find that all strong references are gone, we'll replace our own
/// internal pointer state with the state of the replacement pointer.
///
/// This allows us to avoid having to require callers to provide interior mutability around
/// `Pointer<T>` itself in order to reattach a metric handle to the current recorder.
pub struct Pointer<T: ?Sized> {
    ptr: AtomicPtr<()>,
    update: Mutex<()>,
    _ty: PhantomData<T>,
}

#[repr(align(2))]
struct StrongInner<T: ?Sized>(Arc<T>);

#[repr(align(2))]
struct WeakInner<T: ?Sized>(Weak<T>);

// A few static assertions to show we're maintaing our minimum required alignment which is necessary
// to ensure we can safely tag our inner pointer to indicate if it's a strong or weak inner.
const _: () = assert!(
    std::mem::align_of::<StrongInner<()>>() >= 2,
    "alignment of StrongInner<T> must always be 2 or greater"
);
const _: () = assert!(
    std::mem::align_of::<StrongInner<dyn CounterFn>>() >= 2,
    "alignment of StrongInner<T> must always be 2 or greater"
);
const _: () = assert!(
    std::mem::align_of::<WeakInner<()>>() >= 2,
    "alignment of WeakInner<T> must always be 2 or greater"
);
const _: () = assert!(
    std::mem::align_of::<WeakInner<dyn CounterFn>>() >= 2,
    "alignment of WeakInner<T> must always be 2 or greater"
);

impl<T: ?Sized> Pointer<T> {
    /// Creates a no-op pointer.
    fn noop() -> Self {
        Self { ptr: AtomicPtr::new(std::ptr::null_mut()), update: Mutex::new(()), _ty: PhantomData }
    }

    /// Creates a strong pointer.
    fn strong(strong: Arc<T>) -> Self {
        let ptr = Box::into_raw(Box::new(StrongInner(strong)));
        let tagged_ptr = (ptr as usize | 1) as *mut ();
        Self { ptr: AtomicPtr::new(tagged_ptr), update: Mutex::new(()), _ty: PhantomData }
    }

    /// Creates a weak pointer.
    fn weak(weak: Weak<T>) -> Self {
        let ptr = Box::into_raw(Box::new(WeakInner(weak)));
        Self { ptr: AtomicPtr::new(ptr as *mut ()), update: Mutex::new(()), _ty: PhantomData }
    }

    /// Runs the given closure `f` with a reference to the inner data.
    ///
    /// If this is a weak pointer, and the data wrapped by the pointer has gone away, the internal
    /// pointer state is updated to point to a new, valid pointer (by calling `replace` to create a
    /// new pointer which is consumed) before again trying to take a reference to the inner data in
    /// order to call `f`. This logic happens in a loop until a reference is
    /// successfully taken.
    ///
    /// If the pointer is a no-op pointer, `f` and `replace` are not called.
    fn with_ref(&self, f: impl FnOnce(&T), replace: impl Fn() -> Self) {
        loop {
            let ptr = self.ptr.load(Ordering::Acquire);
            if ptr.is_null() {
                // We don't run the given closure if the pointer is a no-op.
                return;
            }

            match PointerKind::from_ptr(ptr) {
                PointerKind::Strong { ptr: strong_ptr } => {
                    let strong = unsafe { &*strong_ptr };
                    f(&strong.0);
                    return;
                }
                PointerKind::Weak { ptr: weak_ptr } => {
                    let weak = unsafe { &*weak_ptr };
                    match weak.0.upgrade() {
                        Some(strong) => {
                            f(&strong);
                            return;
                        }
                        None => {
                            // If the upgrade fails, try and take the update lock.
                            //
                            // When we get the lock, re-check again to see if the pointer is still
                            // the same. If so, we now have exclusive access to update the pointer
                            // by creating a replacement pointer via `replace()` and pointing to the
                            // inner state that it has, effectively consuming it.
                            //
                            // We wrap it in `ManuallyDrop` first to avoid triggering the drop logic
                            // after consuming it.
                            let _guard = self.update.lock().unwrap();
                            let current_ptr = self.ptr.load(Ordering::Acquire);
                            if current_ptr == ptr {
                                let new_pointer = ManuallyDrop::new(replace());
                                let new_ptr = new_pointer.ptr.load(Ordering::Acquire);
                                self.ptr.store(new_ptr, Ordering::Release);
                            }
                        }
                    }
                }
            }
        }
    }
}

impl<T: ?Sized> Drop for Pointer<T> {
    fn drop(&mut self) {
        let ptr = self.ptr.load(Ordering::Acquire);
        if !ptr.is_null() {
            match PointerKind::<T>::from_ptr(ptr) {
                PointerKind::Strong { ptr: strong_ptr } => {
                    // SAFETY: If the pointer is non-null, then we know it came directly from
                    // `Box::into_raw`, so it's pointing to an initialized value, it's aligned, etc.
                    let strong = unsafe { Box::from_raw(strong_ptr) };
                    drop(strong);
                }
                PointerKind::Weak { ptr: weak_ptr } => {
                    // SAFETY: If the pointer is non-null, then we know it came directly from
                    // `Box::into_raw`, so it's pointing to an initialized value, it's aligned, etc.
                    let weak = unsafe { Box::from_raw(weak_ptr) };
                    drop(weak);
                }
            }
        }
    }
}