Implement cancelable WaitForSingleObject for Windows

trio/_core/_io_windows.py

@@ -12,6 +12,7 @@
 import attr
 
 from .. import _core
+from .. import _timeouts
 from . import _public
 from ._wakeup_socketpair import WakeupSocketpair
 from .._util import is_main_thread
@@ -109,6 +110,65 @@ def _handle(obj):
         return obj
 
 
+async def WaitForSingleObject(handle):
+    """Async and cancellable variant of kernel32.WaitForSingleObject().
+
+    Args:
+      handle: A win32 handle.
+
+    """
+    # Quick check; we might not even need to spawn a thread. The zero
+    # means a zero timeout; this call never blocks. We also exit here
+    # if the handle is already closed for some reason.
+    retcode = kernel32.WaitForSingleObject(handle, 0)
+    if retcode != ErrorCodes.WAIT_TIMEOUT:
+        return
+
+    # :'(  avoid circular imports
+    from .._threads import run_sync_in_worker_thread
+
+    class StubLimiter:
+        def release_on_behalf_of(self, x):
+            pass
+
+        async def acquire_on_behalf_of(self, x):
+            pass
+
+    # 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)
+    try:
+        await run_sync_in_worker_thread(
+            WaitForMultipleObjects_sync,
+            handle,
+            cancel_handle,
+            cancellable=True,
+            limiter=StubLimiter(),
+        )
+    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)
+
+
+def WaitForMultipleObjects_sync(*handles):
+    """Wait for any of the given Windows handles to be signaled.
+
+    """
+    n = len(handles)
+    handle_arr = ffi.new("HANDLE[{}]".format(n))
+    for i in range(n):
+        handle_arr[i] = handles[i]
+    timeout = 1000 * 60 * 60 * 24  # todo: use INF here, whatever that is, and ditch the while
+    while True:
+        retcode = kernel32.WaitForMultipleObjects(
+            n, handle_arr, False, timeout
+        )
+        if retcode != ErrorCodes.WAIT_TIMEOUT:
+            break
+
+
 @attr.s(frozen=True)
 class _WindowsStatistics:
     tasks_waiting_overlapped = attr.ib()

trio/_core/_windows_cffi.py

@@ -35,6 +35,8 @@
 
 typedef OVERLAPPED WSAOVERLAPPED;
 typedef LPOVERLAPPED LPWSAOVERLAPPED;
+typedef PVOID LPSECURITY_ATTRIBUTES;
+typedef PVOID LPCSTR;
 
 typedef struct _OVERLAPPED_ENTRY {
     ULONG_PTR lpCompletionKey;
@@ -80,6 +82,34 @@
   _In_opt_ void*            HandlerRoutine,
   _In_     BOOL             Add
 );
+
+HANDLE CreateEventA(
+  LPSECURITY_ATTRIBUTES lpEventAttributes,
+  BOOL                  bManualReset,
+  BOOL                  bInitialState,
+  LPCSTR                lpName
+);
+
+BOOL SetEvent(
+  HANDLE hEvent
+);
+
+BOOL ResetEvent(
+  HANDLE hEvent
+);
+
+DWORD WaitForSingleObject(
+  HANDLE hHandle,
+  DWORD  dwMilliseconds
+);
+
+DWORD WaitForMultipleObjects(
+  DWORD        nCount,
+  HANDLE       *lpHandles,
+  BOOL         bWaitAll,
+  DWORD        dwMilliseconds
+);
+
 """
 
 # cribbed from pywincffi
@@ -116,6 +146,10 @@ def raise_winerror(winerror=None, *, filename=None, filename2=None):
 
 class ErrorCodes(enum.IntEnum):
     STATUS_TIMEOUT = 0x102
+    WAIT_TIMEOUT = 0x102
+    WAIT_ABANDONED = 0x80
+    WAIT_OBJECT_0 = 0x00  # object is signaled
+    WAIT_FAILED = 0xFFFFFFFF
     ERROR_IO_PENDING = 997
     ERROR_OPERATION_ABORTED = 995
     ERROR_ABANDONED_WAIT_0 = 735

trio/_core/tests/test_windows.py

@@ -1,13 +1,16 @@
 import os
+from threading import Thread
 import pytest
 
 on_windows = (os.name == "nt")
 # Mark all the tests in this file as being windows-only
 pytestmark = pytest.mark.skipif(not on_windows, reason="windows only")
 
 from ... import _core
+from ... import _timeouts
 if on_windows:
     from .._windows_cffi import ffi, kernel32
+    from .._io_windows import WaitForSingleObject, WaitForMultipleObjects_sync
 
 
 async def test_completion_key_listen():
@@ -38,5 +41,109 @@ async def post(key):
             print("end loop")
 
 
+async def test_WaitForMultipleObjects_sync():
+    # One handle
+    handle1 = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    t = Thread(target=WaitForMultipleObjects_sync, args=(handle1,))
+    t.start()
+    kernel32.SetEvent(handle1)
+    t.join()  # the test succeeds if we do not block here :)
+    kernel32.CloseHandle(handle1)
+
+    # Two handles, signal first
+    handle1 = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    handle2 = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    t = Thread(target=WaitForMultipleObjects_sync, args=(handle1, handle2))
+    t.start()
+    kernel32.SetEvent(handle1)
+    t.join()  # the test succeeds if we do not block here :)
+    kernel32.CloseHandle(handle1)
+    kernel32.CloseHandle(handle2)
+
+    # Two handles, signal seconds
+    handle1 = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    handle2 = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    t = Thread(target=WaitForMultipleObjects_sync, args=(handle1, handle2))
+    t.start()
+    kernel32.SetEvent(handle2)
+    t.join()  # the test succeeds if we do not block here :)
+    kernel32.CloseHandle(handle1)
+    kernel32.CloseHandle(handle2)
+
+    # Closing the handle will not stop the thread. Initiating a wait on a
+    # closed handle will fail/return, but closing a handle that is already
+    # being waited on will not stop whatever is waiting for it.
+
+
+async def test_WaitForSingleObject():
+
+    # Set the timeout used in the tests. The resolution of WaitForSingleObject
+    # is 0.01 so anything more than a magnitude larger should probably do.
+    # If too large, the test become slow and we might need to mark it as @slow.
+    TIMEOUT = 0.5
+
+    async def handle_setter(handle):
+        await _timeouts.sleep(TIMEOUT)
+        kernel32.SetEvent(handle)
+
+    # Test 1, handle is SET after 1 sec in separate coroutine
+
+    handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    t0 = _core.current_time()
+
+    async with _core.open_nursery() as nursery:
+        nursery.start_soon(WaitForSingleObject, handle)
+        nursery.start_soon(handle_setter, handle)
+
+    kernel32.CloseHandle(handle)
+    t1 = _core.current_time()
+    assert TIMEOUT <= (t1 - t0) < 1.1 * TIMEOUT
+    print('test_WaitForSingleObject test 1 OK')
+
+    # Test 2, handle is CLOSED after 1 sec - NOPE, wont work unless we use zero timeout
+
+    pass
+
+    # Test 3, cancelation
+
+    handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    t0 = _core.current_time()
+
+    with _timeouts.move_on_after(TIMEOUT):
+        await WaitForSingleObject(handle)
+
+    kernel32.CloseHandle(handle)
+    t1 = _core.current_time()
+    assert TIMEOUT <= (t1 - t0) < 1.1 * TIMEOUT
+    print('test_WaitForSingleObject test 3 OK')
+
+    # Test 4, already cancelled
+
+    handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    kernel32.SetEvent(handle)
+    t0 = _core.current_time()
+
+    with _timeouts.move_on_after(TIMEOUT):
+        await WaitForSingleObject(handle)
+
+    kernel32.CloseHandle(handle)
+    t1 = _core.current_time()
+    assert (t1 - t0) < 0.5 * TIMEOUT
+    print('test_WaitForSingleObject test 4 OK')
+
+    # Test 5, already closed
+
+    handle = kernel32.CreateEventA(ffi.NULL, True, False, ffi.NULL)
+    kernel32.CloseHandle(handle)
+    t0 = _core.current_time()
+
+    with _timeouts.move_on_after(TIMEOUT):
+        await WaitForSingleObject(handle)
+
+    t1 = _core.current_time()
+    assert (t1 - t0) < 0.5 * TIMEOUT
+    print('test_WaitForSingleObject test 5 OK')
+
+
 # XX test setting the iomanager._iocp to something weird to make sure that the
 # IOCP thread can send exceptions back to the main thread