diff options
Diffstat (limited to 'Lib')
-rw-r--r-- | Lib/asyncio/windows_events.py | 40 | ||||
-rw-r--r-- | Lib/test/test_asyncio/test_windows_events.py | 40 |
2 files changed, 80 insertions, 0 deletions
diff --git a/Lib/asyncio/windows_events.py b/Lib/asyncio/windows_events.py index bbeada8..1ffac99 100644 --- a/Lib/asyncio/windows_events.py +++ b/Lib/asyncio/windows_events.py @@ -46,6 +46,22 @@ class _OverlappedFuture(futures.Future): return super().cancel() +class _WaitHandleFuture(futures.Future): + """Subclass of Future which represents a wait handle.""" + + def __init__(self, wait_handle, *, loop=None): + super().__init__(loop=loop) + self._wait_handle = wait_handle + + def cancel(self): + super().cancel() + try: + _overlapped.UnregisterWait(self._wait_handle) + except OSError as e: + if e.winerror != _overlapped.ERROR_IO_PENDING: + raise + + class PipeServer(object): """Class representing a pipe server. @@ -271,6 +287,30 @@ class IocpProactor: return windows_utils.PipeHandle(handle) return self._register(ov, None, finish, wait_for_post=True) + def wait_for_handle(self, handle, timeout=None): + if timeout is None: + ms = _winapi.INFINITE + else: + ms = int(timeout * 1000 + 0.5) + + # We only create ov so we can use ov.address as a key for the cache. + ov = _overlapped.Overlapped(NULL) + wh = _overlapped.RegisterWaitWithQueue( + handle, self._iocp, ov.address, ms) + f = _WaitHandleFuture(wh, loop=self._loop) + + def finish(timed_out, _, ov): + if not f.cancelled(): + try: + _overlapped.UnregisterWait(wh) + except OSError as e: + if e.winerror != _overlapped.ERROR_IO_PENDING: + raise + return not timed_out + + self._cache[ov.address] = (f, ov, None, finish) + return f + def _register_with_iocp(self, obj): # To get notifications of finished ops on this objects sent to the # completion port, were must register the handle. diff --git a/Lib/test/test_asyncio/test_windows_events.py b/Lib/test/test_asyncio/test_windows_events.py index 969360c..553ea34 100644 --- a/Lib/test/test_asyncio/test_windows_events.py +++ b/Lib/test/test_asyncio/test_windows_events.py @@ -5,13 +5,17 @@ import unittest if sys.platform != 'win32': raise unittest.SkipTest('Windows only') +import _winapi + import asyncio from asyncio import windows_events +from asyncio import futures from asyncio import protocols from asyncio import streams from asyncio import transports from asyncio import test_utils +from asyncio import _overlapped class UpperProto(protocols.Protocol): @@ -94,6 +98,42 @@ class ProactorTests(unittest.TestCase): return 'done' + def test_wait_for_handle(self): + event = _overlapped.CreateEvent(None, True, False, None) + self.addCleanup(_winapi.CloseHandle, event) + + # Wait for unset event with 0.2s timeout; + # result should be False at timeout + f = self.loop._proactor.wait_for_handle(event, 0.2) + start = self.loop.time() + self.loop.run_until_complete(f) + elapsed = self.loop.time() - start + self.assertFalse(f.result()) + self.assertTrue(0.18 < elapsed < 0.22, elapsed) + + _overlapped.SetEvent(event) + + # Wait for for set event; + # result should be True immediately + f = self.loop._proactor.wait_for_handle(event, 10) + start = self.loop.time() + self.loop.run_until_complete(f) + elapsed = self.loop.time() - start + self.assertTrue(f.result()) + self.assertTrue(0 <= elapsed < 0.02, elapsed) + + _overlapped.ResetEvent(event) + + # Wait for unset event with a cancelled future; + # CancelledError should be raised immediately + f = self.loop._proactor.wait_for_handle(event, 10) + f.cancel() + start = self.loop.time() + with self.assertRaises(futures.CancelledError): + self.loop.run_until_complete(f) + elapsed = self.loop.time() - start + self.assertTrue(0 <= elapsed < 0.02, elapsed) + if __name__ == '__main__': unittest.main() |