simple atomic lock

.gitignore

@@ -1 +1,2 @@
 /target
+Cargo.lock

Cargo.lock

@@ -2,6 +2,82 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "atomic-wait"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a55b94919229f2c42292fd71ffa4b75e83193bffdd77b1e858cd55fd2d0b0ea8"
+dependencies = [
+ "libc",
+ "windows-sys",
+]
+
+[[package]]
+name = "libc"
+version = "0.2.174"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1171693293099992e19cddea4e8b849964e9846f4acee11b3948bcc337be8776"
+
 [[package]]
 name = "werkzeug"
 version = "0.1.0"
+dependencies = [
+ "atomic-wait",
+]
+
+[[package]]
+name = "windows-sys"
+version = "0.42.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "5a3e1820f08b8513f676f7ab6c1f99ff312fb97b553d30ff4dd86f9f15728aa7"
+dependencies = [
+ "windows_aarch64_gnullvm",
+ "windows_aarch64_msvc",
+ "windows_i686_gnu",
+ "windows_i686_msvc",
+ "windows_x86_64_gnu",
+ "windows_x86_64_gnullvm",
+ "windows_x86_64_msvc",
+]
+
+[[package]]
+name = "windows_aarch64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8"
+
+[[package]]
+name = "windows_aarch64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43"
+
+[[package]]
+name = "windows_i686_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f"
+
+[[package]]
+name = "windows_i686_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060"
+
+[[package]]
+name = "windows_x86_64_gnu"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36"
+
+[[package]]
+name = "windows_x86_64_gnullvm"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3"
+
+[[package]]
+name = "windows_x86_64_msvc"
+version = "0.42.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0"

Cargo.toml

@@ -10,3 +10,8 @@ std = []
 nightly = []
 
 [dependencies]
+# libc = "0.2"
+
+# While I could use libc / windows for this, why not just use this tiny crate
+# which does exactly and only a futex
+atomic-wait = "1.1.0"

rust-toolchain

@@ -0,0 +1 @@
+nightly

src/lib.rs

@@ -13,6 +13,7 @@ pub mod ptr;
 pub mod rand;
 #[cfg(feature = "alloc")]
 pub mod smallbox;
+pub mod sync;
 pub mod util;
 
 pub use cachepadded::CachePadded;

src/sync.rs

@@ -0,0 +1,170 @@
+use core::{
+    mem,
+    sync::atomic::{AtomicU32, Ordering},
+};
+
+const LOCKED_BIT: u32 = 0b001;
+const EMPTY: u32 = 0;
+
+/// A simple lock implementation using an atomic u32.
+#[repr(transparent)]
+pub struct Lock {
+    inner: AtomicU32,
+}
+
+impl Lock {
+    /// Creates a new lock in the unlocked state.
+    pub const fn new() -> Self {
+        Self {
+            inner: AtomicU32::new(0),
+        }
+    }
+
+    pub fn as_ptr(&self) -> *mut u32 {
+        self.inner.as_ptr()
+    }
+
+    pub unsafe fn from_ptr<'a>(ptr: *mut u32) -> &'a Self {
+        // SAFETY: The caller must ensure that `ptr` is not aliased, and lasts
+        // for the lifetime of the `Lock`.
+        unsafe { mem::transmute(AtomicU32::from_ptr(ptr)) }
+    }
+
+    /// Acquires the lock, blocking until it is available.
+    pub fn lock(&self) {
+        // attempt acquiring the lock with no contention.
+        if self
+            .inner
+            .compare_exchange_weak(EMPTY, LOCKED_BIT, Ordering::Acquire, Ordering::Relaxed)
+            .is_ok()
+        {
+            // We successfully acquired the lock.
+            return;
+        } else {
+            self.lock_slow();
+        }
+    }
+
+    pub fn unlock(&self) {
+        self.inner.fetch_and(!LOCKED_BIT, Ordering::Release);
+    }
+
+    fn lock_slow(&self) {
+        // The lock is either locked, or someone is waiting for it:
+
+        let mut spin_wait = SpinWait::new();
+        let mut state = self.inner.load(Ordering::Acquire);
+        loop {
+            // If the lock isn't locked, we can try to acquire it.
+            if state & LOCKED_BIT == 0 {
+                // Try to acquire the lock.
+                match self.inner.compare_exchange_weak(
+                    state,
+                    state | LOCKED_BIT,
+                    Ordering::Acquire,
+                    Ordering::Relaxed,
+                ) {
+                    Ok(_) => {
+                        // We successfully acquired the lock.
+                        return;
+                    }
+                    Err(new_state) => {
+                        // The state changed, we need to check again.
+                        state = new_state;
+                        continue;
+                    }
+                }
+            }
+
+            if {
+                let spun: bool;
+                #[cfg(feature = "std")]
+                {
+                    spun = spin_wait.spin_yield();
+                }
+                #[cfg(not(feature = "std"))]
+                {
+                    spun = spin_wait.spin();
+                }
+
+                spun
+            } {
+                // We can spin for a little while and see if it becomes available.
+                state = self.inner.load(Ordering::Relaxed);
+                continue;
+            }
+
+            // If we reach here, we need to park the thread.
+            atomic_wait::wait(&self.inner, LOCKED_BIT);
+
+            if self
+                .inner
+                .compare_exchange_weak(
+                    state,
+                    state | LOCKED_BIT,
+                    Ordering::Acquire,
+                    Ordering::Relaxed,
+                )
+                .is_ok()
+            {
+                // We successfully acquired the lock after being woken up.
+                return;
+            }
+
+            spin_wait.reset();
+            state = self.inner.load(Ordering::Relaxed);
+        }
+    }
+}
+
+pub struct SpinWait {
+    counter: u32,
+}
+
+impl SpinWait {
+    /// Creates a new `SpinWait` with an initial counter value.
+    pub const fn new() -> Self {
+        Self { counter: 0 }
+    }
+
+    /// Resets the counter to zero.
+    pub fn reset(&mut self) {
+        self.counter = 0;
+    }
+
+    pub fn spin(&mut self) -> bool {
+        if self.counter >= 10 {
+            // If the counter is too high, we signal the caller to potentially park.
+            return false;
+        }
+        self.counter += 1;
+
+        // spin for a small number of iterations based on the counter value.
+        for _ in 0..(1 << self.counter) {
+            core::hint::spin_loop();
+        }
+
+        true
+    }
+
+    #[cfg(feature = "std")]
+    pub fn spin_yield(&mut self) -> bool {
+        if self.counter >= 10 {
+            // If the counter is too high, we signal the caller to potentially park.
+            return false;
+        }
+        self.counter += 1;
+
+        if self.counter >= 3 {
+            // spin for a small number of iterations based on the counter value.
+            for _ in 0..(1 << self.counter) {
+                core::hint::spin_loop();
+            }
+        } else {
+            // yield the thread and wait for the OS to reschedule us.
+            std::thread::yield_now();
+        }
+
+        true
+    }
+}