sync 2

kernel/src/sync/mod.rs

@@ -86,76 +86,197 @@ pub use spin_mutex::SpinMutex;
 
 mod once {
     use core::{
-        cell::UnsafeCell,
+        cell::{Cell, UnsafeCell},
         mem::{ManuallyDrop, MaybeUninit},
         sync::atomic::{AtomicU8, Ordering},
     };
 
-    pub struct Once<T> {
+    const UNINITIALIZED: u8 = 0;
+    const INITIALIZING: u8 = 1;
+    const INITIALIZED: u8 = 2;
+    const PANICKED: u8 = 3;
+
+    pub struct OnceState {
+        poisoned: bool,
+        state_to_set: Cell<u8>,
+    }
+
+    impl OnceState {
+        pub fn is_poisoned(&self) -> bool {
+            self.poisoned
+        }
+        pub fn poison(&self) {
+            self.state_to_set.set(PANICKED);
+        }
+    }
+
+    pub struct Once {
         state: AtomicU8,
+    }
+
+    const impl Default for Once {
+        fn default() -> Self {
+            Self {
+                state: AtomicU8::new(UNINITIALIZED),
+            }
+        }
+    }
+
+    impl Once {
+        pub const fn new() -> Self {
+            Self {
+                state: AtomicU8::new(UNINITIALIZED),
+            }
+        }
+
+        pub fn state(&self) -> u8 {
+            self.state.load(Ordering::Acquire)
+        }
+
+        pub fn set_state(&self, state: u8) {
+            self.state.store(state, Ordering::Release);
+        }
+
+        pub fn is_completed(&self) -> bool {
+            self.state.load(Ordering::Acquire) == INITIALIZED
+        }
+
+        pub fn call_once<F>(&self, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            if self.is_completed() {
+                return;
+            }
+
+            self.call_once_slow(false, f);
+        }
+
+        pub fn call_once_force<F>(&self, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            if self.is_completed() {
+                return;
+            }
+
+            self.call_once_slow(true, f);
+        }
+
+        #[cold]
+        pub fn call_once_slow<F>(&self, ignore_poison: bool, f: F)
+        where
+            F: FnOnce(&OnceState),
+        {
+            let mut state = self.state.load(Ordering::Acquire);
+            loop {
+                match state {
+                    INITIALIZED => return,
+                    PANICKED if !ignore_poison => {
+                        panic!("Once instance has previously been poisoned")
+                    }
+                    PANICKED | UNINITIALIZED => {
+                        match self.state.compare_exchange(
+                            state,
+                            INITIALIZING,
+                            Ordering::Acquire,
+                            // if we get `Err(INITIALIZED)`, we want to have
+                            // acquired what the lock is protecting.
+                            Ordering::Acquire,
+                        ) {
+                            Err(new) => {
+                                state = new;
+                                continue;
+                            }
+                            Ok(_) => {
+                                // even though we don't have unwinding, for
+                                // completeness sake we'll poison the lock if
+                                // the closure panics.
+                                struct Guard<'a>(&'a AtomicU8);
+
+                                impl Drop for Guard<'_> {
+                                    fn drop(&mut self) {
+                                        self.0.store(PANICKED, Ordering::Release);
+                                    }
+                                }
+
+                                let guard = Guard(&self.state);
+
+                                let state = OnceState {
+                                    poisoned: state == PANICKED,
+                                    state_to_set: Cell::new(INITIALIZED),
+                                };
+
+                                f(&state);
+                                core::mem::forget(guard);
+
+                                self.state
+                                    .store(state.state_to_set.take(), Ordering::Release);
+                                return;
+                            }
+                        }
+                    }
+                    _ => {
+                        assert_eq!(state, INITIALIZING);
+                        loop {
+                            state = self.state.load(Ordering::Acquire);
+                            if state != INITIALIZING {
+                                break;
+                            }
+                            core::hint::spin_loop();
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    pub struct OnceLock<T> {
+        once: Once,
         data: UnsafeCell<MaybeUninit<T>>,
     }
 
-    unsafe impl<T: Send> Send for Once<T> {}
-    unsafe impl<T: Send + Sync> Sync for Once<T> {}
+    unsafe impl<T: Send> Send for OnceLock<T> {}
+    unsafe impl<T: Send + Sync> Sync for OnceLock<T> {}
 
-    const impl<T> Default for Once<T> {
+    const impl<T> Default for OnceLock<T> {
         fn default() -> Self {
             Self {
-                state: AtomicU8::new(0),
+                once: Once::new(),
                 data: UnsafeCell::new(MaybeUninit::uninit()),
             }
         }
     }
 
-    #[repr(u8)]
-    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-    enum OnceState {
-        Uninitialized = 0,
-        Initializing = 1,
-        Initialized = 2,
-    }
-    impl OnceState {
-        fn from_u8(value: u8) -> Option<Self> {
-            match value {
-                0 => Some(Self::Uninitialized),
-                1 => Some(Self::Initializing),
-                2 => Some(Self::Initialized),
-                _ => None,
-            }
-        }
-        unsafe fn from_u8_unchecked(value: u8) -> Self {
-            unsafe { core::mem::transmute(value) }
-        }
-        const fn into_u8(self) -> u8 {
-            self as u8
-        }
-    }
-
-    impl<T> Once<T> {
+    impl<T> OnceLock<T> {
         pub const fn new() -> Self {
             Self {
-                state: AtomicU8::new(OnceState::Uninitialized.into_u8()),
+                once: Once::new(),
                 data: UnsafeCell::new(MaybeUninit::uninit()),
             }
         }
 
         pub const fn from(t: T) -> Self {
             Self {
-                state: AtomicU8::new(OnceState::Initialized.into_u8()),
+                once: Once::new(),
                 data: UnsafeCell::new(MaybeUninit::new(t)),
             }
         }
 
+        /// # Safety
+        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
         pub unsafe fn get_unchecked(&self) -> &T {
             unsafe { self.data.get().as_ref_unchecked().assume_init_ref() }
         }
+
+        /// # Safety
+        /// The caller must ensure that the `OnceLock` has been initialized before calling this method.
         pub unsafe fn get_mut_unchecked(&mut self) -> &mut T {
             unsafe { self.data.get().as_mut_unchecked().assume_init_mut() }
         }
 
         pub fn get(&self) -> Option<&T> {
-            if self.state(Ordering::Acquire) == OnceState::Initialized {
+            if self.is_completed() {
                 Some(unsafe { self.get_unchecked() })
             } else {
                 None
@@ -163,7 +284,7 @@ mod once {
         }
 
         pub fn get_mut(&mut self) -> Option<&mut T> {
-            if self.state(Ordering::Acquire) == OnceState::Initialized {
+            if self.is_completed_mut() {
                 Some(unsafe { self.get_mut_unchecked() })
             } else {
                 None
@@ -171,86 +292,52 @@ mod once {
         }
 
         pub fn is_completed(&self) -> bool {
-            self.state(Ordering::Acquire) == OnceState::Initialized
+            self.once.is_completed()
         }
 
-        fn get_spinning(&self) -> Option<&T> {
+        pub fn is_completed_mut(&mut self) -> bool {
+            // avoid atomic load when we have mutable access to self
+            *self.once.state.get_mut() == INITIALIZED
+        }
+
+        pub fn get_spinning(&self) -> Option<&T> {
             loop {
-                match self.state(Ordering::Acquire) {
-                    OnceState::Uninitialized => return None,
-                    OnceState::Initializing => core::hint::spin_loop(),
-                    OnceState::Initialized => return Some(unsafe { self.get_unchecked() }),
+                match self.once.state() {
+                    INITIALIZED => return Some(unsafe { self.get_unchecked() }),
+                    INITIALIZING => core::hint::spin_loop(),
+                    UNINITIALIZED | PANICKED => return None,
+                    _ => unreachable!(),
                 }
             }
         }
 
-        fn should_try_init(&self) -> bool {
-            match unsafe {
-                self.state
-                    .compare_exchange(
-                        OnceState::Uninitialized.into_u8(),
-                        OnceState::Initializing.into_u8(),
-                        Ordering::Acquire,
-                        Ordering::Relaxed,
-                    )
-                    .map(|v| OnceState::from_u8_unchecked(v))
-                    .map_err(|v| OnceState::from_u8_unchecked(v))
-            } {
-                Ok(_) => true,
-                Err(OnceState::Initializing) => self.get_spinning().is_none(),
-                Err(OnceState::Initialized) => false,
-                Err(_) => unreachable!(),
-            }
+        pub fn initialize<F, E>(&self, f: F) -> Result<(), E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            let mut res: Result<(), E> = Ok(());
+            let cell = unsafe { &mut *self.data.get() };
+
+            self.once.call_once_force(|state| match f() {
+                Ok(value) => {
+                    unsafe { cell.as_mut_ptr().write(value) };
+                }
+                Err(e) => {
+                    state.poison();
+                    res = Err(e);
+                }
+            });
+
+            res
         }
 
         pub fn get_or_init<F>(&self, f: F) -> &T
         where
             F: FnOnce() -> T,
         {
-            if self.get().is_none() {
-                self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f()));
-            }
-
-            unsafe { self.get_unchecked() }
-        }
-
-        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
-        where
-            F: FnOnce() -> Result<T, E>,
-        {
-            if !self.should_try_init() {
-                return Ok(unsafe { self.get_unchecked() });
-            }
-
-            struct PanicGuard<'a>(&'a AtomicU8);
-
-            impl Drop for PanicGuard<'_> {
-                fn drop(&mut self) {
-                    self.0
-                        .store(OnceState::Uninitialized.into_u8(), Ordering::Release);
-                }
-            }
-
-            let guard = PanicGuard(&self.state);
-
-            match f() {
-                Ok(value) => {
-                    unsafe { self.data.get().as_mut_unchecked().write(value) };
-
-                    core::mem::forget(guard);
-
-                    self.state
-                        .store(OnceState::Initialized.into_u8(), Ordering::Release);
-
-                    Ok(unsafe { self.get_unchecked() })
-                }
-                Err(e) => {
-                    core::mem::forget(guard);
-                    self.state
-                        .store(OnceState::Uninitialized.into_u8(), Ordering::Release);
-
-                    Err(e)
-                }
+            match self.get_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
+                Ok(value) => value,
+                Err(_) => unreachable!(),
             }
         }
 
@@ -258,11 +345,32 @@ mod once {
         where
             F: FnOnce() -> T,
         {
-            if self.get_mut().is_none() {
-                self.init_once(f);
+            match self.get_mut_or_try_init(|| Ok::<T, core::convert::Infallible>(f())) {
+                Ok(value) => value,
+                Err(_) => unreachable!(),
+            }
+        }
+
+        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            if !self.is_completed() {
+                self.initialize(f)?;
             }
 
-            unsafe { self.get_mut_unchecked() }
+            unsafe { Ok(self.get_unchecked()) }
+        }
+
+        pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
+        where
+            F: FnOnce() -> Result<T, E>,
+        {
+            if self.get_mut().is_none() {
+                self.initialize(f)?;
+            }
+
+            Ok(unsafe { self.get_mut_unchecked() })
         }
 
         pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
@@ -273,26 +381,9 @@ mod once {
                 None => Ok(res),
             }
         }
-
-        fn init_once<F>(&self, f: F) -> &T
-        where
-            F: FnOnce() -> T,
-        {
-            if self.should_try_init() {
-                let value = f();
-                unsafe { self.data.get().as_mut_unchecked().write(value) };
-                self.state
-                    .store(OnceState::Initialized.into_u8(), Ordering::Release);
-            }
-            unsafe { self.get_unchecked() }
-        }
-
-        fn state(&self, ordering: Ordering) -> OnceState {
-            unsafe { OnceState::from_u8_unchecked(self.state.load(ordering)) }
-        }
     }
 
-    impl<T> Drop for Once<T> {
+    impl<T> Drop for OnceLock<T> {
         fn drop(&mut self) {
             if self.is_completed() {
                 unsafe { self.data.get().as_mut_unchecked().assume_init_drop() }
@@ -301,11 +392,12 @@ mod once {
     }
 
     pub struct LazyLock<T, F = fn() -> T> {
-        once: Once<T>,
-        init: UnsafeCell<ManuallyDrop<F>>,
+        once: Once,
+        f: UnsafeCell<ManuallyDrop<F>>,
+        t: UnsafeCell<MaybeUninit<T>>,
     }
 
-    unsafe impl<T, F: Send> Sync for LazyLock<T, F> where Once<T>: Sync {}
+    unsafe impl<T: Sync + Send, F: Send> Sync for LazyLock<T, F> {}
 
     impl<T, F> LazyLock<T, F>
     where
@@ -314,30 +406,41 @@ mod once {
         pub const fn new(f: F) -> Self {
             Self {
                 once: Once::new(),
-                init: UnsafeCell::new(ManuallyDrop::new(f)),
+                f: UnsafeCell::new(ManuallyDrop::new(f)),
+                t: UnsafeCell::new(MaybeUninit::uninit()),
             }
         }
 
         pub fn get(&self) -> Option<&T> {
-            self.once.get()
+            if self.once.is_completed() {
+                Some(unsafe { (&*self.t.get()).assume_init_ref() })
+            } else {
+                None
+            }
         }
 
         fn force(this: &Self) -> &T {
-            // SAFETY: `once` cannot recycle from `Initializing` to `Uninitialized`, so the closure will only be called once.
-            this.once.get_or_init(|| unsafe {
-                let init = &mut *this.init.get();
-                let f = ManuallyDrop::take(init);
-                f()
-            })
+            this.once.call_once(|_| {
+                // SAFETY: because `call_once` will panic if poisoned, this
+                // closure will only be called once.
+                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
+                let val = f();
+                unsafe { (&mut *this.t.get()).write(val) };
+            });
+
+            unsafe { (&*this.t.get()).assume_init_ref() }
         }
 
         fn force_mut(this: &mut Self) -> &mut T {
-            // SAFETY: `once` cannot recycle from `Initializing` to `Uninitialized`, so the closure will only be called once.
-            this.once.get_mut_or_init(|| unsafe {
-                let init = &mut *this.init.get();
-                let f = ManuallyDrop::take(init);
-                f()
-            })
+            this.once.call_once(|_| {
+                // SAFETY: because `call_once` will panic if poisoned, this
+                // closure will only be called once.
+                let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) };
+                let val = f();
+                unsafe { (&mut *this.t.get()).write(val) };
+            });
+
+            unsafe { (&mut *this.t.get()).assume_init_mut() }
         }
     }
 
@@ -357,13 +460,17 @@ mod once {
 
     impl<T, F> Drop for LazyLock<T, F> {
         fn drop(&mut self) {
-            if !self.once.is_completed() {
-                unsafe {
-                    ManuallyDrop::drop(self.init.get_mut());
-                }
+            match self.once.state() {
+                UNINITIALIZED => unsafe {
+                    ManuallyDrop::drop(self.f.get_mut());
+                },
+                INITIALIZED => unsafe {
+                    (&mut *self.t.get()).assume_init_drop();
+                },
+                _ => {}
             }
         }
     }
 }
 
-pub use once::{LazyLock, Once};
+pub use once::{LazyLock, Once, OnceLock};