Re: [01/13] rust: sync: support [pin_]init for `SetOnce`

[lyude@redhat.com]

On Thu, 2026-03-26 at 14:52 +0800, Alvin Sun wrote:
> From: Gary Guo <gary@garyguo.net>
> 
> Make `SetOnce` support initialization via `impl Init` or `impl
> PinInit`.
> 
> This adds a possible failing path if an initializer fails. In such
> case,
> the state is dropped back to 0 instead of keep increasing; so the
> monotonicity
> invariant is dropped. The order for initialization is upgraded from
> Relaxed
> to Acquire so it can observe the effect of a previously failed
> initialization.
> 
> Signed-off-by: Gary Guo <gary@garyguo.net>
> Signed-off-by: Alvin Sun <alvin.sun@linux.dev>
> ---
>  rust/kernel/sync/set_once.rs | 120
> +++++++++++++++++++++++++++++++++++--------
>  1 file changed, 99 insertions(+), 21 deletions(-)
> 
> diff --git a/rust/kernel/sync/set_once.rs
> b/rust/kernel/sync/set_once.rs
> index 139cef05e935f..db9c5423fade3 100644
> --- a/rust/kernel/sync/set_once.rs
> +++ b/rust/kernel/sync/set_once.rs
> @@ -2,11 +2,23 @@
>  
>  //! A container that can be initialized at most once.
>  
> +use core::{
> +    cell::UnsafeCell,
> +    mem::MaybeUninit, //
> +};
> +use pin_init::{
> +    Init,
> +    PinInit, //
> +};
> +
>  use super::atomic::{
> -    ordering::{Acquire, Relaxed, Release},
> -    Atomic,
> +    ordering::{
> +        Acquire,
> +        Release, //
> +    },
> +    Atomic, //
>  };
> -use core::{cell::UnsafeCell, mem::MaybeUninit};
> +use crate::prelude::*;
>  
>  /// A container that can be populated at most once. Thread safe.
>  ///
> @@ -15,13 +27,13 @@
>  ///
>  /// # Invariants
>  ///
> -/// - `init` may only increase in value.
>  /// - `init` may only assume values in the range `0..=2`.
>  /// - `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.
> +/// - once `init == 2`, it must remain so.
>  ///
>  /// # Example
>  ///
> @@ -51,6 +63,35 @@ fn default() -> Self {
>      }
>  }
>  
> +/// Error that can occur during initialization of `SetOnce`.
> +#[derive(Debug)]
> +pub enum InitError<E> {
> +    /// The `Once` has already been initialized.
> +    AlreadyInit,
> +    /// The `Once` is being raced to initialize by another thread.
> +    RacedInit,
> +    /// Error occurs during initialization.
> +    DuringInit(E),
> +}
> +
> +impl<E> From<E> for InitError<E> {
> +    #[inline]
> +    fn from(err: E) -> Self {
> +        InitError::DuringInit(err)
> +    }
> +}
> +
> +impl<E: Into<Error>> From<InitError<E>> for Error {
> +    #[inline]
> +    fn from(this: InitError<E>) -> Self {
> +        match this {
> +            InitError::AlreadyInit => EEXIST,
> +            InitError::RacedInit => EBUSY,
> +            InitError::DuringInit(e) => e.into(),
> +        }
> +    }
> +}
> +
>  impl<T> SetOnce<T> {
>      /// Create a new [`SetOnce`].
>      ///
> @@ -76,31 +117,68 @@ pub fn as_ref(&self) -> Option<&T> {
>          }
>      }
>  
> -    /// Populate the [`SetOnce`].
> +    /// Populate the [`SetOnce`] with an initializer.
>      ///
> -    /// Returns `true` if the [`SetOnce`] was successfully
> populated.
> -    pub fn populate(&self, value: T) -> bool {
> +    /// Returns the initialized reference if the [`SetOnce`] was
> successfully populated.
> +    pub fn init<E>(&self, init: impl Init<T, E>) -> Result<&T,
> InitError<E>> {
>          // INVARIANT: If the swap succeeds:
> -        //  - We increase `init`.
>          //  - We write the valid value `1` to `init`.
> +        //  - The previous value is not `2`, so it is valid to move
> to `1`.
>          //  - 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
> `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 `self.value`
> and it is now valid for shared
> -            //    access.
> -            self.init.store(2, Release);
> -            true
> -        } else {
> -            false
> +        match self.init.cmpxchg(0, 1, Acquire) {
> +            Ok(_) => {
> +                // SAFETY:
> +                // - By the type invariants of `Self`, the fact that
> we succeeded in writing `1`
> +                //   to `self.init` means we obtained exclusive
> access to `self.value`.
> +                // - When `Err` is returned, we did not set
> `self.init` to `2` so the `Drop` is not
> +                //   armed.
> +                match unsafe { init.__init(self.value.get().cast())
> } {
> +                    Ok(()) => {
> +                        // INVARIANT:
> +                        //  - The previous value is `1`, so it is
> valid to move to `2`.
> +                        //  - We write the valid value `2` to
> `init`.
> +                        //  - We release our exclusive access to
> `self.value` and it is now valid
> +                        //    for shared access.
> +                        self.init.store(2, Release);
> +                        // SAFETY: we have just initialized the
> value.
> +                        Ok(unsafe { &*self.value.get().cast() })
> +                    }
> +                    Err(err) => {
> +                        // INVARIANT:
> +                        //  - The previous value is `1`, so it is
> valid to move to `0`.
> +                        //  - We write the valid value `0` to
> `init`.
> +                        //  - We release our exclusive access to
> `self.value` and it is now valid
> +                        //    for shared access.
> +                        self.init.store(0, Release);
> +                        Err(err.into())
> +                    }
> +                }
> +            }
> +            Err(1) => Err(InitError::RacedInit),
> +            Err(_) => Err(InitError::AlreadyInit),
>          }
>      }

So: I've been struggling quite a bit trying to understand what the
purpose of this is. The main problem is I don't see how callers are
actually supposed to make use of this. Or at the very least, maybe this
type should have a more appropriate name because the places it can be
used are a lot more limited then they appear.

Assume for a moment that we have two threads, each of them notices that
a SetOnce is empty and attempts to initialize it. This patch would of
course make it so that one thread successfully initializes it, and the
other gets InitError::RacedInit or InitError::AlreadyInit depending on
luck.

Consider the following function I tried writing, where `self.sg_res` is
`SetOnce<Devres<SGTableMap<T>>>`:

    fn get_sg_table<'a>(
        &'a self,
        dev: &'a device::Device<Bound>,
    ) -> Result<&'a Devres<SGTableMap<T>>> {
        if dev.as_raw() != self.dev().as_ref().as_raw() {
            return Err(EINVAL);
        }

        if let Some(sgt_res) = self.sgt_res.as_ref() {
            return Ok(sgt_res);
        }

        // SAFETY:
        // * `slot` only has a single field, which is initialized in this callback
        // * If any of the steps in this closure fail, we return Err(err)
        // * It is perfectly fine to move a Devres around after initialization
        // * T has no pinning variants to uphold.
        let init = unsafe {
            init_from_closure(|sgt_res| {
                *sgt_res = Devres::new(dev, SGTableMap::new(self))?;
                Ok(())
            })
        };

        match self.sgt_res.init(init) {
            Ok(res) => Ok(res),
            Err(InitError::DuringInit(e)) => Err(e),
            Err(InitError::AlreadyInit) => {
                // SAFETY: We know that `sgt_res` is populated now, so its safe to unwrap
                Ok(unsafe { self.sgt_res.as_ref().unwrap_unchecked() })
            }
            // We need to busy-wait until the other thread finishes init
            Err(InitError::RacedInit) => loop {
		// XXX: The below line will not work for reasons explained
		// later in this email
                if let Some(sgt_res) = self.sgt_res.as_ref() {
                    break Ok(sgt_res);
                }
            }, // (retry)
        }

There's a few issues. First, the most obvious one: there doesn't
actually appear to be any way of actually handling raced initialization
without busy-looping.

The other problem, which I mentioned in an XXX: comment in the above
snipping, is that busy-waiting isn't even possible unless you busy-wait
on .init() instead of .as_ref(). .init() just returns `Some(…)` or
`None`, which means that if `self.sgt_res.populate(…)` actually fails
during initialization we'll just busy-wait forever because we keep
getting `None`.

So, the only correct way of busy-waiting is repeatedly calling .init()
with a value until we get `InitError::AlreadyInit`, and then finally
calling `as_ref()` to retrieve the value that got stored.

I have to assume that this is probably not the way that this API is
intended to be used and as such, maybe this needs to be renamed or we
need to rethink this a bit?

For what it's worth, I'm also looking at potentially making a thread-
safe single-init container that avoids this issue entirely by just
using a Mutex to protect initialization. Which would look something
like this (I hate this naming, part of the reason I wrote this email
:):

#[pin_data]
struct InitOnce<T: Send + Sync, E> {
    #[pin]
    inner: Mutex<Option<Result<T, E>>>
}

where the basic idea is that the `Mutex` is only actually used for
protecting the initial population of the contents of `Option<Result<T,
E>>`, and the actual contents of the mutex can be shared outside of the
critical section. Soundness would be enforced on the latter condition
by making it so that, similar to SetOnce, you can only really write to
it a single time. 

>  
> +    /// Populate the [`SetOnce`] with a pinned initializer.
> +    ///
> +    /// Returns the initialized reference if the [`SetOnce`] was
> successfully populated.
> +    pub fn pin_init<E>(self: Pin<&Self>, init: impl PinInit<T, E>) -
> > Result<&T, InitError<E>> {
> +        // SAFETY:
> +        // - `__pinned_init` satisfy all requirements of
> `init_from_closure`
> +        // - calling `__pinned_init` require additional that the
> slot is pinned, which is satisfied
> +        //   given `self: Pin<&Self>`.
> +        self.get_ref()
> +            .init(unsafe { pin_init::init_from_closure(|slot|
> init.__pinned_init(slot)) })
> +    }
> +
> +    /// Populate the [`SetOnce`].
> +    ///
> +    /// Returns `true` if the [`SetOnce`] was successfully
> populated.
> +    pub fn populate(&self, value: T) -> bool {
> +        self.init(value).is_ok()
> +    }
> +
>      /// Get a copy of the contained object.
>      ///
>      /// Returns [`None`] if the [`SetOnce`] is empty.