asyncio: Add support for running subprocesses on Windows with the IOCP event loop (Richard Oudkerk).

7 files changed, 137 insertions, 195 deletions

diff --git a/Lib/asyncio/__init__.py b/Lib/asyncio/__init__.py
index afc444d..0d288d5 100644
--- a/Lib/asyncio/__init__.py
+++ b/Lib/asyncio/__init__.py
@@ -4,10 +4,18 @@ import sys
 
 # The selectors module is in the stdlib in Python 3.4 but not in 3.3.
 # Do this first, so the other submodules can use "from . import selectors".
+# Prefer asyncio/selectors.py over the stdlib one, as ours may be newer.
 try:
-    import selectors  # Will also be exported.
-except ImportError:
     from . import selectors
+except ImportError:
+    import selectors  # Will also be exported.
+
+if sys.platform == 'win32':
+    # Similar thing for _overlapped.
+    try:
+        from . import _overlapped
+    except ImportError:
+        import _overlapped  # Will also be exported.
 
 # This relies on each of the submodules having an __all__ variable.
 from .futures import *
diff --git a/Lib/asyncio/proactor_events.py b/Lib/asyncio/proactor_events.py
index cb8625d..ce226b9 100644
--- a/Lib/asyncio/proactor_events.py
+++ b/Lib/asyncio/proactor_events.py
@@ -267,8 +267,15 @@ class BaseProactorEventLoop(base_events.BaseEventLoop):
         return _ProactorReadPipeTransport(self, sock, protocol, waiter, extra)
 
     def _make_write_pipe_transport(self, sock, protocol, waiter=None,
-                                   extra=None):
-        return _ProactorWritePipeTransport(self, sock, protocol, waiter, extra)
+                                   extra=None, check_for_hangup=True):
+        if check_for_hangup:
+            # We want connection_lost() to be called when other end closes
+            return _ProactorDuplexPipeTransport(self,
+                                                sock, protocol, waiter, extra)
+        else:
+            # If other end closes we may not notice for a long time
+            return _ProactorWritePipeTransport(self, sock, protocol, waiter,
+                                               extra)
 
     def close(self):
         if self._proactor is not None:
diff --git a/Lib/asyncio/unix_events.py b/Lib/asyncio/unix_events.py
index 3807680..c95ad48 100644
--- a/Lib/asyncio/unix_events.py
+++ b/Lib/asyncio/unix_events.py
@@ -1,6 +1,5 @@
 """Selector eventloop for Unix with signal handling."""
 
-import collections
 import errno
 import fcntl
 import os
@@ -11,6 +10,7 @@ import subprocess
 import sys
 
 
+from . import base_subprocess
 from . import constants
 from . import events
 from . import protocols
@@ -406,159 +406,20 @@ class _UnixWritePipeTransport(transports.WriteTransport):
             self._loop = None
 
 
-class _UnixWriteSubprocessPipeProto(protocols.BaseProtocol):
-    pipe = None
+class _UnixSubprocessTransport(base_subprocess.BaseSubprocessTransport):
 
-    def __init__(self, proc, fd):
-        self.proc = proc
-        self.fd = fd
-        self.connected = False
-        self.disconnected = False
-        proc._pipes[fd] = self
-
-    def connection_made(self, transport):
-        self.connected = True
-        self.pipe = transport
-        self.proc._try_connected()
-
-    def connection_lost(self, exc):
-        self.disconnected = True
-        self.proc._pipe_connection_lost(self.fd, exc)
-
-
-class _UnixReadSubprocessPipeProto(_UnixWriteSubprocessPipeProto,
-                                   protocols.Protocol):
-
-    def data_received(self, data):
-        self.proc._pipe_data_received(self.fd, data)
-
-    def eof_received(self):
-        pass
-
-
-class _UnixSubprocessTransport(transports.SubprocessTransport):
-
-    def __init__(self, loop, protocol, args, shell,
-                 stdin, stdout, stderr, bufsize,
-                 extra=None, **kwargs):
-        super().__init__(extra)
-        self._protocol = protocol
-        self._loop = loop
-
-        self._pipes = {}
+    def _start(self, args, shell, stdin, stdout, stderr, bufsize, **kwargs):
         stdin_w = None
         if stdin == subprocess.PIPE:
-            self._pipes[STDIN] = None
             # Use a socket pair for stdin, since not all platforms
             # support selecting read events on the write end of a
             # socket (which we use in order to detect closing of the
             # other end).  Notably this is needed on AIX, and works
             # just fine on other platforms.
             stdin, stdin_w = self._loop._socketpair()
-        if stdout == subprocess.PIPE:
-            self._pipes[STDOUT] = None
-        if stderr == subprocess.PIPE:
-            self._pipes[STDERR] = None
-        self._pending_calls = collections.deque()
-        self._finished = False
-        self._returncode = None
-
         self._proc = subprocess.Popen(
             args, shell=shell, stdin=stdin, stdout=stdout, stderr=stderr,
             universal_newlines=False, bufsize=bufsize, **kwargs)
         if stdin_w is not None:
             stdin.close()
             self._proc.stdin = open(stdin_w.detach(), 'rb', buffering=bufsize)
-        self._extra['subprocess'] = self._proc
-
-    def close(self):
-        for proto in self._pipes.values():
-            proto.pipe.close()
-        if self._returncode is None:
-            self.terminate()
-
-    def get_pid(self):
-        return self._proc.pid
-
-    def get_returncode(self):
-        return self._returncode
-
-    def get_pipe_transport(self, fd):
-        if fd in self._pipes:
-            return self._pipes[fd].pipe
-        else:
-            return None
-
-    def send_signal(self, signal):
-        self._proc.send_signal(signal)
-
-    def terminate(self):
-        self._proc.terminate()
-
-    def kill(self):
-        self._proc.kill()
-
-    @tasks.coroutine
-    def _post_init(self):
-        proc = self._proc
-        loop = self._loop
-        if proc.stdin is not None:
-            transp, proto = yield from loop.connect_write_pipe(
-                lambda: _UnixWriteSubprocessPipeProto(self, STDIN),
-                proc.stdin)
-        if proc.stdout is not None:
-            transp, proto = yield from loop.connect_read_pipe(
-                lambda: _UnixReadSubprocessPipeProto(self, STDOUT),
-                proc.stdout)
-        if proc.stderr is not None:
-            transp, proto = yield from loop.connect_read_pipe(
-                lambda: _UnixReadSubprocessPipeProto(self, STDERR),
-                proc.stderr)
-        if not self._pipes:
-            self._try_connected()
-
-    def _call(self, cb, *data):
-        if self._pending_calls is not None:
-            self._pending_calls.append((cb, data))
-        else:
-            self._loop.call_soon(cb, *data)
-
-    def _try_connected(self):
-        assert self._pending_calls is not None
-        if all(p is not None and p.connected for p in self._pipes.values()):
-            self._loop.call_soon(self._protocol.connection_made, self)
-            for callback, data in self._pending_calls:
-                self._loop.call_soon(callback, *data)
-            self._pending_calls = None
-
-    def _pipe_connection_lost(self, fd, exc):
-        self._call(self._protocol.pipe_connection_lost, fd, exc)
-        self._try_finish()
-
-    def _pipe_data_received(self, fd, data):
-        self._call(self._protocol.pipe_data_received, fd, data)
-
-    def _process_exited(self, returncode):
-        assert returncode is not None, returncode
-        assert self._returncode is None, self._returncode
-        self._returncode = returncode
-        self._loop._subprocess_closed(self)
-        self._call(self._protocol.process_exited)
-        self._try_finish()
-
-    def _try_finish(self):
-        assert not self._finished
-        if self._returncode is None:
-            return
-        if all(p is not None and p.disconnected
-               for p in self._pipes.values()):
-            self._finished = True
-            self._loop.call_soon(self._call_connection_lost, None)
-
-    def _call_connection_lost(self, exc):
-        try:
-            self._protocol.connection_lost(exc)
-        finally:
-            self._proc = None
-            self._protocol = None
-            self._loop = None
diff --git a/Lib/asyncio/windows_events.py b/Lib/asyncio/windows_events.py
index 1ffac99..b70b353 100644
--- a/Lib/asyncio/windows_events.py
+++ b/Lib/asyncio/windows_events.py
@@ -2,21 +2,19 @@
 
 import errno
 import socket
+import subprocess
 import weakref
 import struct
 import _winapi
 
+from . import base_subprocess
 from . import futures
 from . import proactor_events
 from . import selector_events
 from . import tasks
 from . import windows_utils
 from .log import logger
-
-try:
-    import _overlapped
-except ImportError:
-    from . import _overlapped
+from . import _overlapped
 
 
 __all__ = ['SelectorEventLoop', 'ProactorEventLoop', 'IocpProactor']
@@ -168,6 +166,19 @@ class ProactorEventLoop(proactor_events.BaseProactorEventLoop):
     def _stop_serving(self, server):
         server.close()
 
+    @tasks.coroutine
+    def _make_subprocess_transport(self, protocol, args, shell,
+                                   stdin, stdout, stderr, bufsize,
+                                   extra=None, **kwargs):
+        transp = _WindowsSubprocessTransport(self, protocol, args, shell,
+                                             stdin, stdout, stderr, bufsize,
+                                             extra=None, **kwargs)
+        yield from transp._post_init()
+        return transp
+
+    def _subprocess_closed(self, transport):
+        pass
+
 
 class IocpProactor:
     """Proactor implementation using IOCP."""
@@ -413,3 +424,16 @@ class IocpProactor:
         if self._iocp is not None:
             _winapi.CloseHandle(self._iocp)
             self._iocp = None
+
+
+class _WindowsSubprocessTransport(base_subprocess.BaseSubprocessTransport):
+
+    def _start(self, args, shell, stdin, stdout, stderr, bufsize, **kwargs):
+        self._proc = windows_utils.Popen(
+            args, shell=shell, stdin=stdin, stdout=stdout, stderr=stderr,
+            bufsize=bufsize, **kwargs)
+        def callback(f):
+            returncode = self._proc.poll()
+            self._process_exited(returncode)
+        f = self._loop._proactor.wait_for_handle(int(self._proc._handle))
+        f.add_done_callback(callback)
diff --git a/Lib/asyncio/windows_utils.py b/Lib/asyncio/windows_utils.py
index 04b43e9..2fc3f7a 100644
--- a/Lib/asyncio/windows_utils.py
+++ b/Lib/asyncio/windows_utils.py
@@ -24,6 +24,7 @@ __all__ = ['socketpair', 'pipe', 'Popen', 'PIPE', 'PipeHandle']
 
 BUFSIZE = 8192
 PIPE = subprocess.PIPE
+STDOUT = subprocess.STDOUT
 _mmap_counter = itertools.count()
 
 #
@@ -146,24 +147,34 @@ class Popen(subprocess.Popen):
     The stdin, stdout, stderr are None or instances of PipeHandle.
     """
     def __init__(self, args, stdin=None, stdout=None, stderr=None, **kwds):
+        assert not kwds.get('universal_newlines')
+        assert kwds.get('bufsize', 0) == 0
         stdin_rfd = stdout_wfd = stderr_wfd = None
         stdin_wh = stdout_rh = stderr_rh = None
         if stdin == PIPE:
-            stdin_rh, stdin_wh = pipe(overlapped=(False, True))
+            stdin_rh, stdin_wh = pipe(overlapped=(False, True), duplex=True)
             stdin_rfd = msvcrt.open_osfhandle(stdin_rh, os.O_RDONLY)
+        else:
+            stdin_rfd = stdin
         if stdout == PIPE:
             stdout_rh, stdout_wh = pipe(overlapped=(True, False))
             stdout_wfd = msvcrt.open_osfhandle(stdout_wh, 0)
+        else:
+            stdout_wfd = stdout
         if stderr == PIPE:
             stderr_rh, stderr_wh = pipe(overlapped=(True, False))
             stderr_wfd = msvcrt.open_osfhandle(stderr_wh, 0)
+        elif stderr == STDOUT:
+            stderr_wfd = stdout_wfd
+        else:
+            stderr_wfd = stderr
         try:
-            super().__init__(args, bufsize=0, universal_newlines=False,
-                             stdin=stdin_rfd, stdout=stdout_wfd,
+            super().__init__(args, stdin=stdin_rfd, stdout=stdout_wfd,
                              stderr=stderr_wfd, **kwds)
         except:
             for h in (stdin_wh, stdout_rh, stderr_rh):
-                _winapi.CloseHandle(h)
+                if h is not None:
+                    _winapi.CloseHandle(h)
             raise
         else:
             if stdin_wh is not None:
diff --git a/Lib/test/test_asyncio/test_events.py b/Lib/test/test_asyncio/test_events.py
index 98896e8..fd2af2e 100644
--- a/Lib/test/test_asyncio/test_events.py
+++ b/Lib/test/test_asyncio/test_events.py
@@ -955,8 +955,23 @@ class EventLoopTestsMixin:
         r.close()
         w.close()
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
+
+class SubprocessTestsMixin:
+
+    def check_terminated(self, returncode):
+        if sys.platform == 'win32':
+            self.assertIsInstance(returncode, int)
+            self.assertNotEqual(0, returncode)
+        else:
+            self.assertEqual(-signal.SIGTERM, returncode)
+
+    def check_killed(self, returncode):
+        if sys.platform == 'win32':
+            self.assertIsInstance(returncode, int)
+            self.assertNotEqual(0, returncode)
+        else:
+            self.assertEqual(-signal.SIGKILL, returncode)
+
     def test_subprocess_exec(self):
         proto = None
         transp = None
@@ -980,11 +995,9 @@ class EventLoopTestsMixin:
         self.loop.run_until_complete(proto.got_data[1].wait())
         transp.close()
         self.loop.run_until_complete(proto.completed)
-        self.assertEqual(-signal.SIGTERM, proto.returncode)
+        self.check_terminated(proto.returncode)
         self.assertEqual(b'Python The Winner', proto.data[1])
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_interactive(self):
         proto = None
         transp = None
@@ -1017,10 +1030,8 @@ class EventLoopTestsMixin:
             transp.close()
 
         self.loop.run_until_complete(proto.completed)
-        self.assertEqual(-signal.SIGTERM, proto.returncode)
+        self.check_terminated(proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_shell(self):
         proto = None
         transp = None
@@ -1030,7 +1041,7 @@ class EventLoopTestsMixin:
             nonlocal proto, transp
             transp, proto = yield from self.loop.subprocess_shell(
                 functools.partial(MySubprocessProtocol, self.loop),
-                'echo "Python"')
+                'echo Python')
             self.assertIsInstance(proto, MySubprocessProtocol)
 
         self.loop.run_until_complete(connect())
@@ -1040,10 +1051,9 @@ class EventLoopTestsMixin:
         self.loop.run_until_complete(proto.completed)
         self.assertEqual(0, proto.returncode)
         self.assertTrue(all(f.done() for f in proto.disconnects.values()))
-        self.assertEqual({1: b'Python\n', 2: b''}, proto.data)
+        self.assertEqual(proto.data[1].rstrip(b'\r\n'), b'Python')
+        self.assertEqual(proto.data[2], b'')
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_exitcode(self):
         proto = None
 
@@ -1059,8 +1069,6 @@ class EventLoopTestsMixin:
         self.loop.run_until_complete(proto.completed)
         self.assertEqual(7, proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_close_after_finish(self):
         proto = None
         transp = None
@@ -1081,8 +1089,6 @@ class EventLoopTestsMixin:
         self.assertEqual(7, proto.returncode)
         self.assertIsNone(transp.close())
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_kill(self):
         proto = None
         transp = None
@@ -1102,10 +1108,30 @@ class EventLoopTestsMixin:
 
         transp.kill()
         self.loop.run_until_complete(proto.completed)
-        self.assertEqual(-signal.SIGKILL, proto.returncode)
+        self.check_killed(proto.returncode)
+
+    def test_subprocess_terminate(self):
+        proto = None
+        transp = None
+
+        prog = os.path.join(os.path.dirname(__file__), 'echo.py')
+
+        @tasks.coroutine
+        def connect():
+            nonlocal proto, transp
+            transp, proto = yield from self.loop.subprocess_exec(
+                functools.partial(MySubprocessProtocol, self.loop),
+                sys.executable, prog)
+            self.assertIsInstance(proto, MySubprocessProtocol)
+
+        self.loop.run_until_complete(connect())
+        self.loop.run_until_complete(proto.connected)
+
+        transp.terminate()
+        self.loop.run_until_complete(proto.completed)
+        self.check_terminated(proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
+    @unittest.skipIf(sys.platform == 'win32', "Don't have SIGHUP")
     def test_subprocess_send_signal(self):
         proto = None
         transp = None
@@ -1127,8 +1153,6 @@ class EventLoopTestsMixin:
         self.loop.run_until_complete(proto.completed)
         self.assertEqual(-signal.SIGHUP, proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_stderr(self):
         proto = None
         transp = None
@@ -1156,8 +1180,6 @@ class EventLoopTestsMixin:
         self.assertTrue(proto.data[2].startswith(b'ERR:test'), proto.data[2])
         self.assertEqual(0, proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_stderr_redirect_to_stdout(self):
         proto = None
         transp = None
@@ -1188,8 +1210,6 @@ class EventLoopTestsMixin:
         transp.close()
         self.assertEqual(0, proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_close_client_stream(self):
         proto = None
         transp = None
@@ -1217,14 +1237,18 @@ class EventLoopTestsMixin:
         self.loop.run_until_complete(proto.disconnects[1])
         stdin.write(b'xxx')
         self.loop.run_until_complete(proto.got_data[2].wait())
-        self.assertEqual(b'ERR:BrokenPipeError', proto.data[2])
-
+        if sys.platform != 'win32':
+            self.assertEqual(b'ERR:BrokenPipeError', proto.data[2])
+        else:
+            # After closing the read-end of a pipe, writing to the
+            # write-end using os.write() fails with errno==EINVAL and
+            # GetLastError()==ERROR_INVALID_NAME on Windows!?!  (Using
+            # WriteFile() we get ERROR_BROKEN_PIPE as expected.)
+            self.assertEqual(b'ERR:OSError', proto.data[2])
         transp.close()
         self.loop.run_until_complete(proto.completed)
-        self.assertEqual(-signal.SIGTERM, proto.returncode)
+        self.check_terminated(proto.returncode)
 
-    @unittest.skipIf(sys.platform == 'win32',
-                     "Don't support subprocess for Windows yet")
     def test_subprocess_wait_no_same_group(self):
         proto = None
         transp = None
@@ -1252,7 +1276,10 @@ if sys.platform == 'win32':
         def create_event_loop(self):
             return windows_events.SelectorEventLoop()
 
-    class ProactorEventLoopTests(EventLoopTestsMixin, unittest.TestCase):
+
+    class ProactorEventLoopTests(EventLoopTestsMixin,
+                                 SubprocessTestsMixin,
+                                 unittest.TestCase):
 
         def create_event_loop(self):
             return windows_events.ProactorEventLoop()
@@ -1283,26 +1310,34 @@ else:
     from asyncio import unix_events
 
     if hasattr(selectors, 'KqueueSelector'):
-        class KqueueEventLoopTests(EventLoopTestsMixin, unittest.TestCase):
+        class KqueueEventLoopTests(EventLoopTestsMixin,
+                                   SubprocessTestsMixin,
+                                   unittest.TestCase):
 
             def create_event_loop(self):
                 return unix_events.SelectorEventLoop(
                     selectors.KqueueSelector())
 
     if hasattr(selectors, 'EpollSelector'):
-        class EPollEventLoopTests(EventLoopTestsMixin, unittest.TestCase):
+        class EPollEventLoopTests(EventLoopTestsMixin,
+                                  SubprocessTestsMixin,
+                                  unittest.TestCase):
 
             def create_event_loop(self):
                 return unix_events.SelectorEventLoop(selectors.EpollSelector())
 
     if hasattr(selectors, 'PollSelector'):
-        class PollEventLoopTests(EventLoopTestsMixin, unittest.TestCase):
+        class PollEventLoopTests(EventLoopTestsMixin,
+                                 SubprocessTestsMixin,
+                                 unittest.TestCase):
 
             def create_event_loop(self):
                 return unix_events.SelectorEventLoop(selectors.PollSelector())
 
     # Should always exist.
-    class SelectEventLoopTests(EventLoopTestsMixin, unittest.TestCase):
+    class SelectEventLoopTests(EventLoopTestsMixin,
+                               SubprocessTestsMixin,
+                               unittest.TestCase):
 
         def create_event_loop(self):
             return unix_events.SelectorEventLoop(selectors.SelectSelector())
diff --git a/Lib/test/test_asyncio/test_windows_utils.py b/Lib/test/test_asyncio/test_windows_utils.py
index f721d31..e013fbd 100644
--- a/Lib/test/test_asyncio/test_windows_utils.py
+++ b/Lib/test/test_asyncio/test_windows_utils.py
@@ -11,11 +11,7 @@ if sys.platform != 'win32':
 import _winapi
 
 from asyncio import windows_utils
-
-try:
-    import _overlapped
-except ImportError:
-    from asyncio import _overlapped
+from asyncio import _overlapped
 
 
 class WinsocketpairTests(unittest.TestCase):