Windows: Load synch functions together

Attempt to load all the required sync functions and fail if any one of them fails.

This reintroduces a macro for optional loading of functions but keeps it separate from the fallback macro rather than having that do two different jobs.

3 files changed, 94 insertions, 79 deletions

diff --git a/library/std/src/sys/windows/c.rs b/library/std/src/sys/windows/c.rs
index c5a30f8bac8..ef3f6a9ba17 100644
--- a/library/std/src/sys/windows/c.rs
+++ b/library/std/src/sys/windows/c.rs
@@ -228,6 +228,8 @@ pub const IPV6_ADD_MEMBERSHIP: c_int = 12;
 pub const IPV6_DROP_MEMBERSHIP: c_int = 13;
 pub const MSG_PEEK: c_int = 0x2;
 
+pub const LOAD_LIBRARY_SEARCH_SYSTEM32: u32 = 0x800;
+
 #[repr(C)]
 #[derive(Copy, Clone)]
 pub struct linger {
@@ -1030,6 +1032,7 @@ extern "system" {
     pub fn GetProcAddress(handle: HMODULE, name: LPCSTR) -> *mut c_void;
     pub fn GetModuleHandleA(lpModuleName: LPCSTR) -> HMODULE;
     pub fn GetModuleHandleW(lpModuleName: LPCWSTR) -> HMODULE;
+    pub fn LoadLibraryExA(lplibfilename: *const i8, hfile: HANDLE, dwflags: u32) -> HINSTANCE;
 
     pub fn GetSystemTimeAsFileTime(lpSystemTimeAsFileTime: LPFILETIME);
     pub fn GetSystemInfo(lpSystemInfo: LPSYSTEM_INFO);
@@ -1250,21 +1253,16 @@ compat_fn_with_fallback! {
     }
 }
 
-compat_fn_with_fallback! {
-    pub static SYNCH_API: &CStr = ansi_str!("api-ms-win-core-synch-l1-2-0");
-    #[allow(unused)]
-    fn WakeByAddressSingle(Address: LPVOID) -> () {
-        // This fallback is currently tightly coupled to its use in Parker::unpark.
-        //
-        // FIXME: If `WakeByAddressSingle` needs to be used anywhere other than
-        // Parker::unpark then this fallback will be wrong and will need to be decoupled.
-        crate::sys::windows::thread_parker::unpark_keyed_event(Address)
-    }
+compat_fn_optional! {
+    crate::sys::compat::load_synch_functions();
+    pub fn WaitOnAddress(
+        Address: LPVOID,
+        CompareAddress: LPVOID,
+        AddressSize: SIZE_T,
+        dwMilliseconds: DWORD
+    );
+    pub fn WakeByAddressSingle(Address: LPVOID);
 }
-pub use crate::sys::compat::WaitOnAddress;
-// Change exported name of `WakeByAddressSingle` to make the strange fallback
-// behaviour clear.
-pub use WakeByAddressSingle::call as wake_by_address_single_or_unpark_keyed_event;
 
 compat_fn_with_fallback! {
     pub static NTDLL: &CStr = ansi_str!("ntdll");
diff --git a/library/std/src/sys/windows/compat.rs b/library/std/src/sys/windows/compat.rs
index 473544c4d4f..60cc705cef2 100644
--- a/library/std/src/sys/windows/compat.rs
+++ b/library/std/src/sys/windows/compat.rs
@@ -21,6 +21,7 @@
 
 use crate::ffi::{c_void, CStr};
 use crate::ptr::NonNull;
+use crate::sync::atomic::{AtomicBool, Ordering};
 use crate::sys::c;
 
 /// Helper macro for creating CStrs from literals and symbol names.
@@ -74,6 +75,20 @@ impl Module {
         NonNull::new(module).map(Self)
     }
 
+    /// Load the library (if not already loaded)
+    ///
+    /// # Safety
+    ///
+    /// The module must not be unloaded.
+    pub unsafe fn load_system_library(name: &CStr) -> Option<Self> {
+        let module = c::LoadLibraryExA(
+            name.as_ptr(),
+            crate::ptr::null_mut(),
+            c::LOAD_LIBRARY_SEARCH_SYSTEM32,
+        );
+        NonNull::new(module).map(Self)
+    }
+
     // Try to get the address of a function.
     pub fn proc_address(self, name: &CStr) -> Option<NonNull<c_void>> {
         // SAFETY:
@@ -144,61 +159,66 @@ macro_rules! compat_fn_with_fallback {
     )*)
 }
 
-/// Optionally load `WaitOnAddress`.
-/// Unlike the dynamic loading described above, this does not have a fallback.
+/// Optionally loaded functions.
 ///
-/// This is rexported from sys::c. You should prefer to import
-/// from there in case this changes again in the future.
-pub mod WaitOnAddress {
-    use super::*;
-    use crate::mem;
-    use crate::ptr;
-    use crate::sync::atomic::{AtomicBool, AtomicPtr, Ordering};
-    use crate::sys::c;
-
-    static MODULE_NAME: &CStr = ansi_str!("api-ms-win-core-synch-l1-2-0");
-    static SYMBOL_NAME: &CStr = ansi_str!("WaitOnAddress");
-
-    // WaitOnAddress function signature.
-    type F = unsafe extern "system" fn(
-        Address: c::LPVOID,
-        CompareAddress: c::LPVOID,
-        AddressSize: c::SIZE_T,
-        dwMilliseconds: c::DWORD,
-    );
-
-    // A place to store the loaded function atomically.
-    static WAIT_ON_ADDRESS: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());
-
-    // We can skip trying to load again if we already tried.
-    static LOAD_MODULE: AtomicBool = AtomicBool::new(true);
+/// Actual loading of the function defers to $load_functions.
+macro_rules! compat_fn_optional {
+    ($load_functions:expr;
+    $(
+        $(#[$meta:meta])*
+        $vis:vis fn $symbol:ident($($argname:ident: $argtype:ty),*) $(-> $rettype:ty)?;
+    )+) => (
+        $(
+            pub mod $symbol {
+                use super::*;
+                use crate::ffi::c_void;
+                use crate::mem;
+                use crate::ptr::{self, NonNull};
+                use crate::sync::atomic::{AtomicPtr, Ordering};
+
+                pub(in crate::sys) static PTR: AtomicPtr<c_void> = AtomicPtr::new(ptr::null_mut());
+
+                type F = unsafe extern "system" fn($($argtype),*) $(-> $rettype)?;
+
+                #[inline(always)]
+                pub fn option() -> Option<F> {
+                    let f = PTR.load(Ordering::Acquire);
+                    if !f.is_null() { Some(unsafe { mem::transmute(f) }) } else { try_load() }
+                }
 
-    #[inline(always)]
-    pub fn option() -> Option<F> {
-        let f = WAIT_ON_ADDRESS.load(Ordering::Acquire);
-        if !f.is_null() { Some(unsafe { mem::transmute(f) }) } else { try_load() }
+                #[cold]
+                fn try_load() -> Option<F> {
+                    $load_functions;
+                    NonNull::new(PTR.load(Ordering::Acquire)).map(|f| unsafe { mem::transmute(f) })
+                }
+            }
+        )+
+    )
+}
+
+/// Load all needed functions from "api-ms-win-core-synch-l1-2-0".
+pub(super) fn load_synch_functions() {
+    fn try_load() -> Option<()> {
+        static MODULE_NAME: &CStr = ansi_str!("api-ms-win-core-synch-l1-2-0");
+        static WAIT_ON_ADDRESS: &CStr = ansi_str!("WaitOnAddress");
+        static WAKE_BY_ADDRESS_SINGLE: &CStr = ansi_str!("WakeByAddressSingle");
+
+        // Try loading the library and all the required functions.
+        // If any step fails, then they all fail.
+        let library = unsafe { Module::load_system_library(MODULE_NAME) }?;
+        let wait_on_address = library.proc_address(WAIT_ON_ADDRESS)?;
+        let wake_by_address_single = library.proc_address(WAKE_BY_ADDRESS_SINGLE)?;
+
+        c::WaitOnAddress::PTR.store(wait_on_address.as_ptr(), Ordering::Release);
+        c::WakeByAddressSingle::PTR.store(wake_by_address_single.as_ptr(), Ordering::Release);
+        Some(())
     }
 
-    #[cold]
-    fn try_load() -> Option<F> {
-        if LOAD_MODULE.load(Ordering::Acquire) {
-            // load the module
-            let mut wait_on_address = None;
-            if let Some(func) = try_load_inner() {
-                WAIT_ON_ADDRESS.store(func.as_ptr(), Ordering::Release);
-                wait_on_address = Some(unsafe { mem::transmute(func) });
-            }
-            // Don't try to load the module again even if loading failed.
+    // Shortcut if we've already tried (and failed) to load the library.
+    static LOAD_MODULE: AtomicBool = AtomicBool::new(true);
+    if LOAD_MODULE.load(Ordering::Acquire) {
+        if try_load().is_none() {
             LOAD_MODULE.store(false, Ordering::Release);
-            wait_on_address
-        } else {
-            None
         }
     }
-
-    // In the future this could be a `try` block but until then I think it's a
-    // little bit cleaner as a separate function.
-    fn try_load_inner() -> Option<NonNull<c_void>> {
-        unsafe { Module::new(MODULE_NAME)?.proc_address(SYMBOL_NAME) }
-    }
 }
diff --git a/library/std/src/sys/windows/thread_parker.rs b/library/std/src/sys/windows/thread_parker.rs
index 16863c9903a..2f7ae863b6a 100644
--- a/library/std/src/sys/windows/thread_parker.rs
+++ b/library/std/src/sys/windows/thread_parker.rs
@@ -198,8 +198,18 @@ impl Parker {
         // with park().
         if self.state.swap(NOTIFIED, Release) == PARKED {
             unsafe {
-                // This calls either WakeByAddressSingle or unpark_keyed_event (see below).
-                c::wake_by_address_single_or_unpark_keyed_event(self.ptr());
+                if let Some(wake_by_address_single) = c::WakeByAddressSingle::option() {
+                    wake_by_address_single(self.ptr());
+                } else {
+                    // If we run NtReleaseKeyedEvent before the waiting thread runs
+                    // NtWaitForKeyedEvent, this (shortly) blocks until we can wake it up.
+                    // If the waiting thread wakes up before we run NtReleaseKeyedEvent
+                    // (e.g. due to a timeout), this blocks until we do wake up a thread.
+                    // To prevent this thread from blocking indefinitely in that case,
+                    // park_impl() will, after seeing the state set to NOTIFIED after
+                    // waking up, call NtWaitForKeyedEvent again to unblock us.
+                    c::NtReleaseKeyedEvent(keyed_event_handle(), self.ptr(), 0, ptr::null_mut());
+                }
             }
         }
     }
@@ -209,19 +219,6 @@ impl Parker {
     }
 }
 
-// This function signature makes it compatible with c::WakeByAddressSingle
-// so that it can be used as a fallback for that function.
-pub unsafe extern "C" fn unpark_keyed_event(address: c::LPVOID) {
-    // If we run NtReleaseKeyedEvent before the waiting thread runs
-    // NtWaitForKeyedEvent, this (shortly) blocks until we can wake it up.
-    // If the waiting thread wakes up before we run NtReleaseKeyedEvent
-    // (e.g. due to a timeout), this blocks until we do wake up a thread.
-    // To prevent this thread from blocking indefinitely in that case,
-    // park_impl() will, after seeing the state set to NOTIFIED after
-    // waking up, call NtWaitForKeyedEvent again to unblock us.
-    c::NtReleaseKeyedEvent(keyed_event_handle(), address, 0, ptr::null_mut());
-}
-
 fn keyed_event_handle() -> c::HANDLE {
     const INVALID: c::HANDLE = ptr::invalid_mut(!0);
     static HANDLE: AtomicPtr<libc::c_void> = AtomicPtr::new(INVALID);