thread_aware 0.7.2

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.

mod clone_fn;
mod factory;
pub mod storage;

mod builtin;
#[cfg_attr(coverage_nightly, coverage(off))]
#[cfg(test)]
mod tests;

use std::cmp::Ordering;
use std::hash::Hasher;
use std::ops::Deref;
use std::sync::{self, RwLock};

pub use builtin::{PerCore, PerNuma, PerProcess};
pub use storage::{Storage, Strategy};

use crate::ThreadAware;
use crate::affinity::Affinity;
use crate::cell::factory::Factory;
use crate::closure::{ErasedClosureOnce, ThreadAwareFnOnce, closure_once};

/// Adapter that wraps a `ThreadAwareFnOnce<T>` to produce `Box<T>` instead.
struct BoxedRelocate<F>(F);

impl<F: Clone> Clone for BoxedRelocate<F> {
    fn clone(&self) -> Self {
        Self(self.0.clone())
    }
}

impl<F: ThreadAware> ThreadAware for BoxedRelocate<F> {
    fn relocate(&mut self, source: Option<Affinity>, destination: Affinity) {
        self.0.relocate(source, destination);
    }
}

impl<T, F: ThreadAwareFnOnce<T>> ThreadAwareFnOnce<Box<T>> for BoxedRelocate<F> {
    fn call_once(self) -> Box<T> {
        Box::new(self.0.call_once())
    }
}

/// Transferable reference counted type.
///
/// This type works like a per-affinity (per-thread) [`sync::Arc`]. Each affinity gets a unique value that is shared by clones
/// of the `Arc`, but the [`ThreadAware`] implementation ensures that when moving to another affinity, the resulting
/// `Arc` will point to the value in the destination affinity. See [`new`](`Arc::new`) for information on constructing instances.
///
/// `ThreadAware` of different clones of the `Arc` result in "deduplication" in the destination affinity. The following
/// example demonstrates this using the counter implemented in the documentation for the [`ThreadAware`] trait.
///
/// ```rust
/// # use thread_aware::{Arc, ThreadAware, PerCore};
/// # use thread_aware::affinity::*;
/// # use std::sync::atomic::{AtomicI32, Ordering};
/// # let affinities = pinned_affinities(&[2]);
/// # let affinity1 = Some(affinities[0]);
/// # let affinity2 = affinities[1];
/// # #[derive(Clone)]
/// # struct Counter {
/// #     value: std::sync::Arc<AtomicI32>,
/// # }
/// #
/// # impl Counter {
/// #     fn new() -> Self {
/// #         Self {
/// #             value: std::sync::Arc::new(AtomicI32::new(0)),
/// #         }
/// #     }
/// #
/// #     fn increment_by(&self, v: i32) {
/// #         self.value.fetch_add(v, Ordering::AcqRel);
/// #     }
/// #
/// #     fn value(&self) -> i32 {
/// #         self.value.load(Ordering::Acquire)
/// #     }
/// # }
/// #
/// # impl ThreadAware for Counter {
/// #     fn relocate(&mut self, _source: Option<Affinity>, _destination: Affinity) {
/// #         // Initialize a new value in the destination affinity independent
/// #         // of the source affinity.
/// #         self.value = std::sync::Arc::new(AtomicI32::new(0));
/// #     }
/// # }
///
/// let mut arc_affinity1 = Arc::<_, PerCore>::new(Counter::new);
/// let arc_affinity1_clone = arc_affinity1.clone();
///
/// arc_affinity1.increment_by(42);
/// assert_eq!(arc_affinity1.value(), 42);
///
/// arc_affinity1.relocate(affinity1, affinity2);
/// assert_eq!(arc_affinity1.value(), 0);
/// assert_eq!(arc_affinity1_clone.value(), 42);
///
/// arc_affinity1.increment_by(11);
/// let mut arc_affinity2_clone = arc_affinity1_clone;
/// arc_affinity2_clone.relocate(affinity1, affinity2);
/// assert_eq!(arc_affinity2_clone.value(), 11);
/// ```
#[derive(Debug)]
pub struct Arc<T: ?Sized, S: Strategy> {
    storage: sync::Arc<RwLock<Storage<sync::Arc<T>, S>>>,
    value: sync::Arc<T>,
    factory: Factory<T>,
}

impl<T: PartialEq, S: Strategy> PartialEq for Arc<T, S> {
    fn eq(&self, other: &Self) -> bool {
        self.value == other.value
    }
}

impl<T: Eq, S: Strategy> Eq for Arc<T, S> {}

impl<T: std::hash::Hash + ?Sized, S: Strategy> std::hash::Hash for Arc<T, S> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        self.value.hash(state);
    }
}

impl<T: Ord, S: Strategy> Ord for Arc<T, S> {
    fn cmp(&self, other: &Self) -> Ordering {
        self.value.cmp(&other.value)
    }
}

impl<T: PartialOrd, S: Strategy> PartialOrd for Arc<T, S> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.value.partial_cmp(&other.value)
    }
}

impl<T: ?Sized, S: Strategy> Clone for Arc<T, S> {
    fn clone(&self) -> Self {
        Self {
            storage: sync::Arc::clone(&self.storage),
            value: sync::Arc::clone(&self.value),
            factory: self.factory.clone(),
        }
    }
}

impl<T: ?Sized, S: Strategy> Deref for Arc<T, S> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        &self.value
    }
}

impl<T, S> Arc<T, S>
where
    T: Send + 'static,
    S: Strategy,
{
    /// Creates a new `Arc` with the given value and strategy.
    ///
    /// This variant takes a zero-argument constructor function (`fn() -> T`).
    /// The constructor is invoked lazily and independently for each
    /// processor the first time a `PerCore` is materialized on that processor (i.e. on
    /// the first transfer into that processor). This guarantees that every processor obtains its own
    /// freshly created `T` without requiring `T: Clone` or `T: ThreadAware`.
    ///
    /// Requirements:
    /// * `T` must be `Send + 'static` so it can live in the processor storage.
    /// * The provided function must be pure with respect to per-processor isolation (it should not
    ///   leak references into other processors). Any captured state should therefore be provided via
    ///   globally shareable mechanisms or prefer [`new_with`](Self::new_with) if you need to
    ///   capture data that itself implements [`ThreadAware`].
    ///
    /// When transferring to another affinity which doesn't yet contain a value, the constructor is
    /// called in the destination affinity to create a brand new instance.
    ///
    /// For example, the counter type we implemented in the documentation for [`ThreadAware`] trait
    /// can be used with `new` by passing the constructor function (note the absence of `()`):
    ///
    /// ```rust
    /// # use thread_aware::{Arc, ThreadAware, PerCore};
    /// # use thread_aware::affinity::*;
    /// # use std::sync::atomic::{AtomicI32, Ordering};
    /// # use std::sync;
    /// # #[derive(Clone)]
    /// # struct Counter {
    /// #     value: sync::Arc<AtomicI32>,
    /// # }
    /// #
    /// # impl Counter {
    /// #     fn new() -> Self {
    /// #         Self {
    /// #             value: sync::Arc::new(AtomicI32::new(0)),
    /// #         }
    /// #     }
    /// #
    /// #     fn increment_by(&self, v: i32) {
    /// #         self.value.fetch_add(v, Ordering::AcqRel);
    /// #     }
    /// #
    /// #     fn value(&self) -> i32 {
    /// #         self.value.load(Ordering::Acquire)
    /// #     }
    /// # }
    /// # impl ThreadAware for Counter {
    /// #     fn relocate(&mut self, _source: Option<Affinity>, _destination: Affinity) {
    /// #         // Initialize a new value in the destination affinity independent
    /// #         // of the source affinity.
    /// #         self.value = sync::Arc::new(AtomicI32::new(0));
    /// #     }
    /// # }
    ///
    /// let container = Arc::<_, PerCore>::new(Counter::new);
    /// let container_clone = container.clone();
    /// container.increment_by(42);
    /// assert_eq!(container.value(), 42);
    /// assert_eq!(container_clone.value(), 42);
    /// ```
    pub fn new(ctor: fn() -> T) -> Self {
        // We wrap the function pointer in a tiny ThreadAwareFnOnce implementation that
        // recreates the value independently for each affinity.
        struct Ctor<T> {
            f: fn() -> T,
        }

        impl<T> Clone for Ctor<T> {
            fn clone(&self) -> Self {
                Self { f: self.f }
            }
        }

        impl<T> ThreadAware for Ctor<T> {
            fn relocate(&mut self, _source: Option<Affinity>, _destination: Affinity) {}
        }

        impl<T> ThreadAwareFnOnce<Box<T>> for Ctor<T> {
            fn call_once(self) -> Box<T> {
                Box::new((self.f)())
            }
        }

        // Use Self::with_closure_boxed to ensure Factory::Closure path.
        Self::with_closure_boxed(Ctor { f: ctor })
    }
}

impl<T, S> Arc<T, S>
where
    T: Send + 'static + ?Sized,
    S: Strategy,
{
    /// Creates a new `Arc` with a constructor that returns `Box<T>`.
    ///
    /// This is the `?Sized`-compatible version of [`new`](Self::new). Use this when `T` is a
    /// trait object (e.g., `dyn Trait`) or other unsized type. The constructor produces a
    /// `Box<T>` which is then stored behind a [`sync::Arc`].
    ///
    /// ```rust
    /// # use thread_aware::{Arc, ThreadAware, PerCore};
    /// let arc = Arc::<dyn ThreadAware, PerCore>::new_boxed(|| Box::new(42u32));
    /// ```
    pub fn new_boxed(ctor: fn() -> Box<T>) -> Self {
        struct Ctor<T: ?Sized> {
            f: fn() -> Box<T>,
        }

        impl<T: ?Sized> Clone for Ctor<T> {
            fn clone(&self) -> Self {
                Self { f: self.f }
            }
        }

        impl<T: ?Sized> ThreadAware for Ctor<T> {
            fn relocate(&mut self, _source: Option<Affinity>, _destination: Affinity) {}
        }

        impl<T: ?Sized> ThreadAwareFnOnce<Box<T>> for Ctor<T> {
            fn call_once(self) -> Box<T> {
                (self.f)()
            }
        }

        Self::with_closure_boxed(Ctor { f: ctor })
    }
}

impl<T, S> Arc<T, S>
where
    T: 'static,
    S: Strategy,
{
    /// Creates a new `Arc` with a closure that will be called once per-processor to create the inner value.
    ///
    /// The closure only gets called once for each processor, and it's called only when a `Arc` is actually transferred
    /// to another processor. The closure behaves like a `ThreadAwareFnOnce` to ensure it captures only values that are safe to
    /// transfer themselves.
    ///
    /// This function can be used to create an `Arc` of a type that itself doesn't implement [`ThreadAware`] because
    /// we can ensure that each affinity will get its own, independently-initialized value:
    ///
    /// ```rust
    /// # use std::sync::{self, Mutex};
    /// # use thread_aware::{Arc, PerCore};
    /// struct MyStruct {
    ///     inner: sync::Arc<Mutex<i32>>,
    /// }
    ///
    /// impl MyStruct {
    ///     fn new() -> Self {
    ///         Self {
    ///             inner: sync::Arc::new(Mutex::new(0)),
    ///         }
    ///     }
    /// }
    ///
    /// let container = Arc::<_, PerCore>::new_with((), |_| MyStruct::new());
    /// ```
    ///
    /// The constructor can depend on other values that implement [`ThreadAware`] (this example uses the Counter
    /// defined in [`ThreadAware`] documentation):
    ///
    /// ```rust
    /// # use thread_aware::{ThreadAware, Arc, PerCore};
    /// # use thread_aware::affinity::*;
    /// # use std::sync::atomic::{AtomicI32, Ordering};
    /// # use std::sync;
    /// # #[derive(Clone)]
    /// # struct Counter {
    /// #     value: sync::Arc<AtomicI32>,
    /// # }
    /// #
    /// # impl Counter {
    /// #     fn new() -> Self {
    /// #         Self {
    /// #             value: sync::Arc::new(AtomicI32::new(0)),
    /// #         }
    /// #     }
    /// #
    /// #     fn increment_by(&self, v: i32) {
    /// #         self.value.fetch_add(v, Ordering::AcqRel);
    /// #     }
    /// #
    /// #     fn value(&self) -> i32 {
    /// #         self.value.load(Ordering::Acquire)
    /// #     }
    /// # }
    /// #
    /// # impl ThreadAware for Counter {
    /// #     fn relocate(&mut self, _source: Option<Affinity>, _destination: Affinity) {
    /// #         // Initialize a new value in the destination affinity independent
    /// #         // of the source affinity.
    /// #         self.value = sync::Arc::new(AtomicI32::new(0));
    /// #     }
    /// # }
    ///
    /// struct MyStruct;
    ///
    /// impl MyStruct {
    ///     fn new(value: i32) -> Self {
    ///         Self
    ///     }
    /// }
    ///
    /// let counter = Counter::new();
    /// let container = Arc::<_, PerCore>::new_with(counter, |counter| MyStruct::new(counter.value()));
    /// ```
    pub fn new_with<D>(data: D, f: fn(D) -> T) -> Self
    where
        D: ThreadAware + Send + Sync + Clone + 'static,
    {
        Self::with_closure_boxed(BoxedRelocate(closure_once(data, f)))
    }
}

impl<T, S: Strategy> Arc<T, S>
where
    T: ThreadAware + Clone + 'static + Send,
{
    /// Creates a new `Arc` with the given value.
    ///
    /// The value must implement [`ThreadAware`] and [`Clone`]. When transferring to another affinity
    /// which doesn't yet contain a value, a new value is created by cloning the value in current
    /// affinity and transferring it to the new affinity.
    ///
    /// For example, the counter type we implemented in the documentation for the [`ThreadAware`] trait
    /// can be used with new.
    #[cfg(test)]
    pub(crate) fn with_value(value: T) -> Self {
        let value = sync::Arc::new(value);

        Self {
            storage: sync::Arc::new(RwLock::new(storage::Storage::new())),
            value,
            factory: Factory::Data(|data: &T, source, destination| {
                let mut data = data.clone();
                data.relocate(source, destination);
                Box::new(data)
            }),
        }
    }
}

impl<T, S: Strategy> Arc<T, S>
where
    T: Clone + 'static + Send,
{
    /// Creates a new `Arc` with the given value.
    ///
    /// The value must implement [`Clone`]. When transferring to another affinity
    /// which doesn't yet contain a value, a new value is created by cloning the value in current
    /// affinity and transferring it to the new affinity.
    ///
    /// This is useful for types that do not implement [`ThreadAware`]. In such cases, the same value
    /// is cloned for each affinity without any relocation logic.
    ///
    /// For example, the counter type we implemented in the documentation for [`ThreadAware`] trait
    /// can be used with new:
    ///
    /// ```rust
    /// # use thread_aware::{Arc, PerCore};
    /// # use std::sync::atomic::{AtomicI32, Ordering};
    /// # use std::sync;
    /// # #[derive(Clone)]
    /// # struct Counter {
    /// #     value: sync::Arc<AtomicI32>,
    /// # }
    /// #
    /// # impl Counter {
    /// #     fn new() -> Self {
    /// #         Self {
    /// #             value: sync::Arc::new(AtomicI32::new(0)),
    /// #         }
    /// #     }
    /// #
    /// #     fn increment_by(&self, v: i32) {
    /// #         self.value.fetch_add(v, Ordering::AcqRel);
    /// #     }
    /// #
    /// #     fn value(&self) -> i32 {
    /// #         self.value.load(Ordering::Acquire)
    /// #     }
    /// # }
    ///
    /// let arc = Arc::<_, PerCore>::new(Counter::new);
    /// let arc_clone = arc.clone();
    /// arc.increment_by(42);
    /// assert_eq!(arc.value(), 42);
    /// assert_eq!(arc_clone.value(), 42);
    /// ```
    pub fn from_unaware(value: T) -> Self {
        let value = sync::Arc::new(value);

        Self {
            storage: sync::Arc::new(RwLock::new(storage::Storage::new())),
            value,
            factory: Factory::Data(|data: &T, _source, _destination| Box::new(data.clone())),
        }
    }
}

impl<T, S: Strategy> Arc<T, S>
where
    T: ThreadAware + 'static + ?Sized,
{
    /// Creates a new `Arc` from a value and a clone function, supporting trait objects.
    ///
    /// The object passed will be kept, and serves as the template for all subsequent clones
    /// that happen on relocation. A clone is also performed for the initial `Arc`.
    ///
    /// The `clone_fn` receives `&V` (the concrete type) and returns `Box<T>`, enabling
    /// use with `dyn Trait` where `Clone` is not object-safe. Each clone is
    /// [`relocate`](ThreadAware::relocate) to its target affinity.
    ///
    /// ```rust
    /// # use thread_aware::{Arc, PerCore, ThreadAware};
    /// # #[derive(Clone)]
    /// # struct Foo(u32);
    /// # impl ThreadAware for Foo {
    /// #     fn relocate(&mut self, _: Option<thread_aware::affinity::Affinity>, _: thread_aware::affinity::Affinity) {}
    /// # }
    /// trait MyPlugin: ThreadAware {}
    /// impl MyPlugin for Foo {}
    ///
    /// let arc = Arc::<dyn MyPlugin, PerCore>::with_clone_fn(Foo(42), |v: &Foo| Box::new(v.clone()));
    /// ```
    pub fn with_clone_fn<V: Send + Sync + 'static>(value: V, clone_fn: fn(&V) -> Box<T>) -> Self {
        // In a canonical case, we might have `V = u32`, `T = dyn Foo`, and `clone_fn = |&u32| -> Box<dyn Foo>`.
        let erased = clone_fn::ErasedCloneFn::new(value, clone_fn);
        let value = sync::Arc::clone(erased.arc());

        Self {
            storage: sync::Arc::new(RwLock::new(storage::Storage::new())),
            value,
            factory: Factory::ErasedCloneFn(erased),
        }
    }
}

impl<T, S: Strategy> Arc<T, S>
where
    T: 'static + ?Sized,
{
    /// Creates a new `Arc` with a closure that produces `Box<T>`, called once per-affinity to create the inner value.
    ///
    /// The closure only gets called once for each affinity, and it's called only when an `Arc` is actually transferred
    /// to another affinity. The closure is a [`ThreadAwareFnOnce`] to ensure it captures only values that are safe to
    /// transfer themselves.
    pub(crate) fn with_closure_boxed<F>(closure: F) -> Self
    where
        F: ThreadAwareFnOnce<Box<T>> + Clone + ThreadAware + 'static + Send + Sync,
    {
        let value = sync::Arc::from(closure.clone().call_once());

        Self {
            storage: sync::Arc::new(RwLock::new(storage::Storage::new())),
            value,
            factory: Factory::Closure(sync::Arc::new(ErasedClosureOnce::new(closure)), None),
        }
    }

    /// Creates a new `Arc` from the given storage and the current affinity.
    ///
    /// If the resulting `Arc` is transferred to an affinity which does not have data in the storage,
    /// it will behave like a [`sync::Arc`].
    ///
    /// # Panics
    /// This may panic if the storage does not contain data for the current affinity.
    pub fn from_storage(storage: sync::Arc<RwLock<Storage<sync::Arc<T>, S>>>, current_affinity: Affinity) -> Self {
        let value = storage
            .read()
            .expect("Failed to acquire read lock")
            .get_clone(current_affinity)
            .expect("No data found for the current affinity");

        Self {
            storage,
            value,
            factory: Factory::Manual,
        }
    }
}

impl<T, S: Strategy> Arc<T, S> {
    /// Gets the number of strong references to the value in the current thread/affinity.
    ///
    /// This method returns the strong reference count for the underlying [`sync::Arc`]
    /// that holds the value for the current affinity, excluding any internal references
    /// held by the storage for deduplication purposes. Each affinity maintains its own
    /// separate value with its own reference count.
    ///
    /// # Examples
    ///
    /// ```
    /// use thread_aware::{Arc, PerCore};
    ///
    /// let arc = Arc::<_, PerCore>::new(|| 42);
    /// assert_eq!(Arc::strong_count(&arc), 1);
    ///
    /// let arc2 = arc.clone();
    /// assert_eq!(Arc::strong_count(&arc), 2);
    /// assert_eq!(Arc::strong_count(&arc2), 2);
    /// ```
    #[must_use]
    #[expect(clippy::missing_panics_doc, reason = "this code only panics when the lock is poisoned")]
    pub fn strong_count(this: &Self) -> usize {
        let raw = sync::Arc::strong_count(&this.value);
        let guard = this.storage.read().expect("Failed to acquire read lock");
        let internal = guard.count_where(|stored| sync::Arc::ptr_eq(stored, &this.value));
        raw - internal
    }

    /// Converts the `Arc<T, S>` into an `sync::Arc<T>`.
    #[must_use]
    pub fn into_arc(self) -> sync::Arc<T> {
        self.value
    }
}

impl<T: Send + Sync + ?Sized, S: Strategy + Send + Sync> ThreadAware for Arc<T, S> {
    fn relocate(&mut self, source: Option<Affinity>, destination: Affinity) {
        let mut guard = self.storage.write().expect("Failed to acquire write lock");

        if let Some(value) = guard.get_clone(destination) {
            self.value = value;
        } else {
            // We need to transfer or recreate the data
            let (data, new_factory) = match &self.factory {
                // We can use the closure to create new data
                Factory::Closure(factory, factory_source_affinity) => {
                    let mut factory_clone = (**factory).clone();

                    // In case factory source is stored in factory, use that - it means we already transferred the factory
                    // once, so we know the original source affinity. Otherwise, use source as that means this is the first
                    // time we're transferring the Arc, so source is the source affinity of the factory as well.
                    let factory_source = factory_source_affinity.or(source);

                    factory_clone.relocate(factory_source, destination);
                    (
                        sync::Arc::from(factory_clone.call_once()),
                        Factory::Closure(sync::Arc::clone(factory), factory_source),
                    )
                }

                // We can clone and transfer the data
                Factory::Data(factory) => (sync::Arc::from(factory(&self.value, source, destination)), self.factory.clone()),

                // We can clone the data and closure it
                Factory::ErasedCloneFn(erased) => {
                    let cloned = erased.clone_and_relocate(source, destination);
                    (cloned, self.factory.clone())
                }

                Factory::Manual => {
                    // If we are in manual mode, we just clone the data
                    // This effectively makes it behave like `sync::Arc<T>`
                    (sync::Arc::clone(&self.value), self.factory.clone())
                }
            };

            let old_value = std::mem::replace(&mut self.value, data);

            let old_data = guard.replace(destination, sync::Arc::<T>::clone(&self.value));
            assert!(
                old_data.is_none(),
                "Data already exists for the destination affinity. This should be unreachable due to the the early write lock."
            );

            if let Some(source) = source {
                // Only restore the value to the source slot when source and destination differ.
                // If they are the same slot, the replacement above already stored the new value
                // there; overwriting it here would corrupt storage with the stale pre-relocation value.
                if source != destination {
                    guard.replace(source, old_value);
                }
            }

            self.factory = new_factory;
        }

        drop(guard);
    }
}