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

[Andreas Hindborg <a.hindborg@kernel.org>]

"Benno Lossin" <lossin@kernel.org> writes:

> On Wed Jul 2, 2025 at 3:18 PM CEST, 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;
>
> As Alice already said, we should reexport the type. And then make the
> module private, no need to have `kernel::sync::OnceLock` and
> `kernel::sync::once_lock::OnceLock`...

Will do.

>
> Also, I agree with the name change to `SetOnce` or something similar.

I'm OK with that, but please see my comments on Alice suggestion.

>
>>  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`.
>
> I think we should swap the order and change the ifs to iffs:
>
>     - `init == 0` iff the container is empty.
>     - `init == 1` iff the container is being accessed mutably.

Right, that is better, but I will expand "iff".

>
>> +/// - If the container is populated and ready for shared access, `init` is `2`.
>
> You also need that `init` is only increased and never decreases.
> Otherwise you could read a `2` and then access the value, but `init`
> changed under your nose to `0`.
>
> Then the INVARIANT comments below also need to be updated.

OK.

>
>> +///
>> +/// # 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>,
>> +}
>> +
>> +impl<T> Default for OnceLock<T> {
>> +    fn default() -> Self {
>> +        Self::new()
>> +    }
>> +}
>> +
>> +impl<T> OnceLock<T> {
>> +    /// Create a new [`OnceLock`].
>> +    ///
>> +    /// The returned instance will be empty.
>> +    pub const fn new() -> Self {
>> +        // 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> {
>> +        if self.init.load(Acquire) == 2 {
>> +            // SAFETY: As determined by the load above, the object is ready for shared access.
>
>     // SAFETY: By the safety requirements of `Self`, `self.init == 2` means that `self.value` contains
>     // a valid value.

By the *type invariants* I guess?

>
>> +            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 {
>> +        // 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> {
>> +        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
>> +        }
>
> The impl can just be:
>
>     self.as_ref().copied()

Nice. I was thinking of dropping this method and just have callers do

 my_once_lock.as_ref().map(|v| v.copied())

What do you think?


Best regards,
Andreas Hindborg