Rollup merge of #105903 - joboet:unify_parking, r=m-ou-se

Unify id-based thread parking implementations

Multiple platforms currently use thread-id-based parking implementations (NetBSD and SGX[^1]). Even though the strategy does not differ, these are duplicated for each platform, as the id is encoded into an atomic thread variable in different ways for each platform.

Since `park` is only called by one thread, it is possible to move the thread id into a separate field. By ensuring that the field is only written to once, before any other threads access it, these accesses can be unsynchronized, removing any restrictions on the size and niches of the thread id.

This PR also renames the internal `thread_parker` modules to `thread_parking`, as that name now better reflects their contents. I hope this does not add too much reviewing noise.

r? `@m-ou-se`

`@rustbot` label +T-libs

[^1]: SOLID supports this as well, I will switch it over in a follow-up PR.

6 files changed, 119 insertions, 7 deletions

diff --git a/library/std/src/sys_common/mod.rs b/library/std/src/sys_common/mod.rs
index 73da1ce066c..6b24b0e9aa8 100644
--- a/library/std/src/sys_common/mod.rs
+++ b/library/std/src/sys_common/mod.rs
@@ -30,7 +30,7 @@ pub mod process;
 pub mod thread;
 pub mod thread_info;
 pub mod thread_local_dtor;
-pub mod thread_parker;
+pub mod thread_parking;
 pub mod wstr;
 pub mod wtf8;
 
diff --git a/library/std/src/sys_common/thread_parker/futex.rs b/library/std/src/sys_common/thread_parking/futex.rs
index d9e2f39e345..588e7b27826 100644
--- a/library/std/src/sys_common/thread_parker/futex.rs
+++ b/library/std/src/sys_common/thread_parking/futex.rs
@@ -35,7 +35,7 @@ pub struct Parker {
 impl Parker {
     /// Construct the futex parker. The UNIX parker implementation
     /// requires this to happen in-place.
-    pub unsafe fn new(parker: *mut Parker) {
+    pub unsafe fn new_in_place(parker: *mut Parker) {
         parker.write(Self { state: AtomicU32::new(EMPTY) });
     }
 
diff --git a/library/std/src/sys_common/thread_parker/generic.rs b/library/std/src/sys_common/thread_parking/generic.rs
index f3d8b34d3fd..3209bffe353 100644
--- a/library/std/src/sys_common/thread_parker/generic.rs
+++ b/library/std/src/sys_common/thread_parking/generic.rs
@@ -19,7 +19,7 @@ pub struct Parker {
 impl Parker {
     /// Construct the generic parker. The UNIX parker implementation
     /// requires this to happen in-place.
-    pub unsafe fn new(parker: *mut Parker) {
+    pub unsafe fn new_in_place(parker: *mut Parker) {
         parker.write(Parker {
             state: AtomicUsize::new(EMPTY),
             lock: Mutex::new(()),
diff --git a/library/std/src/sys_common/thread_parking/id.rs b/library/std/src/sys_common/thread_parking/id.rs
new file mode 100644
index 00000000000..e98169597c3
--- /dev/null
+++ b/library/std/src/sys_common/thread_parking/id.rs
@@ -0,0 +1,108 @@
+//! Thread parking using thread ids.
+//!
+//! Some platforms (notably NetBSD) have thread parking primitives whose semantics
+//! match those offered by `thread::park`, with the difference that the thread to
+//! be unparked is referenced by a platform-specific thread id. Since the thread
+//! parker is constructed before that id is known, an atomic state variable is used
+//! to manage the park state and propagate the thread id. This also avoids platform
+//! calls in the case where `unpark` is called before `park`.
+
+use crate::cell::UnsafeCell;
+use crate::pin::Pin;
+use crate::sync::atomic::{
+    fence, AtomicI8,
+    Ordering::{Acquire, Relaxed, Release},
+};
+use crate::sys::thread_parking::{current, park, park_timeout, unpark, ThreadId};
+use crate::time::Duration;
+
+pub struct Parker {
+    state: AtomicI8,
+    tid: UnsafeCell<Option<ThreadId>>,
+}
+
+const PARKED: i8 = -1;
+const EMPTY: i8 = 0;
+const NOTIFIED: i8 = 1;
+
+impl Parker {
+    pub fn new() -> Parker {
+        Parker { state: AtomicI8::new(EMPTY), tid: UnsafeCell::new(None) }
+    }
+
+    /// Create a new thread parker. UNIX requires this to happen in-place.
+    pub unsafe fn new_in_place(parker: *mut Parker) {
+        parker.write(Parker::new())
+    }
+
+    /// # Safety
+    /// * must always be called from the same thread
+    /// * must be called before the state is set to PARKED
+    unsafe fn init_tid(&self) {
+        // The field is only ever written to from this thread, so we don't need
+        // synchronization to read it here.
+        if self.tid.get().read().is_none() {
+            // Because this point is only reached once, before the state is set
+            // to PARKED for the first time, the non-atomic write here can not
+            // conflict with reads by other threads.
+            self.tid.get().write(Some(current()));
+            // Ensure that the write can be observed by all threads reading the
+            // state. Synchronizes with the acquire barrier in `unpark`.
+            fence(Release);
+        }
+    }
+
+    pub unsafe fn park(self: Pin<&Self>) {
+        self.init_tid();
+
+        // Changes NOTIFIED to EMPTY and EMPTY to PARKED.
+        let mut state = self.state.fetch_sub(1, Acquire).wrapping_sub(1);
+        if state == PARKED {
+            // Loop to guard against spurious wakeups.
+            while state == PARKED {
+                park(self.state.as_mut_ptr().addr());
+                state = self.state.load(Acquire);
+            }
+
+            // Since the state change has already been observed with acquire
+            // ordering, the state can be reset with a relaxed store instead
+            // of a swap.
+            self.state.store(EMPTY, Relaxed);
+        }
+    }
+
+    pub unsafe fn park_timeout(self: Pin<&Self>, dur: Duration) {
+        self.init_tid();
+
+        let state = self.state.fetch_sub(1, Acquire).wrapping_sub(1);
+        if state == PARKED {
+            park_timeout(dur, self.state.as_mut_ptr().addr());
+            // Swap to ensure that we observe all state changes with acquire
+            // ordering, even if the state has been changed after the timeout
+            // occured.
+            self.state.swap(EMPTY, Acquire);
+        }
+    }
+
+    pub fn unpark(self: Pin<&Self>) {
+        let state = self.state.swap(NOTIFIED, Release);
+        if state == PARKED {
+            // Synchronize with the release fence in `init_tid` to observe the
+            // write to `tid`.
+            fence(Acquire);
+            // # Safety
+            // The thread id is initialized before the state is set to `PARKED`
+            // for the first time and is not written to from that point on
+            // (negating the need for an atomic read).
+            let tid = unsafe { self.tid.get().read().unwrap_unchecked() };
+            // It is possible that the waiting thread woke up because of a timeout
+            // and terminated before this call is made. This call then returns an
+            // error or wakes up an unrelated thread. The platform API and
+            // environment does allow this, however.
+            unpark(tid, self.state.as_mut_ptr().addr());
+        }
+    }
+}
+
+unsafe impl Send for Parker {}
+unsafe impl Sync for Parker {}
diff --git a/library/std/src/sys_common/thread_parker/mod.rs b/library/std/src/sys_common/thread_parking/mod.rs
index 08a2bdd8229..0ead6633c35 100644
--- a/library/std/src/sys_common/thread_parker/mod.rs
+++ b/library/std/src/sys_common/thread_parking/mod.rs
@@ -11,13 +11,17 @@ cfg_if::cfg_if! {
     ))] {
         mod futex;
         pub use futex::Parker;
+    } else if #[cfg(any(
+        target_os = "netbsd",
+        all(target_vendor = "fortanix", target_env = "sgx"),
+    ))] {
+        mod id;
+        pub use id::Parker;
     } else if #[cfg(target_os = "solid_asp3")] {
         mod wait_flag;
         pub use wait_flag::Parker;
     } else if #[cfg(any(windows, target_family = "unix"))] {
-        pub use crate::sys::thread_parker::Parker;
-    } else if #[cfg(all(target_vendor = "fortanix", target_env = "sgx"))] {
-        pub use crate::sys::thread_parker::Parker;
+        pub use crate::sys::thread_parking::Parker;
     } else {
         mod generic;
         pub use generic::Parker;
diff --git a/library/std/src/sys_common/thread_parker/wait_flag.rs b/library/std/src/sys_common/thread_parking/wait_flag.rs
index 6561c186655..d0f8899a94e 100644
--- a/library/std/src/sys_common/thread_parker/wait_flag.rs
+++ b/library/std/src/sys_common/thread_parking/wait_flag.rs
@@ -41,7 +41,7 @@ pub struct Parker {
 impl Parker {
     /// Construct a parker for the current thread. The UNIX parker
     /// implementation requires this to happen in-place.
-    pub unsafe fn new(parker: *mut Parker) {
+    pub unsafe fn new_in_place(parker: *mut Parker) {
         parker.write(Parker { state: AtomicI8::new(EMPTY), wait_flag: WaitFlag::new() })
     }