Add wrapper around dispatch_once_f

madsmtm · madsmtm · commit 32453189250c · 2024-12-23T16:12:52.000+01:00
Part of #77.
diff --git a/crates/dispatch2/CHANGELOG.md b/crates/dispatch2/CHANGELOG.md
@@ -46,7 +46,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
   - `dispatch_write`.
 - Added `#[must_use]` in bindings where the source uses it.
 - Added `Queue::main` for fetching the main queue.
-- Moved `run_on_main` and `MainThreadBound` from `objc2-foundation` to this crate.
+- Moved `run_on_main` and `MainThreadBound` from `objc2-foundation` to this
+  crate.
+- Added `Once`, a wrapper over `dispatch_once_f` which works similarly to
+  `std::sync::Once`.
 
 ### Changed
 - Moved to the `objc2` project.
diff --git a/crates/dispatch2/TODO.md b/crates/dispatch2/TODO.md
@@ -6,7 +6,6 @@
 - CI test on Windows using https://github.com/apple/swift-corelibs-libdispatch
 - Safe wrapper for ``dispatch_source_*`` + ``set_target_queue/activate/suspend/resume`` for it
 - Safe wrapper for ``dispatch_data_*``
-- Safe wrapper for ``dispatch_once_f`` (is that relevant?)
 - Safe wrapper for ``dispatch_get_context/dispatch_set_context`` (quite impossible without big overhead => wrap dispatch object destructor to release the boxed value)
 - All blocks related bindings and ``dispatch_block_*`` functions with compat with ``block2`` on Apple platforms.
 - Integrate conversion from SystemTime to dispatch_time_t via dispatch_walltime and safe APIs using that.
diff --git a/crates/dispatch2/src/lib.rs b/crates/dispatch2/src/lib.rs
@@ -33,6 +33,7 @@ pub mod group;
 #[cfg(feature = "objc2")]
 mod main_thread_bound;
 pub mod object;
+mod once;
 pub mod queue;
 pub mod semaphore;
 mod utils;
@@ -85,5 +86,6 @@ pub use self::group::*;
 #[cfg(feature = "objc2")]
 pub use self::main_thread_bound::{run_on_main, MainThreadBound};
 pub use self::object::*;
+pub use self::once::*;
 pub use self::queue::*;
 pub use self::semaphore::*;
diff --git a/crates/dispatch2/src/once.rs b/crates/dispatch2/src/once.rs
@@ -0,0 +1,336 @@
+use core::ffi::c_void;
+use core::fmt;
+use core::ptr::NonNull;
+use std::cell::UnsafeCell;
+use std::panic::{RefUnwindSafe, UnwindSafe};
+use std::sync::atomic::AtomicIsize;
+use std::sync::atomic::Ordering;
+
+use crate::ffi;
+
+/// A low-level synchronization primitive for one-time global execution.
+///
+/// This is equivalent to [`std::sync::Once`], except that this uses the
+/// underlying system primitives from `libdispatch`, which:
+/// - Might result in less code-size overhead.
+/// - Aborts on panics in the initialization closure.
+///
+/// Generally, prefer [`std::sync::Once`] unless you have a specific need for
+/// this.
+///
+///
+/// # Example
+///
+/// Run a closure once for the duration of the entire program, without using
+/// [`std::sync::Once`].
+///
+/// ```
+/// use dispatch2::Once;
+///
+/// static INIT: Once = Once::new();
+///
+/// INIT.call_once(|| {
+///     // run initialization here
+/// });
+/// ```
+#[doc(alias = "dispatch_once_t")]
+pub struct Once {
+    predicate: UnsafeCell<ffi::dispatch_once_t>,
+}
+
+// This is intentionally `extern "C"`, since libdispatch will not propagate an
+// internal panic, but will simply abort.
+extern "C" fn invoke_closure<F>(context: *mut c_void)
+where
+    F: FnOnce(),
+{
+    let context: *mut Option<F> = context.cast();
+    // SAFETY: Context was created below in `invoke_dispatch_once` from
+    // `&mut Option<F>`.
+    let closure: &mut Option<F> = unsafe { &mut *context };
+
+    // SAFETY: libdispatch is implemented correctly, and will only call this
+    // once (and we set it to be available before calling dispatch_once).
+    let closure = unsafe { closure.take().unwrap_unchecked() };
+
+    (closure)();
+}
+
+#[cfg_attr(
+    // DISPATCH_ONCE_INLINE_FASTPATH, see Once::call_once below.
+    any(target_arch = "x86", target_arch = "x86_64", target_vendor = "apple"),
+    cold,
+    inline(never)
+)]
+fn invoke_dispatch_once<F>(predicate: NonNull<ffi::dispatch_once_t>, closure: F)
+where
+    F: FnOnce(),
+{
+    // Convert closure data to context parameter.
+    let mut closure = Some(closure);
+    let context: *mut Option<F> = &mut closure;
+    let context: *mut c_void = context.cast();
+
+    // SAFETY: The function and context are valid, and the predicate pointer
+    // is valid.
+    //
+    // NOTE: The documentation says:
+    // > The predicate must point to a variable stored in global or static
+    // > scope. The result of using a predicate with automatic or dynamic
+    // > storage (including Objective-C instance variables) is undefined.
+    //
+    // In Rust though, we have stronger guarantees, and can guarantee that the
+    // predicate is never moved while in use, because the `Once` itself is not
+    // cloneable.
+    //
+    // Even if libdispatch may sometimes use the pointer as a condition
+    // variable, or may internally store a self-referential pointer, it can
+    // only do that while the Once is in use somewhere (i.e. it should not be
+    // able to do that while the Once is being moved).
+    //
+    // Outside of being moved, the Once can only be in two states:
+    // - Initialized.
+    // - Done.
+    //
+    // And those two states are freely movable.
+    unsafe { ffi::dispatch_once_f(predicate, context, invoke_closure::<F>) };
+
+    // Closure is dropped here, depending on if it was executed (and taken
+    // from the `Option`) by `dispatch_once_f` or not.
+}
+
+impl Once {
+    /// Creates a new `Once`.
+    #[inline]
+    #[allow(clippy::new_without_default)] // `std::sync::Once` doesn't have it either
+    pub const fn new() -> Self {
+        Self {
+            predicate: UnsafeCell::new(0),
+        }
+    }
+
+    /// Executes a closure once for the lifetime of the application.
+    ///
+    /// If called simultaneously from multiple threads, this function waits
+    /// synchronously until the work function has completed.
+    ///
+    ///
+    /// # Aborts
+    ///
+    /// The process will trap or abort if:
+    /// - The given initialization closure unwinds.
+    /// - The given closure recursively invokes `call_once` on the same
+    ///   `Once` instance.
+    #[inline]
+    #[doc(alias = "dispatch_once")]
+    #[doc(alias = "dispatch_once_f")]
+    pub fn call_once<F>(&self, work: F)
+    where
+        F: FnOnce(),
+    {
+        // Unwrap is fine, the pointer is valid so can never be NULL.
+        let predicate = NonNull::new(self.predicate.get()).unwrap();
+
+        // DISPATCH_ONCE_INLINE_FASTPATH
+        if cfg!(any(
+            target_arch = "x86",
+            target_arch = "x86_64",
+            target_vendor = "apple"
+        )) {
+            // On certain platforms, the ABI of the predicate is stable enough
+            // that we are allowed to read it to check if the condition is
+            // done yet.
+            //
+            // The code in C is inside `_dispatch_once_f` in dispatch/once.h:
+            //
+            //     if (DISPATCH_EXPECT(*predicate, ~0l) != ~0l) {
+            //         dispatch_once_f(predicate, context, function);
+            //     } else {
+            //         dispatch_compiler_barrier();
+            //     }
+            //     DISPATCH_COMPILER_CAN_ASSUME(*predicate == ~0l);
+
+            // NOTE: To uphold the rules set by the Rust AM, we use an atomic
+            // comparison here to avoid a possible tear, even though the
+            // equivalent C code just loads the predicate un-atomically.
+            //
+            // SAFETY: The predicate is a valid atomic pointer.
+            // TODO: Use `AtomicIsize::from_ptr` once in MSRV.
+            let atomic_predicate: &AtomicIsize = unsafe { predicate.cast().as_ref() };
+
+            // We use an acquire load, as that's what's done internally in
+            // libdispatch, and matches what's done in Rust's std too:
+            // <https://github.com/swiftlang/swift-corelibs-libdispatch/blob/swift-6.0.3-RELEASE/src/once.c#L57>
+            // <https://github.com/rust-lang/rust/blob/1.83.0/library/std/src/sys/sync/once/queue.rs#L130>
+            if atomic_predicate.load(Ordering::Acquire) != !0 {
+                invoke_dispatch_once(predicate, work);
+            }
+
+            // NOTE: Unlike in C, we cannot use `std::hint::assert_unchecked`,
+            // since that would actually be lying from a language perspective;
+            // the value seems to only settle on being !0 after some time
+            // (something about the _COMM_PAGE_CPU_QUIESCENT_COUNTER?)
+            //
+            // TODO: Investigate this further!
+            // std::hint::assert_unchecked(atomic_predicate.load(Ordering::Acquire) == !0);
+        } else {
+            invoke_dispatch_once(predicate, work);
+        }
+    }
+}
+
+// SAFETY: Same as `std::sync::Once`
+unsafe impl Send for Once {}
+
+// SAFETY: Same as `std::sync::Once`
+unsafe impl Sync for Once {}
+
+impl UnwindSafe for Once {}
+impl RefUnwindSafe for Once {}
+
+impl fmt::Debug for Once {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("Once").finish_non_exhaustive()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::cell::Cell;
+    use std::mem::ManuallyDrop;
+    use std::thread;
+
+    use super::*;
+
+    #[test]
+    fn test_static() {
+        static ONCE: Once = Once::new();
+        let mut num = 0;
+        ONCE.call_once(|| num += 1);
+        ONCE.call_once(|| num += 1);
+        assert!(num == 1);
+    }
+
+    #[test]
+    fn test_in_loop() {
+        let once = Once::new();
+
+        let mut call_count = 0;
+        for _ in 0..10 {
+            once.call_once(|| call_count += 1);
+        }
+
+        assert_eq!(call_count, 1);
+    }
+
+    #[test]
+    fn test_move() {
+        let once = Once::new();
+
+        let mut call_count = 0;
+        for _ in 0..10 {
+            once.call_once(|| call_count += 1);
+        }
+
+        let once = once;
+        for _ in 0..10 {
+            once.call_once(|| call_count += 1);
+        }
+
+        let once = Once {
+            predicate: UnsafeCell::new(once.predicate.into_inner()),
+        };
+        for _ in 0..10 {
+            once.call_once(|| call_count += 1);
+        }
+
+        assert_eq!(call_count, 1);
+    }
+
+    #[test]
+    fn test_threaded() {
+        let once = Once::new();
+
+        let num = AtomicIsize::new(0);
+
+        thread::scope(|scope| {
+            scope.spawn(|| {
+                once.call_once(|| {
+                    num.fetch_add(1, Ordering::Relaxed);
+                });
+            });
+            scope.spawn(|| {
+                once.call_once(|| {
+                    num.fetch_add(1, Ordering::Relaxed);
+                });
+            });
+            scope.spawn(|| {
+                once.call_once(|| {
+                    num.fetch_add(1, Ordering::Relaxed);
+                });
+            });
+        });
+
+        assert!(num.load(Ordering::Relaxed) == 1);
+    }
+
+    #[derive(Clone)]
+    struct DropTest<'a>(&'a Cell<usize>);
+
+    impl Drop for DropTest<'_> {
+        fn drop(&mut self) {
+            self.0.set(self.0.get() + 1);
+        }
+    }
+
+    #[test]
+    fn test_drop_in_closure() {
+        let amount_of_drops = Cell::new(0);
+        let once = Once::new();
+
+        let tester = DropTest(&amount_of_drops);
+        once.call_once(move || {
+            let _tester = tester;
+        });
+        assert_eq!(amount_of_drops.get(), 1);
+
+        let tester = DropTest(&amount_of_drops);
+        once.call_once(move || {
+            let _tester = tester;
+        });
+        assert_eq!(amount_of_drops.get(), 2);
+    }
+
+    #[test]
+    fn test_drop_in_closure_with_leak() {
+        let amount_of_drops = Cell::new(0);
+        let once = Once::new();
+
+        // Not dropped here, since we ManuallyDrop inside the closure (and the
+        // closure is executed).
+        let tester = DropTest(&amount_of_drops);
+        once.call_once(move || {
+            let _tester = ManuallyDrop::new(tester);
+        });
+        assert_eq!(amount_of_drops.get(), 0);
+
+        // Still dropped here, since the once is not executed
+        let tester = DropTest(&amount_of_drops);
+        once.call_once(move || {
+            let _tester = ManuallyDrop::new(tester);
+        });
+        assert_eq!(amount_of_drops.get(), 1);
+    }
+
+    #[test]
+    #[ignore = "traps the process (as expected)"]
+    fn test_recursive_invocation() {
+        let once = Once::new();
+
+        once.call_once(|| {
+            once.call_once(|| {
+                println!("foo");
+            });
+        });
+    }
+}
diff --git a/crates/dispatch2/translation-config.toml b/crates/dispatch2/translation-config.toml
@@ -68,3 +68,7 @@ typedef.dispatch_workloop_t.skipped = true
 
 # Inline, defined manually for now
 fn.dispatch_get_main_queue.skipped = true
+
+# Fast paths for dispatch_once, done manually
+fn._dispatch_once.skipped = true
+fn._dispatch_once_f.skipped = true
diff --git a/generated b/generated
@@ -1 +1 @@
-Subproject commit 4ff3aeed44b045ea5a9332679cf8591ce8ec28cb
+Subproject commit 1c10c0d11e922302f9a69f8627e6f201aef5d65a
