Re: [PATCH v14 1/7] rust: sync: add `OnceLock`

[Wren Turkal <wt@penguintechs.org>]

On 7/2/25 6:18 AM, Andreas Hindborg wrote:
> Introduce the `OnceLock` type, a container that can only be written once.
> The container uses an internal atomic to synchronize writes to the internal
> value.
> 
> Signed-off-by: Andreas Hindborg <a.hindborg@kernel.org>
> ---
>   rust/kernel/sync.rs           |   1 +
>   rust/kernel/sync/once_lock.rs | 104 ++++++++++++++++++++++++++++++++++++++++++
>   2 files changed, 105 insertions(+)
> 
> diff --git a/rust/kernel/sync.rs b/rust/kernel/sync.rs
> index c7c0e552bafe..f2ee07315091 100644
> --- a/rust/kernel/sync.rs
> +++ b/rust/kernel/sync.rs
> @@ -15,6 +15,7 @@
>   mod condvar;
>   pub mod lock;
>   mod locked_by;
> +pub mod once_lock;
>   pub mod poll;
>   pub mod rcu;
>   
> diff --git a/rust/kernel/sync/once_lock.rs b/rust/kernel/sync/once_lock.rs
> new file mode 100644
> index 000000000000..cd311bea3919
> --- /dev/null
> +++ b/rust/kernel/sync/once_lock.rs
> @@ -0,0 +1,104 @@
> +//! A container that can be initialized at most once.
> +
> +use super::atomic::ordering::Acquire;
> +use super::atomic::ordering::Release;
> +use super::atomic::Atomic;
> +use kernel::types::Opaque;
> +
> +/// A container that can be populated at most once. Thread safe.
> +///
> +/// Once the a [`OnceLock`] is populated, it remains populated by the same object for the
> +/// lifetime `Self`.
> +///
> +/// # Invariants
> +///
> +/// `init` tracks the state of the container:
> +///
> +/// - If the container is empty, `init` is `0`.
> +/// - If the container is mutably accessed, `init` is `1`.
> +/// - If the container is populated and ready for shared access, `init` is `2`.
> +///
> +/// # Example
> +///
> +/// ```
> +/// # use kernel::sync::once_lock::OnceLock;
> +/// let value = OnceLock::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 OnceLock<T> {
> +    init: Atomic<u32>,
> +    value: Opaque<T>,
> +}

This type looks very much like the Once type in rust's stdlib. I am 
wondering if the api could be changed to match that api. I know that 
this type is trying to provide a version subset of std::sync::OnceLock 
that doesn't allow resetting the type like these apis:

* https://doc.rust-lang.org/std/sync/struct.OnceLock.html#method.get_mut
* https://doc.rust-lang.org/std/sync/struct.OnceLock.html#method.take

However, these methods can only be used on mut. See here for failing 
example: 
https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=a78e51203c5b9555e3c151e162f0acab

I think it might make more sense to match the api of the stdlib API and 
maybe only implement the methods you need.

> +
> +impl<T> Default for OnceLock<T> {
> +    fn default() -> Self {
> +        Self::new()
> +    }
> +}

Any reason not to use #[derive(Default)]?
> +
> +impl<T> OnceLock<T> {
> +    /// Create a new [`OnceLock`].
> +    ///
> +    /// The returned instance will be empty.
> +    pub const fn new() -> Self {

Like new in std OnceLock. Matches. Good.

> +        // INVARIANT: The container is empty and we set `init` to `0`.
> +        Self {
> +            value: Opaque::uninit(),
> +            init: Atomic::new(0),
> +        }
> +    }
> +
> +    /// Get a reference to the contained object.
> +    ///
> +    /// Returns [`None`] if this [`OnceLock`] is empty.
> +    pub fn as_ref(&self) -> Option<&T> {

Looks like the get method in the OnceLock.

> +        if self.init.load(Acquire) == 2 {
> +            // SAFETY: As determined by the load above, the object is ready for shared access.
> +            Some(unsafe { &*self.value.get() })
> +        } else {
> +            None
> +        }
> +    }
> +
> +    /// Populate the [`OnceLock`].
> +    ///
> +    /// Returns `true` if the [`OnceLock`] was successfully populated.
> +    pub fn populate(&self, value: T) -> bool {

Looks like set in OnceLock.

Might also be worth implementing get_or_{try,}init, which get the value 
while initializing.

> +        // INVARIANT: We obtain exclusive access to the contained allocation and write 1 to
> +        // `init`.
> +        if let Ok(0) = self.init.cmpxchg(0, 1, Acquire) {
> +            // SAFETY: We obtained exclusive access to the contained object.
> +            unsafe { core::ptr::write(self.value.get(), value) };
> +            // INVARIANT: We release our exclusive access and transition the object to shared
> +            // access.
> +            self.init.store(2, Release);
> +            true
> +        } else {
> +            false
> +        }
> +    }
> +}
> +
> +impl<T: Copy> OnceLock<T> {
> +    /// Get a copy of the contained object.
> +    ///
> +    /// Returns [`None`] if the [`OnceLock`] is empty.
> +    pub fn copy(&self) -> Option<T> {

No equivalent in OnceLock. Similar to something like this:

x.get().copied().unwrap(); // x is a OnceLock

Example:
https://play.rust-lang.org/?version=nightly&mode=debug&edition=2021&gist=f21068e55f73722544fb5ad341bce1c5

Maybe not specifically needed?

> +        if self.init.load(Acquire) == 2 {
> +            // SAFETY: As determined by the load above, the object is ready for shared access.
> +            Some(unsafe { *self.value.get() })
> +        } else {
> +            None
> +        }
> +    }
> +}
> 

-- 
You're more amazing than you think!