Implement cancelable WaitForSingleObject for Windows

docs/source/reference-hazmat.rst

@@ -268,6 +268,8 @@ anything real. See `#26
 .. function:: monitor_completion_key()
    :with: queue
 
+.. function:: WaitForSingleObject()
+    :async:
 
 Unbounded queues
 ================

newsfragments/233.feature.rst

@@ -0,0 +1 @@
+Add ``trio.hazmat.WaitForSingleObject()`` async function to await Windows handles.

trio/__init__.py

@@ -1,3 +1,6 @@
+"""Trio - Pythonic async I/O for humans and snake people.
+"""
+
 # General layout:
 #
 # trio/_core/... is the self-contained core library. It does various

trio/_core/__init__.py

@@ -1,3 +1,10 @@
+"""
+This namespace represents the core functionality that has to be built-in
+and deal with private internal data structures. Things in this namespace
+are publicly available in either trio, trio.hazmat, or trio.testing.
+"""
+
+
 # Needs to be defined early so it can be imported:
 def _public(fn):
     # Used to mark methods on _Runner and on IOManager implementations that

trio/_core/_io_windows.py

@@ -22,6 +22,7 @@
     INVALID_HANDLE_VALUE,
     raise_winerror,
     ErrorCodes,
+    _handle,
 )
 
 # There's a lot to be said about the overall design of a Windows event
@@ -96,19 +97,6 @@ def _check(success):
     return success
 
 
-def _handle(obj):
-    # For now, represent handles as either cffi HANDLEs or as ints.  If you
-    # try to pass in a file descriptor instead, it's not going to work
-    # out. (For that msvcrt.get_osfhandle does the trick, but I don't know if
-    # we'll actually need that for anything...) For sockets this doesn't
-    # matter, Python never allocates an fd. So let's wait until we actually
-    # encounter the problem before worrying about it.
-    if type(obj) is int:
-        return ffi.cast("HANDLE", obj)
-    else:
-        return obj
-
-
 @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
@@ -104,6 +134,19 @@
 INVALID_HANDLE_VALUE = ffi.cast("HANDLE", -1)
 
 
+def _handle(obj):
+    # For now, represent handles as either cffi HANDLEs or as ints.  If you
+    # try to pass in a file descriptor instead, it's not going to work
+    # out. (For that msvcrt.get_osfhandle does the trick, but I don't know if
+    # we'll actually need that for anything...) For sockets this doesn't
+    # matter, Python never allocates an fd. So let's wait until we actually
+    # encounter the problem before worrying about it.
+    if type(obj) is int:
+        return ffi.cast("HANDLE", obj)
+    else:
+        return obj
+
+
 def raise_winerror(winerror=None, *, filename=None, filename2=None):
     if winerror is None:
         winerror, msg = ffi.getwinerror()
@@ -116,6 +159,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,4 +1,5 @@
 import os
+
 import pytest
 
 on_windows = (os.name == "nt")

trio/_wait_for_object.py

@@ -0,0 +1,67 @@
+from . import _timeouts
+from . import _core
+from ._threads import run_sync_in_worker_thread
+from ._core._windows_cffi import ffi, kernel32, ErrorCodes, raise_winerror, _handle
+
+__all__ = ["WaitForSingleObject"]
+
+
+class StubLimiter:
+    def release_on_behalf_of(self, x):
+        pass
+
+    async def acquire_on_behalf_of(self, x):
+        pass
+
+
+async def WaitForSingleObject(obj):
+    """Async and cancellable variant of kernel32.WaitForSingleObject(). Windows only.
+
+    Args:
+      handle: A win32 handle in the form of an int or cffi HANDLE object.
+
+    """
+    # Allow ints or whatever we can convert to a win handle
+    handle = _handle(obj)
+
+    # 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_FAILED:
+        raise_winerror()
+    elif retcode != ErrorCodes.WAIT_TIMEOUT:
+        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)
+    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 = 0xffffffff  # INFINITE
+    retcode = kernel32.WaitForMultipleObjects(
+        n, handle_arr, False, timeout
+    )  # blocking
+    if retcode == ErrorCodes.WAIT_FAILED:
+        raise_winerror()

trio/hazmat.py

@@ -1,3 +1,11 @@
+"""
+This namespace represents low-level functionality not intended for daily use,
+but useful for extending Trio's functionality. It is the union of
+a subset of trio/_core/ and some things from trio/*.py.
+"""
+
+import sys
+
 # These are all re-exported from trio._core. See comments in trio/__init__.py
 # for details. To make static analysis easier, this lists all possible
 # symbols, and then we prune some below if they aren't available on this
@@ -56,3 +64,8 @@
         # who knows.
         remove_from_all = __all__.remove
         remove_from_all(_sym)
+
+# Import bits from trio/*.py
+if sys.platform.startswith("win"):
+    from ._wait_for_object import WaitForSingleObject
+    __all__ += ["WaitForSingleObject"]