From f1c99041c25b08821a0a18fd677f008da24ac2ad Mon Sep 17 00:00:00 2001
From: Christopher Berner <me@cberner.com>
Date: Fri, 18 Oct 2024 21:08:25 -0700
Subject: [PATCH] Support lock() and lock_shared() on async IO Files

---
 std/src/fs/tests.rs           | 38 +++++++++++++++++++++
 std/src/sys/pal/windows/fs.rs | 62 +++++++++++++++++++++++++++--------
 2 files changed, 87 insertions(+), 13 deletions(-)

diff --git a/std/src/fs/tests.rs b/std/src/fs/tests.rs
index fe3dc4adff0ba..05efed6b5dfc6 100644
--- a/std/src/fs/tests.rs
+++ b/std/src/fs/tests.rs
@@ -283,6 +283,44 @@ fn file_lock_double_unlock() {
     assert!(check!(f2.try_lock()));
 }
 
+#[test]
+#[cfg(windows)]
+fn file_lock_blocking_async() {
+    use crate::thread::{sleep, spawn};
+    const FILE_FLAG_OVERLAPPED: u32 = 0x40000000;
+
+    let tmpdir = tmpdir();
+    let filename = &tmpdir.join("file_lock_blocking_async.txt");
+    let f1 = check!(File::create(filename));
+    let f2 =
+        check!(OpenOptions::new().custom_flags(FILE_FLAG_OVERLAPPED).write(true).open(filename));
+
+    check!(f1.lock());
+
+    // Ensure that lock() is synchronous when the file is opened for asynchronous IO
+    let t = spawn(move || {
+        check!(f2.lock());
+    });
+    sleep(Duration::from_secs(1));
+    assert!(!t.is_finished());
+    check!(f1.unlock());
+    t.join().unwrap();
+
+    // Ensure that lock_shared() is synchronous when the file is opened for asynchronous IO
+    let f2 =
+        check!(OpenOptions::new().custom_flags(FILE_FLAG_OVERLAPPED).write(true).open(filename));
+    check!(f1.lock());
+
+    // Ensure that lock() is synchronous when the file is opened for asynchronous IO
+    let t = spawn(move || {
+        check!(f2.lock_shared());
+    });
+    sleep(Duration::from_secs(1));
+    assert!(!t.is_finished());
+    check!(f1.unlock());
+    t.join().unwrap();
+}
+
 #[test]
 fn file_test_io_seek_shakedown() {
     //                   01234567890123
diff --git a/std/src/sys/pal/windows/fs.rs b/std/src/sys/pal/windows/fs.rs
index bc728d498c3c3..138ac59c1613b 100644
--- a/std/src/sys/pal/windows/fs.rs
+++ b/std/src/sys/pal/windows/fs.rs
@@ -346,27 +346,63 @@ impl File {
         self.fsync()
     }
 
-    pub fn lock(&self) -> io::Result<()> {
-        cvt(unsafe {
-            let mut overlapped = mem::zeroed();
-            c::LockFileEx(
+    fn acquire_lock(&self, flags: c::LOCK_FILE_FLAGS) -> io::Result<()> {
+        unsafe {
+            let mut overlapped: c::OVERLAPPED = mem::zeroed();
+            let event = c::CreateEventW(ptr::null_mut(), c::FALSE, c::FALSE, ptr::null());
+            if event.is_null() {
+                return Err(io::Error::last_os_error());
+            }
+            overlapped.hEvent = event;
+            let lock_result = cvt(c::LockFileEx(
                 self.handle.as_raw_handle(),
-                c::LOCKFILE_EXCLUSIVE_LOCK,
+                flags,
                 0,
                 u32::MAX,
                 u32::MAX,
                 &mut overlapped,
-            )
-        })?;
-        Ok(())
+            ));
+
+            let final_result = match lock_result {
+                Ok(_) => Ok(()),
+                Err(err) => {
+                    if err.raw_os_error() == Some(c::ERROR_IO_PENDING as i32) {
+                        // Wait for the lock to be acquired. This can happen asynchronously,
+                        // if the file handle was opened for async IO
+                        let wait_result = c::WaitForSingleObject(overlapped.hEvent, c::INFINITE);
+                        if wait_result == c::WAIT_OBJECT_0 {
+                            // Wait completed successfully, get the lock operation status
+                            let mut bytes_transferred = 0;
+                            cvt(c::GetOverlappedResult(
+                                self.handle.as_raw_handle(),
+                                &mut overlapped,
+                                &mut bytes_transferred,
+                                c::TRUE,
+                            ))
+                            .map(|_| ())
+                        } else if wait_result == c::WAIT_FAILED {
+                            // Wait failed
+                            Err(io::Error::last_os_error())
+                        } else {
+                            // WAIT_ABANDONED and WAIT_TIMEOUT should not be possible
+                            unreachable!()
+                        }
+                    } else {
+                        Err(err)
+                    }
+                }
+            };
+            c::CloseHandle(overlapped.hEvent);
+            final_result
+        }
+    }
+
+    pub fn lock(&self) -> io::Result<()> {
+        self.acquire_lock(c::LOCKFILE_EXCLUSIVE_LOCK)
     }
 
     pub fn lock_shared(&self) -> io::Result<()> {
-        cvt(unsafe {
-            let mut overlapped = mem::zeroed();
-            c::LockFileEx(self.handle.as_raw_handle(), 0, 0, u32::MAX, u32::MAX, &mut overlapped)
-        })?;
-        Ok(())
+        self.acquire_lock(0)
     }
 
     pub fn try_lock(&self) -> io::Result<bool> {