添加与rust std接口相同的once库 (DragonOS-Community#353)

kernel/Cargo.toml

@@ -25,6 +25,7 @@ num-derive = "0.3"
 # 一个no_std的hashmap、hashset
 hashbrown = "0.13.2"
 elf = { version = "0.7.2", default-features = false }
+atomic_enum = "0.2.0"
 
 # 构建时依赖项
 [build-dependencies]

kernel/src/libs/mod.rs

@@ -11,6 +11,8 @@ pub mod lib_ui;
 pub mod list;
 pub mod lockref;
 pub mod mutex;
+pub mod notifier;
+pub mod once;
 pub mod printk;
 pub mod rbtree;
 #[macro_use]
@@ -21,5 +23,4 @@ pub mod spinlock;
 pub mod vec_cursor;
 #[macro_use]
 pub mod volatile;
-pub mod notifier;
 pub mod wait_queue;

kernel/src/libs/once.rs

@@ -0,0 +1,179 @@
+use core::{
+    fmt::{self, Debug, Formatter},
+    sync::atomic::Ordering,
+};
+
+use atomic_enum::atomic_enum;
+
+pub struct Once {
+    inner: AtomicOnceState,
+}
+
+#[atomic_enum]
+#[derive(PartialEq, Eq)]
+pub enum OnceState {
+    Incomplete,
+    Posioned,
+    Complete,
+}
+
+#[allow(dead_code)]
+impl Once {
+    pub const fn new() -> Self {
+        Self {
+            inner: AtomicOnceState::new(OnceState::Incomplete),
+        }
+    }
+
+    #[track_caller]
+    pub fn call_once<F: FnOnce()>(&self, f: F) {
+        if self.is_completed() {
+            return;
+        }
+
+        // set initialized
+        let r = self.inner.compare_exchange(
+            OnceState::Incomplete,
+            OnceState::Posioned,
+            Ordering::SeqCst,
+            Ordering::SeqCst,
+        );
+        if r.is_err() {
+            return;
+        }
+        // call function
+        f();
+        // set completed
+        self.inner.store(OnceState::Complete, Ordering::SeqCst);
+    }
+
+    /// Performs the same function as [`call_once()`] except ignores poisoning.
+    ///
+    /// Unlike [`call_once()`], if this [`Once`] has been poisoned (i.e., a previous
+    /// call to [`call_once()`] or [`call_once_force()`] caused a panic), calling
+    /// [`call_once_force()`] will still invoke the closure `f` and will _not_
+    /// result in an immediate panic. If `f` panics, the [`Once`] will remain
+    /// in a poison state. If `f` does _not_ panic, the [`Once`] will no
+    /// longer be in a poison state and all future calls to [`call_once()`] or
+    /// [`call_once_force()`] will be no-ops.
+    ///
+    /// The closure `f` is yielded a [`OnceState`] structure which can be used
+    /// to query the poison status of the [`Once`].
+    ///
+    /// [`call_once()`]: Once::call_once
+    /// [`call_once_force()`]: Once::call_once_force
+    ///
+    /// # Examples
+    ///
+    /// ```
+    /// use std::sync::Once;
+    /// use std::thread;
+    ///
+    /// static INIT: Once = Once::new();
+    ///
+    /// // poison the once
+    /// let handle = thread::spawn(|| {
+    ///     INIT.call_once(|| panic!());
+    /// });
+    /// assert!(handle.join().is_err());
+    ///
+    /// // poisoning propagates
+    /// let handle = thread::spawn(|| {
+    ///     INIT.call_once(|| {});
+    /// });
+    /// assert!(handle.join().is_err());
+    ///
+    /// // call_once_force will still run and reset the poisoned state
+    /// INIT.call_once_force(|state| {
+    ///     assert!(state.is_poisoned());
+    /// });
+    ///
+    /// // once any success happens, we stop propagating the poison
+    /// INIT.call_once(|| {});
+    /// ```
+    pub fn call_once_force<F>(&self, f: F)
+    where
+        F: FnOnce(&OnceState),
+    {
+        // fast path check
+        if self.is_completed() {
+            return;
+        }
+
+        // set poisoned
+        self.inner
+            .compare_exchange(
+                OnceState::Incomplete,
+                OnceState::Posioned,
+                Ordering::SeqCst,
+                Ordering::SeqCst,
+            )
+            .ok();
+
+        // call function
+        f(&self.inner.load(Ordering::SeqCst));
+
+        // set initialized
+        self.inner.store(OnceState::Complete, Ordering::SeqCst);
+    }
+
+    /// Fast path check
+    #[inline]
+    pub fn is_completed(&self) -> bool {
+        self.inner.load(Ordering::SeqCst) == OnceState::Complete
+    }
+
+    /// Returns the current state of the `Once` instance.
+    #[inline]
+    pub fn state(&self) -> OnceState {
+        self.inner.load(Ordering::SeqCst)
+    }
+}
+
+impl Debug for Once {
+    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Once").finish_non_exhaustive()
+    }
+}
+
+#[allow(dead_code)]
+impl OnceState {
+    /// Returns `true` if the associated [`Once`] was poisoned prior to the
+    /// invocation of the closure passed to [`Once::call_once_force()`].
+    ///
+    /// # Examples
+    ///
+    /// A poisoned [`Once`]:
+    ///
+    /// ```
+    /// use std::sync::Once;
+    /// use std::thread;
+    ///
+    /// static INIT: Once = Once::new();
+    ///
+    /// // poison the once
+    /// let handle = thread::spawn(|| {
+    ///     INIT.call_once(|| panic!());
+    /// });
+    /// assert!(handle.join().is_err());
+    ///
+    /// INIT.call_once_force(|state| {
+    ///     assert!(state.is_poisoned());
+    /// });
+    /// ```
+    ///
+    /// An unpoisoned [`Once`]:
+    ///
+    /// ```
+    /// use std::sync::Once;
+    ///
+    /// static INIT: Once = Once::new();
+    ///
+    /// INIT.call_once_force(|state| {
+    ///     assert!(!state.is_poisoned());
+    /// });
+    #[inline]
+    pub fn is_poisoned(&self) -> bool {
+        *self == OnceState::Posioned
+    }
+}