avoid cancel_handle close race with weakref.finalize

trio/_core/_io_windows.py

@@ -20,6 +20,7 @@
     INVALID_HANDLE_VALUE,
     raise_winerror,
     _handle,
+    _check,
     ErrorCodes,
     FileFlags,
     AFDPollFlags,
@@ -180,12 +181,6 @@ class CKeys(enum.IntEnum):
     USER_DEFINED = 4  # and above
 
 
-def _check(success):
-    if not success:
-        raise_winerror()
-    return success
-
-
 def _get_underlying_socket(sock, *, which=WSAIoctls.SIO_BASE_HANDLE):
     if hasattr(sock, "fileno"):
         sock = sock.fileno()

trio/_core/_windows_cffi.py

@@ -321,3 +321,9 @@ def raise_winerror(winerror=None, *, filename=None, filename2=None):
         _, msg = ffi.getwinerror(winerror)
     # https://docs.python.org/3/library/exceptions.html#OSError
     raise OSError(0, msg, filename, winerror, filename2)
+
+
+def _check(success):
+    if not success:
+        raise_winerror()
+    return success

trio/_wait_for_object.py

@@ -1,15 +1,18 @@
-import math
-from . import _timeouts
+import weakref
+
 import trio
+from ._core import enable_ki_protection
 from ._core._windows_cffi import (
     ffi,
     kernel32,
     ErrorCodes,
     raise_winerror,
     _handle,
+    _check,
 )
 
 
+@enable_ki_protection
 async def WaitForSingleObject(obj):
     """Async and cancellable variant of WaitForSingleObject. Windows only.
 
@@ -30,24 +33,36 @@ async def WaitForSingleObject(obj):
     if retcode == ErrorCodes.WAIT_FAILED:
         raise_winerror()
     elif retcode != ErrorCodes.WAIT_TIMEOUT:
+        await trio.lowlevel.checkpoint()
         return
 
     # Wait for a thread that waits for two handles: the handle plus a handle
     # that we can use to cancel the thread.
-    cancel_handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    cancel_handle = _check(kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL))
+
+    # Close the handle via weakref because (a) managing within Trio had a tricky
+    # race condition (GH-1271) and (b) we don't care precisely when it happens
+    cancel_handle_address = int(ffi.cast("intptr_t", cancel_handle))
+
+    def cleanup():
+        # regenerate the cancel_handle from an int, so we don't create
+        # a reference cycle that keeps cancel_handle alive
+        _check(kernel32.CloseHandle(_handle(cancel_handle_address)))
+
+    weakref.finalize(cancel_handle, cleanup)
+
     try:
         await trio.to_thread.run_sync(
             WaitForMultipleObjects_sync,
             handle,
             cancel_handle,
             cancellable=True,
-            limiter=trio.CapacityLimiter(math.inf),
+            limiter=trio.CapacityLimiter(1),
         )
     finally:
-        # Clean up our cancel handle. In case we get here because this task was
-        # cancelled, we also want to set the cancel_handle to stop the thread.
-        kernel32.SetEvent(cancel_handle)
-        kernel32.CloseHandle(cancel_handle)
+        # Clean up our thread. If we get here because this task was cancelled,
+        # we also want to set the cancel handle to stop the thread.
+        _check(kernel32.SetEvent(cancel_handle))
 
 
 def WaitForMultipleObjects_sync(*handles):

trio/tests/test_wait_for_object.py

@@ -211,8 +211,9 @@ async def signal_soon_async(handle):
     handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
     t0 = _core.current_time()
 
-    with _timeouts.move_on_after(TIMEOUT):
+    with _timeouts.move_on_after(TIMEOUT) as cs:
         await WaitForSingleObject(handle)
+    assert cs.cancelled_caught
 
     kernel32.CloseHandle(handle)
     t1 = _core.current_time()