Skip to main content

netstack3_sync/
rc.rs

1// Copyright 2022 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Synchronized reference counting primitives.
6//!
7//! This module introduces a family of reference counted types that allows
8//! marking the underlying data for destruction before all strongly references
9//! to the data are dropped. This enables the following features:
10//!   * Upgrading a weak reference to a strong reference succeeds iff at least
11//!     one strong reference exists _and_ the data has not been marked for
12//!     destruction.
13//!   * Allow waiting for all strongly-held references to be dropped after
14//!     marking the data.
15
16use core::cell::UnsafeCell;
17use core::fmt::Debug;
18use core::hash::{Hash, Hasher};
19use core::ops::Deref;
20use core::panic::Location;
21use core::sync::atomic::{AtomicBool, Ordering};
22
23use alloc::boxed::Box;
24use derivative::Derivative;
25
26/// Returns `true` if the current thread is unwinding due to a panic.
27///
28/// Always returns `false` when panics abort, since no code runs while
29/// panicking in that case.
30#[inline]
31fn thread_panicking() -> bool {
32    #[cfg(panic = "abort")]
33    {
34        false
35    }
36    #[cfg(not(panic = "abort"))]
37    {
38        extern crate std;
39        std::thread::panicking()
40    }
41}
42
43mod caller {
44    //! Provides tracking of instances via tracked caller location.
45    //!
46    //! Callers are only tracked in debug builds. All operations and types
47    //! are no-ops and empty unless the `rc-debug-names` feature is enabled.
48
49    use core::fmt::Debug;
50    use core::panic::Location;
51
52    /// Records reference-counted names of instances.
53    #[derive(Default)]
54    pub(super) struct Callers {
55        /// The names that were inserted and aren't known to be gone.
56        ///
57        /// This holds weak references to allow callers to drop without
58        /// synchronizing. Invalid weak pointers are cleaned up periodically but
59        /// are not logically present.
60        ///
61        /// Note that using [`std::sync::Mutex`] here is intentional to opt this
62        /// out of loom checking, which makes testing with `rc-debug-names`
63        /// impossibly slow.
64        #[cfg(feature = "rc-debug-names")]
65        pub(super) callers: std::sync::Mutex<std::collections::HashMap<Location<'static>, usize>>,
66    }
67
68    impl Debug for Callers {
69        #[cfg(not(feature = "rc-debug-names"))]
70        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
71            write!(f, "(Not Tracked)")
72        }
73        #[cfg(feature = "rc-debug-names")]
74        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
75            let Self { callers } = self;
76            let callers = callers.lock().unwrap();
77            write!(f, "[\n")?;
78            for (l, c) in callers.iter() {
79                write!(f, "   {l} => {c},\n")?;
80            }
81            write!(f, "]")
82        }
83    }
84
85    impl Callers {
86        /// Creates a new [`Callers`] from the given [`Location`].
87        ///
88        /// On non-debug builds, this is a no-op.
89        pub(super) fn insert(&self, caller: &Location<'static>) -> TrackedCaller {
90            #[cfg(not(feature = "rc-debug-names"))]
91            {
92                let _ = caller;
93                TrackedCaller {}
94            }
95            #[cfg(feature = "rc-debug-names")]
96            {
97                let Self { callers } = self;
98                let mut callers = callers.lock().unwrap();
99                let count = callers.entry(caller.clone()).or_insert(0);
100                *count += 1;
101                TrackedCaller { location: caller.clone() }
102            }
103        }
104    }
105
106    #[derive(Debug)]
107    pub(super) struct TrackedCaller {
108        #[cfg(feature = "rc-debug-names")]
109        pub(super) location: Location<'static>,
110    }
111
112    impl TrackedCaller {
113        #[cfg(not(feature = "rc-debug-names"))]
114        pub(super) fn release(&mut self, Callers {}: &Callers) {
115            let Self {} = self;
116        }
117
118        #[cfg(feature = "rc-debug-names")]
119        pub(super) fn release(&mut self, Callers { callers }: &Callers) {
120            let Self { location } = self;
121            let mut callers = callers.lock().unwrap();
122            let mut entry = match callers.entry(location.clone()) {
123                std::collections::hash_map::Entry::Vacant(_) => {
124                    panic!("location {location:?} was not in the callers map")
125                }
126                std::collections::hash_map::Entry::Occupied(o) => o,
127            };
128
129            let sub = entry
130                .get()
131                .checked_sub(1)
132                .unwrap_or_else(|| panic!("zero-count location {location:?} in map"));
133            if sub == 0 {
134                let _: usize = entry.remove();
135            } else {
136                *entry.get_mut() = sub;
137            }
138        }
139    }
140}
141
142mod resource_token {
143    use core::fmt::Debug;
144    use core::marker::PhantomData;
145    use core::num::NonZeroU64;
146    use core::sync::atomic::{AtomicU64, Ordering};
147
148    /// An opaque token associated with a resource.
149    ///
150    /// It can be used to create debug and trace identifiers for the resource,
151    /// but it should not be used as a unique identifier of the resource inside
152    /// the netstack.
153    ///
154    /// By default the lifetime of a token is bound the resource that token
155    /// belongs to, but it can be extended by calling
156    /// [`ResourceToken::extend_lifetime`].
157    ///
158    /// Internally the value is stored as `NonZeroU64`. This is an optimization
159    /// to save memory when the value is wrapped in `Option` (it allows the
160    /// compiler to fit `Option<ResourceToken>` in 8 bytes).
161    #[cfg_attr(any(test, feature = "testutils"), derive(PartialEq, Eq, PartialOrd, Ord))]
162    #[derive(Clone)]
163    pub struct ResourceToken<'a> {
164        value: NonZeroU64,
165        _marker: PhantomData<&'a ()>,
166    }
167
168    impl<'a> ResourceToken<'a> {
169        /// Extends lifetime of the token.
170        ///
171        /// # Discussion
172        ///
173        /// It's generally okay to extend the lifetime of the token, but prefer
174        /// to use tokens bound to the resource's lifetime whenever possible,
175        /// since it provides guardrails against identifiers that outlive the
176        /// resource itself.
177        pub fn extend_lifetime(self) -> ResourceToken<'static> {
178            ResourceToken { value: self.value, _marker: PhantomData }
179        }
180
181        /// Returns internal value. Consumes `self`.
182        ///
183        /// # Discussion
184        ///
185        /// Export to `u64` when a representation is needed for interaction with
186        /// other processes or components such as trace identifiers and eBPF
187        /// socket cookies.
188        ///
189        /// Refrain from using the returned value within the netstack otherwise.
190        pub fn export_value(self) -> u64 {
191            self.value.get()
192        }
193    }
194
195    impl<'a> Debug for ResourceToken<'a> {
196        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
197            write!(f, "{}", self.value)
198        }
199    }
200
201    /// Holder of a value for `ResourceToken`. Vends `ResourceToken` instances
202    /// with the same value and the lifetime bound to the lifetime of the holder.
203    ///
204    /// The [`Default`] implementation generates a new unique value.
205    pub struct ResourceTokenValue(NonZeroU64);
206
207    impl ResourceTokenValue {
208        /// Creates a new token.
209        pub fn token(&self) -> ResourceToken<'_> {
210            let ResourceTokenValue(value) = self;
211            ResourceToken { value: *value, _marker: PhantomData }
212        }
213    }
214
215    impl core::fmt::Debug for ResourceTokenValue {
216        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
217            let ResourceTokenValue(value) = self;
218            write!(f, "{}", value)
219        }
220    }
221
222    impl Default for ResourceTokenValue {
223        fn default() -> Self {
224            static NEXT_TOKEN: AtomicU64 = AtomicU64::new(1);
225            // NB: Fetch add will cause the counter to rollback to 0 if we
226            // happen to exceed `u64::MAX` instantiations. In practice, that's
227            // an impossibility (at 1 billion instantiations per second, the
228            // counter is valid for > 500 years). Spare the CPU cycles and don't
229            // bother attempting to handle overflow.
230            Self(NonZeroU64::new(NEXT_TOKEN.fetch_add(1, Ordering::Relaxed)).unwrap())
231        }
232    }
233}
234
235pub use resource_token::{ResourceToken, ResourceTokenValue};
236
237mod debug_id {
238    use super::ResourceToken;
239    use core::fmt::Debug;
240
241    /// A debug identifier for the RC types exposed in the parent module.
242    ///
243    /// Encompasses the underlying pointer for the RC type, as well as
244    /// (optionally) the globally unique [`ResourceToken`].
245    pub(super) enum DebugId<T> {
246        /// Used in contexts that have access to the [`ResourceToken`], e.g.
247        /// [`Primary`], [`Strong`], and sometimes [`Weak`] RC types.
248        WithToken { ptr: *const T, token: ResourceToken<'static> },
249        /// Used in contexts that don't have access to the [`ResourceToken`], e.g.
250        /// [`Weak`] RC types that cannot be upgraded.
251        WithoutToken { ptr: *const T },
252    }
253
254    impl<T> Debug for DebugId<T> {
255        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> Result<(), core::fmt::Error> {
256            match self {
257                DebugId::WithToken { ptr, token } => write!(f, "{:?}:{:?}", token, ptr),
258                DebugId::WithoutToken { ptr } => write!(f, "?:{:?}", ptr),
259            }
260        }
261    }
262}
263
264/// A slot for the destruction notifier.
265///
266/// Wraps an [`UnsafeCell`] to isolate the [`Sync`] implementation so that
267/// [`Inner<T>`] can automatically derive `Sync` based on its fields without
268/// requiring a blanket manual `unsafe impl Sync`.
269struct NotifierSlot<T>(UnsafeCell<Option<Box<dyn Notifier<T>>>>);
270
271// SAFETY: `NotifierSlot<T>` wraps an `UnsafeCell` for interior mutability.
272// The `UnsafeCell` in `notifier` is only ever accessed through a shared
273// reference when consuming the unique `Primary<T>` instance (in
274// `Primary::unwrap_with_notifier` and `Primary::unwrap_or_notify_with`), which
275// cannot race with any other shared access because only one `Primary<T>` ever
276// exists for a given `Inner<T>`, nor with `Inner::drop` or `Inner::unwrap`
277// because `Primary<T>`'s `Arc<Inner<T>>` is still alive during the write and
278// its subsequent `Arc::drop` Release-Acquire synchronization happens-before
279// `Inner::drop`. `Notifier<T>: Send` ensures the boxed notifier can be moved
280// across threads from `Primary` to whichever thread drops the last `Strong`.
281unsafe impl<T: Sync> Sync for NotifierSlot<T> {}
282
283#[derive(Derivative)]
284#[derivative(Debug)]
285struct Inner<T> {
286    marked_for_destruction: AtomicBool,
287    callers: caller::Callers,
288    data: core::mem::ManuallyDrop<T>,
289    // Safety invariant: `notifier` is initialized to `None` when `Inner<T>` is
290    // constructed and may only be mutated through a shared reference by the
291    // unique `Primary<T>` owner while its `Arc<Inner<T>>` is still alive
292    // (before dropping that `Arc`), or through `&mut Inner<T>` during
293    // destruction.
294    #[derivative(Debug = "ignore")]
295    notifier: NotifierSlot<T>,
296    resource_token: ResourceTokenValue,
297}
298
299impl<T> Inner<T> {
300    fn pre_drop_check(marked_for_destruction: &AtomicBool) {
301        // `Ordering::Acquire` because we want to synchronize with with the
302        // `Ordering::Release` write to `marked_for_destruction` so that all
303        // memory writes before the reference was marked for destruction is
304        // visible here.
305        assert!(marked_for_destruction.load(Ordering::Acquire), "Must be marked for destruction");
306    }
307
308    fn unwrap(mut self) -> T {
309        // We cannot destructure `self` by value since `Inner` implements
310        // `Drop`. So we must manually drop all the fields but data and then
311        // forget self.
312        let Inner { marked_for_destruction, data, callers: holders, notifier, resource_token } =
313            &mut self;
314
315        // Make sure that `inner` is in a valid state for destruction.
316        //
317        // Note that we do not actually destroy all of `self` here; we decompose
318        // it into its parts, keeping what we need & throwing away what we
319        // don't. Regardless, we perform the same checks.
320        Inner::<T>::pre_drop_check(marked_for_destruction);
321
322        // SAFETY: Safe since we own `self` and `self` is immediately forgotten
323        // below so the its destructor (and those of its fields) will not be run
324        // as a result of `self` being dropped.
325        let data = unsafe {
326            // Explicitly drop since we do not need these anymore.
327            core::ptr::drop_in_place(marked_for_destruction);
328            core::ptr::drop_in_place(holders);
329            core::ptr::drop_in_place(notifier);
330            core::ptr::drop_in_place(resource_token);
331
332            core::mem::ManuallyDrop::take(data)
333        };
334        // Forget self now to prevent its `Drop::drop` impl from being run which
335        // will attempt to destroy `data` but still perform pre-drop checks on
336        // `Inner`'s state.
337        core::mem::forget(self);
338
339        data
340    }
341
342    /// Sets the notifier for this `Inner`.
343    ///
344    /// # Safety
345    ///
346    /// The caller must ensure that `set_notifier` is called at most once on
347    /// `this`, using the `Arc<Self>` extracted from consuming the unique
348    /// [`Primary<T>`] associated with `this`.
349    unsafe fn set_notifier<N: Notifier<T> + 'static>(this: alloc::sync::Arc<Self>, notifier: N) {
350        let Self { notifier: NotifierSlot(slot), .. } = &*this;
351
352        // Using dynamic dispatch to notify allows us to not have to know the
353        // notifier that will be used from creation and spread the type on all
354        // reference types in this crate. The assumption is that the allocation
355        // and dynamic dispatch costs here are tiny compared to the overall work
356        // of destroying the resources this module is targeting.
357        let boxed: Box<dyn Notifier<T>> = Box::new(notifier);
358
359        // SAFETY: By the caller's `# Safety` contract, `set_notifier` is called
360        // at most once on `this` (from the unique `Primary<T>`). No other
361        // shared reference accesses `slot`, and exclusive accesses in
362        // `Inner::drop` and `Inner::unwrap` cannot run while `this: Arc<Self>`
363        // is alive, so we have unique access to the contents of `slot` for the
364        // duration of this statement.
365        let prev_notifier = unsafe { (*slot.get()).replace(boxed) };
366        // Uphold invariant that this can only be done once from Primary.
367        assert!(prev_notifier.is_none(), "can't have a notifier already installed");
368
369        // Explicitly drop `this` after writing `slot` to uphold the `Drop`
370        // invariant: `Arc::drop` either runs `Inner::drop` on this thread
371        // (if `this` was the last strong reference) or performs a Release
372        // decrement that synchronizes-with the Acquire fence in `Arc::drop`
373        // before `Inner::drop` runs on whichever thread drops the last strong
374        // reference.
375        core::mem::drop(this);
376    }
377}
378
379impl<T> Drop for Inner<T> {
380    fn drop(&mut self) {
381        let Inner {
382            marked_for_destruction,
383            data,
384            callers: _,
385            notifier: NotifierSlot(notifier),
386            resource_token: _,
387        } = self;
388        // Take data out of ManuallyDrop in case we panic in pre_drop_check.
389        // That'll ensure data is dropped if we hit the panic.
390        //
391        //  SAFETY: Safe because ManuallyDrop is not referenced again after
392        // taking.
393        let data = unsafe { core::mem::ManuallyDrop::take(data) };
394        Self::pre_drop_check(marked_for_destruction);
395        if let Some(mut notifier) = notifier.get_mut().take() {
396            notifier.notify(data);
397        }
398    }
399}
400
401/// A primary reference.
402///
403/// Note that only one `Primary` may be associated with data. This is
404/// enforced by not implementing [`Clone`].
405///
406/// For now, this reference is no different than a [`Strong`] but later changes
407/// will enable blocking the destruction of a primary reference until all
408/// strongly held references are dropped.
409#[derive(Debug)]
410pub struct Primary<T> {
411    inner: core::mem::ManuallyDrop<alloc::sync::Arc<Inner<T>>>,
412}
413
414impl<T> Drop for Primary<T> {
415    fn drop(&mut self) {
416        let was_marked = self.mark_for_destruction();
417        let Self { inner } = self;
418        // Take the inner out of ManuallyDrop early so its Drop impl will run in
419        // case we panic here.
420        // SAFETY: Safe because we don't reference ManuallyDrop again.
421        let inner = unsafe { core::mem::ManuallyDrop::take(inner) };
422
423        // Make debugging easier: don't panic if a panic is already happening
424        // since double-panics are annoying to debug. This means that the
425        // invariants provided by Primary are possibly violated during an
426        // unwind, but we're sidestepping that problem because Fuchsia is our
427        // only audience here.
428        if !thread_panicking() {
429            assert_eq!(was_marked, false, "Must not be marked for destruction yet");
430
431            let Inner {
432                marked_for_destruction: _,
433                callers,
434                data: _,
435                notifier: _,
436                resource_token: _,
437            } = &*inner;
438
439            // Make sure that this `Primary` is the last thing to hold a strong
440            // reference to the underlying data when it is being dropped.
441            let refs = alloc::sync::Arc::strong_count(&inner).checked_sub(1).unwrap();
442            assert!(
443                refs == 0,
444                "dropped Primary with {refs} strong refs remaining, \
445                            Callers={callers:?}"
446            );
447        }
448    }
449}
450
451impl<T> AsRef<T> for Primary<T> {
452    fn as_ref(&self) -> &T {
453        self.deref()
454    }
455}
456
457impl<T> Deref for Primary<T> {
458    type Target = T;
459
460    fn deref(&self) -> &T {
461        let Self { inner } = self;
462        let Inner { marked_for_destruction: _, data, callers: _, notifier: _, resource_token: _ } =
463            &***inner;
464        data
465    }
466}
467
468impl<T> Primary<T> {
469    // Marks this primary reference as ready for destruction. Used by all
470    // dropping flows. We take &mut self here to ensure we have the only
471    // possible reference to Primary. Returns whether it was already marked for
472    // destruction.
473    fn mark_for_destruction(&mut self) -> bool {
474        let Self { inner } = self;
475        // `Ordering::Release` because want to make sure that all memory writes
476        // before dropping this `Primary` synchronizes with later attempts to
477        // upgrade weak pointers and the `Drop::drop` impl of `Inner`.
478        inner.marked_for_destruction.swap(true, Ordering::Release)
479    }
480
481    /// Returns a new strongly-held reference.
482    pub fn new(data: T) -> Primary<T> {
483        Primary {
484            inner: core::mem::ManuallyDrop::new(alloc::sync::Arc::new(Inner {
485                marked_for_destruction: AtomicBool::new(false),
486                callers: caller::Callers::default(),
487                data: core::mem::ManuallyDrop::new(data),
488                notifier: NotifierSlot(UnsafeCell::new(None)),
489                resource_token: ResourceTokenValue::default(),
490            })),
491        }
492    }
493
494    /// Constructs a new `Primary<T>` while giving you a Weak<T> to the
495    /// allocation, to allow you to construct a `T` which holds a weak pointer
496    /// to itself.
497    ///
498    /// Like for [`Arc::new_cyclic`], the `Weak` reference provided to `data_fn`
499    /// cannot be upgraded until the [`Primary`] is constructed.
500    pub fn new_cyclic(data_fn: impl FnOnce(Weak<T>) -> T) -> Primary<T> {
501        Primary {
502            inner: core::mem::ManuallyDrop::new(alloc::sync::Arc::new_cyclic(move |weak| Inner {
503                marked_for_destruction: AtomicBool::new(false),
504                callers: caller::Callers::default(),
505                data: core::mem::ManuallyDrop::new(data_fn(Weak(weak.clone()))),
506                notifier: NotifierSlot(UnsafeCell::new(None)),
507                resource_token: ResourceTokenValue::default(),
508            })),
509        }
510    }
511
512    /// Clones a strongly-held reference.
513    #[cfg_attr(feature = "rc-debug-names", track_caller)]
514    pub fn clone_strong(Self { inner }: &Self) -> Strong<T> {
515        let Inner { data: _, callers, marked_for_destruction: _, notifier: _, resource_token: _ } =
516            &***inner;
517        let caller = callers.insert(Location::caller());
518        Strong { inner: alloc::sync::Arc::clone(inner), caller }
519    }
520
521    /// Returns a weak reference pointing to the same underlying data.
522    pub fn downgrade(Self { inner }: &Self) -> Weak<T> {
523        Weak(alloc::sync::Arc::downgrade(inner))
524    }
525
526    /// Returns true if the two pointers point to the same allocation.
527    pub fn ptr_eq(
528        Self { inner: this }: &Self,
529        Strong { inner: other, caller: _ }: &Strong<T>,
530    ) -> bool {
531        alloc::sync::Arc::ptr_eq(this, other)
532    }
533
534    /// Returns [`Debug`] implementation that is stable and unique
535    /// for the data held behind this [`Primary`].
536    pub fn debug_id(&self) -> impl Debug + '_ {
537        let Self { inner } = self;
538
539        // The lifetime of the returned `DebugId` is bound to the lifetime
540        // of `self`.
541        let token = inner.resource_token.token().extend_lifetime();
542
543        debug_id::DebugId::WithToken { ptr: alloc::sync::Arc::as_ptr(inner), token }
544    }
545
546    fn mark_for_destruction_and_take_inner(mut this: Self) -> alloc::sync::Arc<Inner<T>> {
547        // Prepare for destruction.
548        assert!(!this.mark_for_destruction());
549        let Self { inner } = &mut this;
550        // SAFETY: Safe because inner can't be used after this. We forget
551        // our Primary reference to prevent its Drop impl from running.
552        let inner = unsafe { core::mem::ManuallyDrop::take(inner) };
553        core::mem::forget(this);
554        inner
555    }
556
557    fn try_unwrap(this: Self) -> Result<T, alloc::sync::Arc<Inner<T>>> {
558        let inner = Self::mark_for_destruction_and_take_inner(this);
559        alloc::sync::Arc::try_unwrap(inner).map(Inner::unwrap)
560    }
561
562    /// Returns the inner value if no [`Strong`] references are held.
563    ///
564    /// # Panics
565    ///
566    /// Panics if [`Strong`] references are held when this function is called.
567    pub fn unwrap(this: Self) -> T {
568        Self::try_unwrap(this).unwrap_or_else(|inner| {
569            let callers = &inner.callers;
570            let refs = alloc::sync::Arc::strong_count(&inner).checked_sub(1).unwrap();
571            panic!("can't unwrap, still had {refs} strong refs: {callers:?}");
572        })
573    }
574
575    /// Marks this [`Primary`] for destruction and uses `notifier` as a signaler
576    /// for when destruction of all strong references is terminated. After
577    /// calling `unwrap_with_notifier` [`Weak`] references can no longer be
578    /// upgraded.
579    pub fn unwrap_with_notifier<N: Notifier<T> + 'static>(this: Self, notifier: N) {
580        let inner = Self::mark_for_destruction_and_take_inner(this);
581        // SAFETY: `this: Primary<T>` is the unique `Primary` reference for
582        // `inner` and is consumed by value here, so `set_notifier` is called at
583        // most once on `inner`.
584        unsafe { Inner::set_notifier(inner, notifier) };
585    }
586
587    /// Marks this [`Primary`] for destruction and returns `Ok` if this was the
588    /// last strong reference standing for it. Otherwise `new_notifier` is
589    /// called to create a new notifier to observe deferred destruction.
590    ///
591    /// Like [`Primary::unwrap_with_notifier`], [`Weak`] references can no
592    /// longer be upgraded after calling `unwrap_or_notify_with`.
593    pub fn unwrap_or_notify_with<N: Notifier<T> + 'static, O, F: FnOnce() -> (N, O)>(
594        this: Self,
595        new_notifier: F,
596    ) -> Result<T, O> {
597        Self::try_unwrap(this).map_err(move |inner| {
598            let (notifier, output) = new_notifier();
599            // SAFETY: `this: Primary<T>` is the unique `Primary` reference for
600            // `inner` and is consumed by value here, so `set_notifier` is
601            // called at most once on `inner`.
602            unsafe { Inner::set_notifier(inner, notifier) };
603            output
604        })
605    }
606
607    /// Creates a [`DebugReferences`] instance.
608    pub fn debug_references(this: &Self) -> DebugReferences<T> {
609        let Self { inner } = this;
610        DebugReferences(alloc::sync::Arc::downgrade(&*inner))
611    }
612}
613
614/// A strongly-held reference.
615///
616/// Similar to an [`alloc::sync::Arc`] but holding a `Strong` acts as a witness
617/// to the live-ness of the underlying data. That is, holding a `Strong` implies
618/// that the underlying data has not yet been destroyed.
619///
620/// Note that `Strong`'s implementation of [`Hash`] and [`PartialEq`] operate on
621/// the pointer itself and not the underlying data.
622#[derive(Debug, Derivative)]
623pub struct Strong<T> {
624    inner: alloc::sync::Arc<Inner<T>>,
625    caller: caller::TrackedCaller,
626}
627
628impl<T> Drop for Strong<T> {
629    fn drop(&mut self) {
630        let Self { inner, caller } = self;
631        let Inner { marked_for_destruction: _, callers, data: _, notifier: _, resource_token: _ } =
632            &**inner;
633        caller.release(callers);
634    }
635}
636
637impl<T> AsRef<T> for Strong<T> {
638    fn as_ref(&self) -> &T {
639        self.deref()
640    }
641}
642
643impl<T> Deref for Strong<T> {
644    type Target = T;
645
646    fn deref(&self) -> &T {
647        let Self { inner, caller: _ } = self;
648        let Inner { marked_for_destruction: _, data, callers: _, notifier: _, resource_token: _ } =
649            inner.deref();
650        data
651    }
652}
653
654impl<T> core::cmp::Eq for Strong<T> {}
655
656impl<T> core::cmp::PartialEq for Strong<T> {
657    fn eq(&self, other: &Self) -> bool {
658        Self::ptr_eq(self, other)
659    }
660}
661
662impl<T> Hash for Strong<T> {
663    fn hash<H: Hasher>(&self, state: &mut H) {
664        let Self { inner, caller: _ } = self;
665        alloc::sync::Arc::as_ptr(inner).hash(state)
666    }
667}
668
669impl<T> Clone for Strong<T> {
670    #[cfg_attr(feature = "rc-debug-names", track_caller)]
671    fn clone(&self) -> Self {
672        let Self { inner, caller: _ } = self;
673        let Inner { data: _, marked_for_destruction: _, callers, notifier: _, resource_token: _ } =
674            &**inner;
675        let caller = callers.insert(Location::caller());
676        Self { inner: alloc::sync::Arc::clone(inner), caller }
677    }
678}
679
680impl<T> Strong<T> {
681    /// Returns a weak reference pointing to the same underlying data.
682    pub fn downgrade(Self { inner, caller: _ }: &Self) -> Weak<T> {
683        Weak(alloc::sync::Arc::downgrade(inner))
684    }
685
686    /// Returns [`Debug`] implementation that is stable and unique
687    /// for the data held behind this [`Strong`].
688    pub fn debug_id(&self) -> impl Debug + '_ {
689        let Self { inner, caller: _ } = self;
690
691        // The lifetime of the returned `DebugId` is bound to the lifetime
692        // of `self`.
693        let token = inner.resource_token.token().extend_lifetime();
694
695        debug_id::DebugId::WithToken { ptr: alloc::sync::Arc::as_ptr(inner), token }
696    }
697
698    /// Returns a [`ResourceToken`] that corresponds to this object.
699    pub fn resource_token(&self) -> ResourceToken<'_> {
700        self.inner.resource_token.token()
701    }
702
703    /// Returns true if the inner value has since been marked for destruction.
704    pub fn marked_for_destruction(Self { inner, caller: _ }: &Self) -> bool {
705        let Inner { marked_for_destruction, data: _, callers: _, notifier: _, resource_token: _ } =
706            inner.as_ref();
707        // `Ordering::Acquire` because we want to synchronize with with the
708        // `Ordering::Release` write to `marked_for_destruction` so that all
709        // memory writes before the reference was marked for destruction is
710        // visible here.
711        marked_for_destruction.load(Ordering::Acquire)
712    }
713
714    /// Returns true if the two pointers point to the same allocation.
715    pub fn weak_ptr_eq(Self { inner: this, caller: _ }: &Self, Weak(other): &Weak<T>) -> bool {
716        core::ptr::eq(alloc::sync::Arc::as_ptr(this), other.as_ptr())
717    }
718
719    /// Returns true if the two pointers point to the same allocation.
720    pub fn ptr_eq(
721        Self { inner: this, caller: _ }: &Self,
722        Self { inner: other, caller: _ }: &Self,
723    ) -> bool {
724        alloc::sync::Arc::ptr_eq(this, other)
725    }
726
727    /// Compares the two pointers.
728    pub fn ptr_cmp(
729        Self { inner: this, caller: _ }: &Self,
730        Self { inner: other, caller: _ }: &Self,
731    ) -> core::cmp::Ordering {
732        let this = alloc::sync::Arc::as_ptr(this);
733        let other = alloc::sync::Arc::as_ptr(other);
734        this.cmp(&other)
735    }
736
737    /// Creates a [`DebugReferences`] instance.
738    pub fn debug_references(this: &Self) -> DebugReferences<T> {
739        let Self { inner, caller: _ } = this;
740        DebugReferences(alloc::sync::Arc::downgrade(inner))
741    }
742}
743
744/// A weakly-held reference.
745///
746/// Similar to an [`alloc::sync::Weak`].
747///
748/// A `Weak` does not make any claim to the live-ness of the underlying data.
749/// Holders of a `Weak` must attempt to upgrade to a [`Strong`] through
750/// [`Weak::upgrade`] to access the underlying data.
751///
752/// Note that `Weak`'s implementation of [`Hash`] and [`PartialEq`] operate on
753/// the pointer itself and not the underlying data.
754#[derive(Debug)]
755pub struct Weak<T>(alloc::sync::Weak<Inner<T>>);
756
757impl<T> core::cmp::Eq for Weak<T> {}
758
759impl<T> core::cmp::PartialEq for Weak<T> {
760    fn eq(&self, other: &Self) -> bool {
761        Self::ptr_eq(self, other)
762    }
763}
764
765impl<T> Hash for Weak<T> {
766    fn hash<H: Hasher>(&self, state: &mut H) {
767        let Self(this) = self;
768        this.as_ptr().hash(state)
769    }
770}
771
772impl<T> Clone for Weak<T> {
773    fn clone(&self) -> Self {
774        let Self(this) = self;
775        Weak(this.clone())
776    }
777}
778
779impl<T> Weak<T> {
780    /// Returns true if the two pointers point to the same allocation.
781    pub fn ptr_eq(&self, Self(other): &Self) -> bool {
782        let Self(this) = self;
783        this.ptr_eq(other)
784    }
785
786    /// Returns [`Debug`] implementation that is stable and unique
787    /// for the data held behind this [`Weak`].
788    pub fn debug_id(&self) -> impl Debug + '_ {
789        match self.upgrade() {
790            Some(strong) => {
791                let Strong { inner, caller: _ } = &strong;
792
793                // The lifetime of the returned `DebugId` is still bound to the
794                // lifetime of `self`.
795                let token = inner.resource_token.token().extend_lifetime();
796
797                debug_id::DebugId::WithToken { ptr: alloc::sync::Arc::as_ptr(&inner), token }
798            }
799            None => {
800                let Self(this) = self;
801                // NB: If we can't upgrade the socket, we can't know the token.
802                debug_id::DebugId::WithoutToken { ptr: this.as_ptr() }
803            }
804        }
805    }
806
807    /// Attempts to upgrade to a [`Strong`].
808    ///
809    /// Returns `None` if the inner value has since been marked for destruction.
810    #[cfg_attr(feature = "rc-debug-names", track_caller)]
811    pub fn upgrade(&self) -> Option<Strong<T>> {
812        let Self(weak) = self;
813        let arc = weak.upgrade()?;
814        let Inner { marked_for_destruction, data: _, callers, notifier: _, resource_token: _ } =
815            arc.deref();
816
817        // `Ordering::Acquire` because we want to synchronize with with the
818        // `Ordering::Release` write to `marked_for_destruction` so that all
819        // memory writes before the reference was marked for destruction is
820        // visible here.
821        if !marked_for_destruction.load(Ordering::Acquire) {
822            let caller = callers.insert(Location::caller());
823            Some(Strong { inner: arc, caller })
824        } else {
825            None
826        }
827    }
828
829    /// Gets the number of [`Primary`] and [`Strong`] references to this allocation.
830    pub fn strong_count(&self) -> usize {
831        let Self(weak) = self;
832        weak.strong_count()
833    }
834
835    /// Creates a [`DebugReferences`] instance.
836    pub fn debug_references(&self) -> DebugReferences<T> {
837        let Self(inner) = self;
838        DebugReferences(inner.clone())
839    }
840}
841
842fn debug_refs(
843    refs: Option<(usize, &AtomicBool, &caller::Callers)>,
844    name: &'static str,
845    f: &mut core::fmt::Formatter<'_>,
846) -> core::fmt::Result {
847    let mut f = f.debug_struct(name);
848    match refs {
849        Some((strong_count, marked_for_destruction, callers)) => f
850            .field("strong_count", &strong_count)
851            .field("marked_for_destruction", marked_for_destruction)
852            .field("callers", callers)
853            .finish(),
854        None => {
855            let strong_count = 0_usize;
856            f.field("strong_count", &strong_count).finish_non_exhaustive()
857        }
858    }
859}
860
861/// Provides a [`Debug`] implementation that contains information helpful for
862/// debugging dangling references.
863#[derive(Clone)]
864pub struct DebugReferences<T>(alloc::sync::Weak<Inner<T>>);
865
866impl<T> Debug for DebugReferences<T> {
867    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
868        let Self(inner) = self;
869        let inner = inner.upgrade();
870        let refs = inner.as_ref().map(|inner| {
871            (alloc::sync::Arc::strong_count(inner), &inner.marked_for_destruction, &inner.callers)
872        });
873        debug_refs(refs, "DebugReferences", f)
874    }
875}
876
877impl<T: Send + Sync + 'static> DebugReferences<T> {
878    /// Transforms this `DebugReferences` into a [`DynDebugReferences`].
879    pub fn into_dyn(self) -> DynDebugReferences {
880        let Self(w) = self;
881        DynDebugReferences(w)
882    }
883}
884
885/// Like [`DebugReferences`], but type-erases the contained type.
886#[derive(Clone)]
887pub struct DynDebugReferences(alloc::sync::Weak<dyn ExposeRefs>);
888
889impl Debug for DynDebugReferences {
890    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
891        let Self(inner) = self;
892        let inner = inner.upgrade();
893        let refs = inner.as_ref().map(|inner| {
894            let (marked_for_destruction, callers) = inner.refs_info();
895            (alloc::sync::Arc::strong_count(inner), marked_for_destruction, callers)
896        });
897        debug_refs(refs, "DynDebugReferences", f)
898    }
899}
900
901/// A trait allowing [`DynDebugReferences`] to erase the `T` type on [`Inner`].
902trait ExposeRefs: Send + Sync + 'static {
903    fn refs_info(&self) -> (&AtomicBool, &caller::Callers);
904}
905
906impl<T: Send + Sync + 'static> ExposeRefs for Inner<T> {
907    fn refs_info(&self) -> (&AtomicBool, &caller::Callers) {
908        (&self.marked_for_destruction, &self.callers)
909    }
910}
911
912/// Provides delegated notification of all strong references of a [`Primary`]
913/// being dropped.
914///
915/// See [`Primary::unwrap_with_notifier`].
916pub trait Notifier<T>: Send {
917    /// Called when the data contained in the [`Primary`] reference can be
918    /// extracted out because there are no more strong references to it.
919    fn notify(&mut self, data: T);
920}
921
922/// An implementation of [`Notifier`] that stores the unwrapped data in a
923/// `Clone` type.
924///
925/// Useful for tests where completion assertions are possible and useful.
926#[derive(Derivative)]
927#[derivative(Clone(bound = ""), Debug(bound = "T: Debug, R: lock_api::RawMutex"))]
928pub struct ArcNotifier<T, R>(alloc::sync::Arc<crate::Mutex<Option<T>, R>>);
929
930impl<T, R: lock_api::RawMutex> ArcNotifier<T, R> {
931    /// Creates a new `ArcNotifier`.
932    pub fn new() -> Self {
933        Self(alloc::sync::Arc::new(crate::Mutex::new(None)))
934    }
935
936    /// Takes the notified value, if any.
937    pub fn take(&self) -> Option<T> {
938        let Self(inner) = self;
939        inner.lock().take()
940    }
941}
942
943impl<T: Send, R: lock_api::RawMutex + Send + Sync> Notifier<T> for ArcNotifier<T, R> {
944    fn notify(&mut self, data: T) {
945        let Self(inner) = self;
946        assert!(inner.lock().replace(data).is_none(), "notified twice");
947    }
948}
949
950/// An implementation of [`Notifier`] that wraps another `Notifier` and applies
951/// a function on notified objects.
952pub struct MapNotifier<N, F> {
953    inner: N,
954    map: Option<F>,
955}
956
957impl<N, F> MapNotifier<N, F> {
958    /// Creates a new [`MapNotifier`] that wraps `notifier` with a mapping
959    /// function `F`.
960    pub fn new(notifier: N, map: F) -> Self {
961        Self { inner: notifier, map: Some(map) }
962    }
963}
964
965impl<A, B, N: Notifier<B>, F: FnOnce(A) -> B> Notifier<A> for MapNotifier<N, F>
966where
967    Self: Send,
968{
969    fn notify(&mut self, data: A) {
970        let Self { inner, map } = self;
971        let map = map.take().expect("notified twice");
972        inner.notify(map(data))
973    }
974}
975
976/// A handy implementation for the "never" type.
977impl<T> Notifier<T> for ! {
978    fn notify(&mut self, _data: T) {
979        match *self {}
980    }
981}
982
983#[cfg(test)]
984mod tests {
985    use super::*;
986
987    use alloc::format;
988    use fuchsia_sync::RawMutex;
989
990    type ArcNotifier<T> = super::ArcNotifier<T, RawMutex>;
991
992    #[test]
993    fn zombie_weak() {
994        let primary = Primary::new(());
995        let weak = {
996            let strong = Primary::clone_strong(&primary);
997            Strong::downgrade(&strong)
998        };
999        core::mem::drop(primary);
1000
1001        assert!(weak.upgrade().is_none());
1002    }
1003
1004    #[test]
1005    fn rcs() {
1006        const INITIAL_VAL: u8 = 1;
1007        const NEW_VAL: u8 = 2;
1008
1009        let primary = Primary::new(crate::Mutex::<_, RawMutex>::new(INITIAL_VAL));
1010        let strong = Primary::clone_strong(&primary);
1011        let weak = Strong::downgrade(&strong);
1012
1013        *primary.lock() = NEW_VAL;
1014        assert_eq!(*primary.deref().lock(), NEW_VAL);
1015        assert_eq!(*strong.deref().lock(), NEW_VAL);
1016        assert_eq!(*weak.upgrade().unwrap().deref().lock(), NEW_VAL);
1017    }
1018
1019    #[test]
1020    fn unwrap_primary_without_strong_held() {
1021        const VAL: u16 = 6;
1022        let primary = Primary::new(VAL);
1023        assert_eq!(Primary::unwrap(primary), VAL);
1024    }
1025
1026    #[test]
1027    #[should_panic(expected = "can't unwrap, still had 1 strong refs")]
1028    fn unwrap_primary_with_strong_held() {
1029        let primary = Primary::new(8);
1030        let _strong: Strong<_> = Primary::clone_strong(&primary);
1031        let _: u16 = Primary::unwrap(primary);
1032    }
1033
1034    #[test]
1035    #[should_panic(expected = "dropped Primary with 1 strong refs remaining")]
1036    fn drop_primary_with_strong_held() {
1037        let primary = Primary::new(9);
1038        let _strong: Strong<_> = Primary::clone_strong(&primary);
1039        core::mem::drop(primary);
1040    }
1041
1042    // This test trips LSAN on Fuchsia for some unknown reason. The host-side
1043    // test should be enough to protect us against regressing on the panicking
1044    // check.
1045    #[cfg(not(target_os = "fuchsia"))]
1046    #[test]
1047    #[should_panic(expected = "oopsie")]
1048    fn double_panic_protect() {
1049        let primary = Primary::new(9);
1050        let strong = Primary::clone_strong(&primary);
1051        // This will cause primary to be dropped before strong and would yield a
1052        // double panic if we didn't protect against it in Primary's Drop impl.
1053        let _tuple_to_invert_drop_order = (primary, strong);
1054        panic!("oopsie");
1055    }
1056
1057    #[cfg(feature = "rc-debug-names")]
1058    #[test]
1059    fn tracked_callers() {
1060        let primary = Primary::new(10);
1061        // Mark this position so we ensure all track_caller marks are correct in
1062        // the methods that support it.
1063        let here = Location::caller();
1064        let strong1 = Primary::clone_strong(&primary);
1065        let strong2 = strong1.clone();
1066        let weak = Strong::downgrade(&strong2);
1067        let strong3 = weak.upgrade().unwrap();
1068
1069        let Primary { inner } = &primary;
1070        let Inner { marked_for_destruction: _, callers, data: _, notifier: _, resource_token: _ } =
1071            &***inner;
1072
1073        let strongs = [strong1, strong2, strong3];
1074        let _: &Location<'_> = strongs.iter().enumerate().fold(here, |prev, (i, cur)| {
1075            let Strong { inner: _, caller: caller::TrackedCaller { location: cur } } = cur;
1076            assert_eq!(prev.file(), cur.file(), "{i}");
1077            assert!(prev.line() < cur.line(), "{prev} < {cur}, {i}");
1078            {
1079                let callers = callers.callers.lock().unwrap();
1080                assert_eq!(callers.get(cur).copied(), Some(1));
1081            }
1082
1083            cur
1084        });
1085
1086        // All callers must be removed from the callers map on drop.
1087        std::mem::drop(strongs);
1088        {
1089            let callers = callers.callers.lock().unwrap();
1090            let callers = callers.deref();
1091            assert!(callers.is_empty(), "{callers:?}");
1092        }
1093    }
1094    #[cfg(feature = "rc-debug-names")]
1095    #[test]
1096    fn same_location_caller_tracking() {
1097        fn clone_in_fn<T>(p: &Primary<T>) -> Strong<T> {
1098            Primary::clone_strong(p)
1099        }
1100
1101        let primary = Primary::new(10);
1102        let strong1 = clone_in_fn(&primary);
1103        let strong2 = clone_in_fn(&primary);
1104        assert_eq!(strong1.caller.location, strong2.caller.location);
1105
1106        let Primary { inner } = &primary;
1107        let Inner { marked_for_destruction: _, callers, data: _, notifier: _, resource_token: _ } =
1108            &***inner;
1109
1110        {
1111            let callers = callers.callers.lock().unwrap();
1112            assert_eq!(callers.get(&strong1.caller.location).copied(), Some(2));
1113        }
1114
1115        std::mem::drop(strong1);
1116        std::mem::drop(strong2);
1117
1118        {
1119            let callers = callers.callers.lock().unwrap();
1120            let callers = callers.deref();
1121            assert!(callers.is_empty(), "{callers:?}");
1122        }
1123    }
1124
1125    #[cfg(feature = "rc-debug-names")]
1126    #[test]
1127    #[should_panic(expected = "sync/src/rc.rs")]
1128    fn callers_in_panic() {
1129        let primary = Primary::new(10);
1130        let _strong = Primary::clone_strong(&primary);
1131        drop(primary);
1132    }
1133
1134    #[test]
1135    fn unwrap_with_notifier() {
1136        let primary = Primary::new(10);
1137        let strong = Primary::clone_strong(&primary);
1138        let notifier = ArcNotifier::new();
1139        Primary::unwrap_with_notifier(primary, notifier.clone());
1140        // Strong reference is still alive.
1141        assert_eq!(notifier.take(), None);
1142        core::mem::drop(strong);
1143        assert_eq!(notifier.take(), Some(10));
1144    }
1145
1146    #[test]
1147    fn unwrap_or_notify_with_immediate() {
1148        let primary = Primary::new(10);
1149        let result = Primary::unwrap_or_notify_with::<ArcNotifier<_>, (), _>(primary, || {
1150            panic!("should not try to create notifier")
1151        });
1152        assert_eq!(result, Ok(10));
1153    }
1154
1155    #[test]
1156    fn unwrap_or_notify_with_deferred() {
1157        let primary = Primary::new(10);
1158        let strong = Primary::clone_strong(&primary);
1159        let result = Primary::unwrap_or_notify_with(primary, || {
1160            let notifier = ArcNotifier::new();
1161            (notifier.clone(), notifier)
1162        });
1163        let notifier = result.unwrap_err();
1164        assert_eq!(notifier.take(), None);
1165        core::mem::drop(strong);
1166        assert_eq!(notifier.take(), Some(10));
1167    }
1168
1169    #[test]
1170    fn map_notifier() {
1171        let primary = Primary::new(10);
1172        let notifier = ArcNotifier::new();
1173        let map_notifier = MapNotifier::new(notifier.clone(), |data| (data, data + 1));
1174        Primary::unwrap_with_notifier(primary, map_notifier);
1175        assert_eq!(notifier.take(), Some((10, 11)));
1176    }
1177
1178    #[test]
1179    fn new_cyclic() {
1180        #[derive(Debug)]
1181        struct Data {
1182            value: i32,
1183            weak: Weak<Data>,
1184        }
1185
1186        let primary = Primary::new_cyclic(|weak| Data { value: 2, weak });
1187        assert_eq!(primary.value, 2);
1188        let strong = primary.weak.upgrade().unwrap();
1189        assert_eq!(strong.value, 2);
1190        assert!(Primary::ptr_eq(&primary, &strong));
1191    }
1192
1193    macro_rules! assert_debug_id_eq {
1194        ($id1:expr, $id2:expr) => {
1195            assert_eq!(alloc::format!("{:?}", $id1), alloc::format!("{:?}", $id2))
1196        };
1197    }
1198    macro_rules! assert_debug_id_ne {
1199        ($id1:expr, $id2:expr) => {
1200            assert_ne!(alloc::format!("{:?}", $id1), alloc::format!("{:?}", $id2))
1201        };
1202    }
1203
1204    #[test]
1205    fn debug_ids_are_stable() {
1206        // Verify that transforming a given RC doesn't change it's debug_id.
1207        let primary = Primary::new(1);
1208        let strong = Primary::clone_strong(&primary);
1209        let weak_p = Primary::downgrade(&primary);
1210        let weak_s = Strong::downgrade(&strong);
1211        let weak_c = weak_p.clone();
1212        assert_debug_id_eq!(&primary.debug_id(), &strong.debug_id());
1213        assert_debug_id_eq!(&primary.debug_id(), &weak_p.debug_id());
1214        assert_debug_id_eq!(&primary.debug_id(), &weak_s.debug_id());
1215        assert_debug_id_eq!(&primary.debug_id(), &weak_c.debug_id());
1216    }
1217
1218    #[test]
1219    fn debug_ids_are_unique() {
1220        // Verify that RCs to different data have different debug_ids.
1221        let primary1 = Primary::new(1);
1222        let primary2 = Primary::new(1);
1223        assert_debug_id_ne!(&primary1.debug_id(), &primary2.debug_id());
1224
1225        // Verify that dropping an RC does not allow it's debug_id to be reused.
1226        let id1 = format!("{:?}", primary1.debug_id());
1227        std::mem::drop(primary1);
1228        let primary3 = Primary::new(1);
1229        assert_ne!(id1, format!("{:?}", primary3.debug_id()));
1230    }
1231}