Global: use the wheel, dont reinvent it

changed to wrap Mutex<OnceLock>

Author: bend-n

godot-ffi/src/global.rs

@@ -5,7 +5,7 @@
  * file, You can obtain one at https://mozilla.org/MPL/2.0/.
  */
 
-use std::sync::{Mutex, MutexGuard, TryLockError};
+use std::sync::{Mutex, OnceLock, PoisonError, TryLockError};
 
 /// Ergonomic global variables.
 ///
@@ -23,9 +23,10 @@ use std::sync::{Mutex, MutexGuard, TryLockError};
 ///
 /// There are two `const` methods for construction: [`new()`](Self::new) and [`default()`](Self::default).
 /// For access, you should primarily use [`lock()`](Self::lock). There is also [`try_lock()`](Self::try_lock) for special cases.
-pub struct Global<T> {
+pub struct Global<T, F = fn() -> T> {
     // When needed, this could be changed to use RwLock and separate read/write guards.
-    value: Mutex<InitState<T>>,
+    value: Mutex<OnceLock<T>>,
+    init_fn: F,
 }
 
 impl<T> Global<T> {
@@ -35,7 +36,8 @@ impl<T> Global<T> {
     pub const fn new(init_fn: fn() -> T) -> Self {
         // Note: could be generalized to F: FnOnce() -> T + Send. See also once_cell::Lazy<T, F>.
         Self {
-            value: Mutex::new(InitState::Pending(init_fn)),
+            value: Mutex::new(OnceLock::new()),
+            init_fn,
         }
     }
 
@@ -57,72 +59,30 @@ impl<T> Global<T> {
     /// If the initialization function panics. Once that happens, the global is considered poisoned and all future calls to `lock()` will panic.
     /// This can currently not be recovered from.
     pub fn lock(&self) -> GlobalGuard<'_, T> {
-        let mutex_guard = self
-            .value
-            .lock()
-            .expect("Global<T> poisoned; a thread has panicked while holding a lock to it");
-
-        let guard = Self::ensure_init(mutex_guard, true)
-            .unwrap_or_else(|| panic!("previous Global<T> initialization failed due to panic"));
-
-        guard
+        let guard = self.value.lock().unwrap_or_else(PoisonError::into_inner);
+        guard.get_or_init(self.init_fn);
+        // SAFETY: above line guarantees init
+        unsafe { GlobalGuard::new_unchecked(guard) }
     }
 
-    /// Non-panicking access with error introspection.
+    /// Non-blocking access with error introspection.
     pub fn try_lock(&self) -> Result<GlobalGuard<'_, T>, GlobalLockError<'_, T>> {
-        let guard = match self.value.try_lock() {
-            Ok(mutex_guard) => Self::ensure_init(mutex_guard, false),
-            Err(TryLockError::WouldBlock) => {
-                return Err(GlobalLockError::WouldBlock);
-            }
-            Err(TryLockError::Poisoned(poisoned)) => {
-                return Err(GlobalLockError::Poisoned {
-                    // We can likely use `new_unchecked` here, but verifying that it's safe would need somewhat tricky reasoning.
-                    // Since this error condition isn't very common, it is likely not very important to optimize access to the value here.
-                    // Especially since most users will likely not want to access it anyway.
-                    circumvent: GlobalGuard::new(poisoned.into_inner())
-                        .expect("Poisoned global guard should always be initialized"),
-                });
-            }
-        };
-
-        guard.ok_or(GlobalLockError::InitFailed)
-    }
-
-    fn ensure_init(
-        mut mutex_guard: MutexGuard<'_, InitState<T>>,
-        may_panic: bool,
-    ) -> Option<GlobalGuard<'_, T>> {
-        let init_fn = match &mut *mutex_guard {
-            InitState::Initialized(_) => {
-                // SAFETY: `mutex_guard` is `Initialized`.
-                return Some(unsafe { GlobalGuard::new_unchecked(mutex_guard) });
-            }
-            InitState::Failed => {
-                return None;
-            }
-            InitState::Pending(init_fn) => init_fn,
-        };
-
-        // Unwinding should be safe here, as there is no unsafe code relying on it.
-        let init_fn = std::panic::AssertUnwindSafe(init_fn);
-        match std::panic::catch_unwind(init_fn) {
-            Ok(value) => *mutex_guard = InitState::Initialized(value),
-            Err(e) => {
-                eprintln!("panic during Global<T> initialization");
-                *mutex_guard = InitState::Failed;
-
-                if may_panic {
-                    std::panic::resume_unwind(e);
-                } else {
-                    // Note: this currently swallows panic.
-                    return None;
+        let guard = self.value.try_lock().map_err(|x| match x {
+            TryLockError::Poisoned(x) => {
+                let g = x.into_inner();
+                match std::panic::catch_unwind(|| g.get_or_init(self.init_fn)) {
+                    Ok(_) => GlobalLockError::Poisoned {
+                        // SAFETY: see get_or_init.
+                        circumvent: unsafe { GlobalGuard::new_unchecked(g) },
+                    },
+                    Err(_) => GlobalLockError::InitFailed,
                 }
             }
-        };
-
-        // SAFETY: `mutex_guard` was either set to `Initialized` above, or we returned from the function.
-        Some(unsafe { GlobalGuard::new_unchecked(mutex_guard) })
+            TryLockError::WouldBlock => GlobalLockError::WouldBlock,
+        })?;
+        std::panic::catch_unwind(|| guard.get_or_init(self.init_fn))
+            .map_err(|_| GlobalLockError::InitFailed)?;
+        Ok(unsafe { GlobalGuard::new_unchecked(guard) })
     }
 }
 
@@ -131,57 +91,50 @@ impl<T> Global<T> {
 
 // Encapsulate private fields.
 mod global_guard {
-    use std::ops::{Deref, DerefMut};
-    use std::sync::MutexGuard;
-
-    use super::InitState;
+    use std::{
+        ops::{Deref, DerefMut},
+        sync::{MutexGuard, OnceLock},
+    };
 
     /// Guard that temporarily gives access to a `Global<T>`'s inner value.
     pub struct GlobalGuard<'a, T> {
-        // Safety invariant: Is `Initialized`.
-        mutex_guard: MutexGuard<'a, InitState<T>>,
+        mutex_guard: MutexGuard<'a, OnceLock<T>>,
     }
 
     impl<'a, T> GlobalGuard<'a, T> {
-        pub(super) fn new(mutex_guard: MutexGuard<'a, InitState<T>>) -> Option<Self> {
-            match &*mutex_guard {
-                InitState::Initialized(_) => Some(Self { mutex_guard }),
+        pub(super) fn new(mutex_guard: MutexGuard<'a, OnceLock<T>>) -> Option<Self> {
+            match mutex_guard.get() {
+                Some(_) => Some(Self { mutex_guard }),
                 _ => None,
             }
         }
 
         /// # Safety
         ///
-        /// The value must be `Initialized`.
-        pub(super) unsafe fn new_unchecked(mutex_guard: MutexGuard<'a, InitState<T>>) -> Self {
-            debug_assert!(matches!(*mutex_guard, InitState::Initialized(_)));
-
+        /// The value must be initialized.
+        pub(super) unsafe fn new_unchecked(mutex_guard: MutexGuard<'a, OnceLock<T>>) -> Self {
+            debug_assert!(mutex_guard.get().is_some());
             Self::new(mutex_guard).unwrap_unchecked()
         }
     }
 
-    impl<T> Deref for GlobalGuard<'_, T> {
+    impl<'a, T> Deref for GlobalGuard<'a, T> {
         type Target = T;
-
         fn deref(&self) -> &Self::Target {
-            // SAFETY: `self` is `Initialized`.
-            unsafe { self.mutex_guard.as_initialized().unwrap_unchecked() }
+            // SAFETY: must have been init (see GlobalGuard::new)
+            unsafe { self.mutex_guard.get().unwrap_unchecked() }
         }
     }
 
     impl<T> DerefMut for GlobalGuard<'_, T> {
         fn deref_mut(&mut self) -> &mut Self::Target {
-            // SAFETY: `self` is `Initialized`.
-            unsafe { self.mutex_guard.as_initialized_mut().unwrap_unchecked() }
+            unsafe { self.mutex_guard.get_mut().unwrap_unchecked() }
         }
     }
 }
 
 pub use global_guard::GlobalGuard;
 
-// ----------------------------------------------------------------------------------------------------------------------------------------------
-// Errors
-
 /// Guard that temporarily gives access to a `Global<T>`'s inner value.
 pub enum GlobalLockError<'a, T> {
     /// The mutex is currently locked by another thread.
@@ -194,31 +147,6 @@ pub enum GlobalLockError<'a, T> {
     InitFailed,
 }
 
-// ----------------------------------------------------------------------------------------------------------------------------------------------
-// Internals
-
-enum InitState<T> {
-    Initialized(T),
-    Pending(fn() -> T),
-    Failed,
-}
-
-impl<T> InitState<T> {
-    fn as_initialized(&self) -> Option<&T> {
-        match self {
-            InitState::Initialized(t) => Some(t),
-            _ => None,
-        }
-    }
-
-    fn as_initialized_mut(&mut self) -> Option<&mut T> {
-        match self {
-            InitState::Initialized(t) => Some(t),
-            _ => None,
-        }
-    }
-}
-
 // ----------------------------------------------------------------------------------------------------------------------------------------------
 // Tests
 
@@ -283,6 +211,7 @@ mod tests {
 
     #[test]
     fn test_global_poison() {
+        std::panic::set_hook(Box::new(|_| ()));
         let result = std::panic::catch_unwind(|| {
             let guard = POISON.lock();
             panic!("poison injection");