rust: sync: add `SetOnce` (821fe7bf) · Commits · git / linux-net

rust/kernel/sync.rs

+2 −0

Original line number	Diff line number	Diff line
		@@ -20,6 +20,7 @@
		pub mod poll;
		pub mod rcu;
		mod refcount;
		mod set_once;

		pub use arc::{Arc, ArcBorrow, UniqueArc};
		pub use completion::Completion;
		@@ -29,6 +30,7 @@
		pub use lock::spinlock::{new_spinlock, SpinLock, SpinLockGuard};
		pub use locked_by::LockedBy;
		pub use refcount::Refcount;
		pub use set_once::SetOnce;

		/// Represents a lockdep class. It's a wrapper around C's `lock_class_key`.
		#[repr(transparent)]

rust/kernel/sync/set_once.rs

0 → 100644

+125 −0

Original line number	Diff line number	Diff line
		// 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};

		/// 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 `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`
		// 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 `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
		}
		}

		/// 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 {
		let value = self.value.get_mut();
		// 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 { value.assume_init_drop() };
		}
		}
		}