Implement cancelable WaitForSingleObject for Windows

docs/source/reference-hazmat.rst

@@ -253,6 +253,19 @@ anything real. See `#26
 Windows-specific API
 --------------------
 
+.. function:: WaitForSingleObject(handle)
+    :async:
+
+    Async and cancellable variant of `WaitForSingleObject
+    <https://msdn.microsoft.com/en-us/library/windows/desktop/ms687032(v=vs.85).aspx>`__.
+    Windows only.
+
+    :arg handle:
+        A Win32 object handle, as a Python integer.
+    :raises OSError:
+        If the handle is invalid, e.g. when it is already closed.
+
+
 TODO: these are implemented, but are currently more of a sketch than
 anything real. See `#26
 <https://github.com/python-trio/trio/issues/26>`__ and `#52

newsfragments/233.feature.rst

@@ -0,0 +1 @@
+Add :func:`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,70 @@
+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 WaitForSingleObject. Windows only.
+
+    Args:
+      handle: A Win32 handle, as a Python integer.
+
+    Raises:
+      OSError: If the handle is invalid, e.g. when it is already closed.
+
+    """
+    # 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,7 +1,15 @@
-# 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
-# system.
+"""
+This namespace represents low-level functionality not intended for daily use,
+but useful for extending Trio's functionality.
+"""
+
+import sys
+
+# This is the union of a subset of trio/_core/ and some things from trio/*.py.
+# See comments in trio/__init__.py for details. To make static analysis easier,
+# this lists all possible symbols from trio._core, and then we prune those that
+# aren't available on this system. After that we add some symbols from trio/*.py.
+
 __all__ = [
     "cancel_shielded_checkpoint",
     "Abort",
@@ -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"]