Re: [PATCH v16 1/7] rust: sync: add `SetOnce`

["Benno Lossin" <lossin@kernel.org>]

On Wed Jul 9, 2025 at 7:52 PM CEST, Andreas Hindborg wrote:
> Introduce the `SetOnce` type, a container that can only be written once.
> The container uses an internal atomic to synchronize writes to the internal
> value.
>
> Reviewed-by: Alice Ryhl <aliceryhl@google.com>
> Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>

A couple notes on safety documentation below. Also one pretty subtle
functionality change from last version. With everything fixed:

Reviewed-by: Benno Lossin <lossin@kernel.org>

> ---
>  rust/kernel/sync.rs          |   2 +
>  rust/kernel/sync/set_once.rs | 122 +++++++++++++++++++++++++++++++++++++++++++
>  2 files changed, 124 insertions(+)
>
> diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
> index 81e3a806e57e2..13e6bc7fa87ac 100644
> --- a/rust/kernel/sync.rs
> +++ b/rust/kernel/sync.rs
> @@ -18,6 +18,7 @@
>  mod locked_by;
>  pub mod poll;
>  pub mod rcu;
> +mod set_once;
>  
>  pub use arc::{Arc, ArcBorrow, UniqueArc};
>  pub use completion::Completion;
> @@ -26,6 +27,7 @@
>  pub use lock::mutex::{new_mutex, Mutex, MutexGuard};
>  pub use lock::spinlock::{new_spinlock, SpinLock, SpinLockGuard};
>  pub use locked_by::LockedBy;
> +pub use set_once::SetOnce;
>  
>  /// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
>  #[repr(transparent)]
> diff --git a/rust/kernel/sync/set_once.rs b/rust/kernel/sync/set_once.rs
> new file mode 100644
> index 0000000000000..73706abfe9991
> --- /dev/null
> +++ b/rust/kernel/sync/set_once.rs
> @@ -0,0 +1,122 @@
> +// SPDX-License-Identifier: GPL-2.0
> +
> +//! A container that can be initialized at most once.
> +
> +use super::atomic::{
> +    ordering::{Acquire, Relaxed, Release},
> +    Atomic,
> +};
> +use core::{cell::UnsafeCell, mem::MaybeUninit, ptr::drop_in_place};
> +
> +/// A container that can be populated at most once. Thread safe.
> +///
> +/// Once the a [`SetOnce`] is populated, it remains populated by the same object for the
> +/// lifetime `Self`.
> +///
> +/// # Invariants
> +///
> +/// - `init` may only increase in value.
> +/// - `init` may only assume values in the range `0..=2`.
> +/// - `init == 0` if and only if the container is empty.
> +/// - `init == 1` if and only if being initialized.
> +/// - `init == 2` if and only if the container is populated and valid for shared access.

I think I have a better idea for the last three invariants:

- `init == 0` if and only if `value` is uninitialized.
- `init == 1` if and only if there is exactly one thread with exclusive
  access to `self.value`.
- `init == 2` if and only if `value` is initialized and valid for shared
  access.

> +///
> +/// # Example
> +///
> +/// ```
> +/// # use kernel::sync::SetOnce;
> +/// let value = SetOnce::new();
> +/// assert_eq!(None, value.as_ref());
> +///
> +/// let status = value.populate(42u8);
> +/// assert_eq!(true, status);
> +/// assert_eq!(Some(&42u8), value.as_ref());
> +/// assert_eq!(Some(42u8), value.copy());
> +///
> +/// let status = value.populate(101u8);
> +/// assert_eq!(false, status);
> +/// assert_eq!(Some(&42u8), value.as_ref());
> +/// assert_eq!(Some(42u8), value.copy());
> +/// ```
> +pub struct SetOnce<T> {
> +    init: Atomic<u32>,
> +    value: UnsafeCell<MaybeUninit<T>>,
> +}
> +
> +impl<T> Default for SetOnce<T> {
> +    fn default() -> Self {
> +        Self::new()
> +    }
> +}
> +
> +impl<T> SetOnce<T> {
> +    /// Create a new [`SetOnce`].
> +    ///
> +    /// The returned instance will be empty.
> +    pub const fn new() -> Self {
> +        // INVARIANT: The container is empty and we initialize `init` to `0`.
> +        Self {
> +            value: UnsafeCell::new(MaybeUninit::uninit()),
> +            init: Atomic::new(0),
> +        }
> +    }
> +
> +    /// Get a reference to the contained object.
> +    ///
> +    /// Returns [`None`] if this [`SetOnce`] is empty.
> +    pub fn as_ref(&self) -> Option<&T> {
> +        if self.init.load(Acquire) == 2 {
> +            // SAFETY: By the type invariants of `Self`, `self.init == 2` means that `self.value`
> +            // contains a valid value.

s/contains a valid value/is initialized and valid for shared access/

> +            Some(unsafe { &*self.value.get().cast() })
> +        } else {
> +            None
> +        }
> +    }
> +
> +    /// Populate the [`SetOnce`].
> +    ///
> +    /// Returns `true` if the [`SetOnce`] was successfully populated.
> +    pub fn populate(&self, value: T) -> bool {
> +        // INVARIANT: If the swap succeeds:
> +        //  - We increase `init`.
> +        //  - We write the valid value `1` to `init`.
> +        //  - Only one thread can succeed in this write, so we have exclusive access after the
> +        //    write.
> +        if let Ok(0) = self.init.cmpxchg(0, 1, Relaxed) {
> +            // SAFETY: By the type invariants of `Self`, the fact that we succeeded in writing `1`
> +            // to `self.init` means we obtained exclusive access to the contained object.

s/to the contained object/to `self.value`/

> +            unsafe { core::ptr::write(self.value.get().cast(), value) };
> +            // INVARIANT:
> +            //  - We increase `init`.
> +            //  - We write the valid value `2` to `init`.
> +            //  - We release our exclusive access to the contained object and the object is now
> +            //    valid for shared access.
> +            self.init.store(2, Release);
> +            true
> +        } else {
> +            false
> +        }
> +    }
> +
> +    /// Get a copy of the contained object.
> +    ///
> +    /// Returns [`None`] if the [`SetOnce`] is empty.
> +    pub fn copy(&self) -> Option<T>
> +    where
> +        T: Copy,
> +    {
> +        self.as_ref().copied()
> +    }
> +}
> +
> +impl<T> Drop for SetOnce<T> {
> +    fn drop(&mut self) {
> +        if *self.init.get_mut() == 2 {
> +            // SAFETY: By the type invariants of `Self`, `self.init == 2` means that `self.value`
> +            // contains a valid value. We have exclusive access, as we hold a `mut` reference to
> +            // `self`.
> +            unsafe { drop_in_place(self.value.get()) };

This is sadly doing the wrong thing now since you changed the type of
`value`: `self.value.get()` is of type `MaybeUninit<T>` and dropping
that has (obviously) no effect. So we probably need to do

    let value = unsafe { &mut *self.value.get() };
    unsafe { value.assume_init_drop() };

I almost overlooked this :)

---
Cheers,
Benno

> +        }
> +    }
> +}