Re: [RFC v2 04/13] rust: sync: atomic: Add generic atomics

[Boqun Feng <boqun.feng@gmail.com>]

On Thu, Dec 12, 2024 at 11:57:07AM +0100, Alice Ryhl wrote:
[...]
> > diff --git a/rust/kernel/sync/atomic/generic.rs b/rust/kernel/sync/atomic/generic.rs
> > new file mode 100644
> > index 000000000000..204da38e2691
> > --- /dev/null
> > +++ b/rust/kernel/sync/atomic/generic.rs
> > @@ -0,0 +1,253 @@
> > +// SPDX-License-Identifier: GPL-2.0
> > +
> > +//! Generic atomic primitives.
> > +
> > +use super::ops::*;
> > +use super::ordering::*;
> > +use crate::types::Opaque;
> > +
> > +/// A generic atomic variable.
> > +///
> > +/// `T` must impl [`AllowAtomic`], that is, an [`AtomicImpl`] has to be chosen.
> > +///
> > +/// # Invariants
> > +///
> > +/// Doing an atomic operation while holding a reference of [`Self`] won't cause a data race, this
> > +/// is guaranteed by the safety requirement of [`Self::from_ptr`] and the extra safety requirement
> > +/// of the usage on pointers returned by [`Self::as_ptr`].
> > +#[repr(transparent)]
> > +pub struct Atomic<T: AllowAtomic>(Opaque<T>);
> > +
> > +// SAFETY: `Atomic<T>` is safe to share among execution contexts because all accesses are atomic.
> > +unsafe impl<T: AllowAtomic> Sync for Atomic<T> {}
> 
> Surely it should also be Send?
> 

It's `Send` here because `Opaque<T>` is `Send` when `T` is `Send`. And
in patch #9, I changed the definition of `AllowAtomic`, which is not a
subtrait of `Send` anymore, and an `impl Send` block was added there.

> > +/// Atomics that support basic atomic operations.
> > +///
> > +/// TODO: Unless the `impl` is a `#[repr(transparet)]` new type of an existing [`AllowAtomic`], the
> > +/// impl block should be only done in atomic mod. And currently only basic integer types can
> > +/// implement this trait in atomic mod.
> 
> What's up with this TODO? Can't you just write an appropriate safety
> requirement?
> 

Because the limited scope of types that allows atomic is an artificial
choice, i.e. we want to start with a limited number of types and make
forward progress, and the types that we don't want to support atomics
for now are not because of safety reasons, but more of a lack of
users/motivations. So I don't think this is something we should use
safety requirement to describe.

> > +/// # Safety
> > +///
> > +/// [`Self`] must have the same size and alignment as [`Self::Repr`].
> > +pub unsafe trait AllowAtomic: Sized + Send + Copy {
> > +    /// The backing atomic implementation type.
> > +    type Repr: AtomicImpl;
> > +
> > +    /// Converts into a [`Self::Repr`].
> > +    fn into_repr(self) -> Self::Repr;
> > +
> > +    /// Converts from a [`Self::Repr`].
> > +    fn from_repr(repr: Self::Repr) -> Self;
> 
> What do you need these methods for?
> 

Converting a `AtomicImpl` value (currently only `i32` and `i64`) to a
`AllowAtomic` value without using transmute in `impl` block of
`Atomic<T>`. Any better idea?

Regards,
Boqun

> > +}
> > +
> > +// SAFETY: `T::Repr` is `Self` (i.e. `T`), so they have the same size and alignment.
> > +unsafe impl<T: AtomicImpl> AllowAtomic for T {
> > +    type Repr = Self;
> > +
> > +    fn into_repr(self) -> Self::Repr {
> > +        self
> > +    }
> > +
> > +    fn from_repr(repr: Self::Repr) -> Self {
> > +        repr
> > +    }
> > +}
> > +
> > +impl<T: AllowAtomic> Atomic<T> {
> > +    /// Creates a new atomic.
> > +    pub const fn new(v: T) -> Self {
> > +        Self(Opaque::new(v))
> > +    }
> > +
> > +    /// Creates a reference to [`Self`] from a pointer.
> > +    ///
> > +    /// # Safety
> > +    ///
> > +    /// - `ptr` has to be a valid pointer.
> > +    /// - `ptr` has to be valid for both reads and writes for the whole lifetime `'a`.
> > +    /// - For the whole lifetime of '`a`, other accesses to the object cannot cause data races
> > +    ///   (defined by [`LKMM`]) against atomic operations on the returned reference.
> > +    ///
> > +    /// [`LKMM`]: srctree/tools/memory-model
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// Using [`Atomic::from_ptr()`] combined with [`Atomic::load()`] or [`Atomic::store()`] can
> > +    /// achieve the same functionality as `READ_ONCE()`/`smp_load_acquire()` or
> > +    /// `WRITE_ONCE()`/`smp_store_release()` in C side:
> > +    ///
> > +    /// ```rust
> > +    /// # use kernel::types::Opaque;
> > +    /// use kernel::sync::atomic::{Atomic, Relaxed, Release};
> > +    ///
> > +    /// // Assume there is a C struct `Foo`.
> > +    /// mod cbindings {
> > +    ///     #[repr(C)]
> > +    ///     pub(crate) struct foo { pub(crate) a: i32, pub(crate) b: i32 }
> > +    /// }
> > +    ///
> > +    /// let tmp = Opaque::new(cbindings::foo { a: 1, b: 2});
> > +    ///
> > +    /// // struct foo *foo_ptr = ..;
> > +    /// let foo_ptr = tmp.get();
> > +    ///
> > +    /// // SAFETY: `foo_ptr` is a valid pointer, and `.a` is inbound.
> > +    /// let foo_a_ptr = unsafe { core::ptr::addr_of_mut!((*foo_ptr).a) };
> > +    ///
> > +    /// // a = READ_ONCE(foo_ptr->a);
> > +    /// //
> > +    /// // SAFETY: `foo_a_ptr` is a valid pointer for read, and all accesses on it is atomic, so no
> > +    /// // data race.
> > +    /// let a = unsafe { Atomic::from_ptr(foo_a_ptr) }.load(Relaxed);
> > +    /// # assert_eq!(a, 1);
> > +    ///
> > +    /// // smp_store_release(&foo_ptr->a, 2);
> > +    /// //
> > +    /// // SAFETY: `foo_a_ptr` is a valid pointer for write, and all accesses on it is atomic, so no
> > +    /// // data race.
> > +    /// unsafe { Atomic::from_ptr(foo_a_ptr) }.store(2, Release);
> > +    /// ```
> > +    ///
> > +    /// However, this should be only used when communicating with C side or manipulating a C struct.
> > +    pub unsafe fn from_ptr<'a>(ptr: *mut T) -> &'a Self
> > +    where
> > +        T: Sync,
> > +    {
> > +        // CAST: `T` is transparent to `Atomic<T>`.
> > +        // SAFETY: Per function safety requirement, `ptr` is a valid pointer and the object will
> > +        // live long enough. It's safe to return a `&Atomic<T>` because function safety requirement
> > +        // guarantees other accesses won't cause data races.
> > +        unsafe { &*ptr.cast::<Self>() }
> > +    }
> > +
> > +    /// Returns a pointer to the underlying atomic variable.
> > +    ///
> > +    /// Extra safety requirement on using the return pointer: the operations done via the pointer
> > +    /// cannot cause data races defined by [`LKMM`].
> > +    ///
> > +    /// [`LKMM`]: srctree/tools/memory-model
> > +    pub const fn as_ptr(&self) -> *mut T {
> > +        self.0.get()
> > +    }
> > +
> > +    /// Returns a mutable reference to the underlying atomic variable.
> > +    ///
> > +    /// This is safe because the mutable reference of the atomic variable guarantees the exclusive
> > +    /// access.
> > +    pub fn get_mut(&mut self) -> &mut T {
> > +        // SAFETY: `self.as_ptr()` is a valid pointer to `T`, and the object has already been
> > +        // initialized. `&mut self` guarantees the exclusive access, so it's safe to reborrow
> > +        // mutably.
> > +        unsafe { &mut *self.as_ptr() }
> > +    }
> > +}
> > +
> > +impl<T: AllowAtomic> Atomic<T>
> > +where
> > +    T::Repr: AtomicHasBasicOps,
> > +{
> > +    /// Loads the value from the atomic variable.
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// Simple usages:
> > +    ///
> > +    /// ```rust
> > +    /// use kernel::sync::atomic::{Atomic, Relaxed};
> > +    ///
> > +    /// let x = Atomic::new(42i32);
> > +    ///
> > +    /// assert_eq!(42, x.load(Relaxed));
> > +    ///
> > +    /// let x = Atomic::new(42i64);
> > +    ///
> > +    /// assert_eq!(42, x.load(Relaxed));
> > +    /// ```
> > +    ///
> > +    /// Customized new types in [`Atomic`]:
> > +    ///
> > +    /// ```rust
> > +    /// use kernel::sync::atomic::{generic::AllowAtomic, Atomic, Relaxed};
> > +    ///
> > +    /// #[derive(Clone, Copy)]
> > +    /// #[repr(transparent)]
> > +    /// struct NewType(u32);
> > +    ///
> > +    /// // SAFETY: `NewType` is transparent to `u32`, which has the same size and alignment as
> > +    /// // `i32`.
> > +    /// unsafe impl AllowAtomic for NewType {
> > +    ///     type Repr = i32;
> > +    ///
> > +    ///     fn into_repr(self) -> Self::Repr {
> > +    ///         self.0 as i32
> > +    ///     }
> > +    ///
> > +    ///     fn from_repr(repr: Self::Repr) -> Self {
> > +    ///         NewType(repr as u32)
> > +    ///     }
> > +    /// }
> > +    ///
> > +    /// let n = Atomic::new(NewType(0));
> > +    ///
> > +    /// assert_eq!(0, n.load(Relaxed).0);
> > +    /// ```
> > +    #[inline(always)]
> > +    pub fn load<Ordering: AcquireOrRelaxed>(&self, _: Ordering) -> T {
> > +        let a = self.as_ptr().cast::<T::Repr>();
> > +
> > +        // SAFETY:
> > +        // - For calling the atomic_read*() function:
> > +        //   - `self.as_ptr()` is a valid pointer, and per the safety requirement of `AllocAtomic`,
> > +        //      a `*mut T` is a valid `*mut T::Repr`. Therefore `a` is a valid pointer,
> > +        //   - per the type invariants, the following atomic operation won't cause data races.
> > +        // - For extra safety requirement of usage on pointers returned by `self.as_ptr():
> > +        //   - atomic operations are used here.
> > +        let v = unsafe {
> > +            if Ordering::IS_RELAXED {
> > +                T::Repr::atomic_read(a)
> > +            } else {
> > +                T::Repr::atomic_read_acquire(a)
> > +            }
> > +        };
> > +
> > +        T::from_repr(v)
> > +    }
> > +
> > +    /// Stores a value to the atomic variable.
> > +    ///
> > +    /// # Examples
> > +    ///
> > +    /// ```rust
> > +    /// use kernel::sync::atomic::{Atomic, Relaxed};
> > +    ///
> > +    /// let x = Atomic::new(42i32);
> > +    ///
> > +    /// assert_eq!(42, x.load(Relaxed));
> > +    ///
> > +    /// x.store(43, Relaxed);
> > +    ///
> > +    /// assert_eq!(43, x.load(Relaxed));
> > +    /// ```
> > +    ///
> > +    #[inline(always)]
> > +    pub fn store<Ordering: ReleaseOrRelaxed>(&self, v: T, _: Ordering) {
> > +        let v = T::into_repr(v);
> > +        let a = self.as_ptr().cast::<T::Repr>();
> > +
> > +        // SAFETY:
> > +        // - For calling the atomic_set*() function:
> > +        //   - `self.as_ptr()` is a valid pointer, and per the safety requirement of `AllocAtomic`,
> > +        //      a `*mut T` is a valid `*mut T::Repr`. Therefore `a` is a valid pointer,
> > +        //   - per the type invariants, the following atomic operation won't cause data races.
> > +        // - For extra safety requirement of usage on pointers returned by `self.as_ptr():
> > +        //   - atomic operations are used here.
> > +        unsafe {
> > +            if Ordering::IS_RELAXED {
> > +                T::Repr::atomic_set(a, v)
> > +            } else {
> > +                T::Repr::atomic_set_release(a, v)
> > +            }
> > +        };
> > +    }
> > +}
> > --
> > 2.45.2
> >