Rewrite IPC to use multiprocessing.connection

mypy/dmypy.py

@@ -422,7 +422,7 @@ def request(command: str, *, timeout: Optional[int] = None,
     _, name = get_status()
     try:
         with IPCClient(name, timeout) as client:
-                client.write(bdata)
+                client.send_bytes(bdata)
                 response = receive(client)
     except (OSError, IPCException) as err:
         return {'error': str(err)}

mypy/dmypy_server.py

@@ -200,8 +200,8 @@ def serve(self) -> None:
                 json.dump({'pid': os.getpid(), 'connection_name': server.connection_name}, f)
                 f.write('\n')  # I like my JSON with a trailing newline
             while True:
-                with server:
-                    data = receive(server)
+                with server as connection:
+                    data = receive(connection)
                     resp = {}  # type: Dict[str, Any]
                     if 'command' not in data:
                         resp = {'error': "No command found in request"}
@@ -217,10 +217,10 @@ def serve(self) -> None:
                                 # If we are crashing, report the crash to the client
                                 tb = traceback.format_exception(*sys.exc_info())
                                 resp = {'error': "Daemon crashed!\n" + "".join(tb)}
-                                server.write(json.dumps(resp).encode('utf8'))
+                                connection.send_bytes(json.dumps(resp).encode('utf8'))
                                 raise
                     try:
-                        server.write(json.dumps(resp).encode('utf8'))
+                        connection.send_bytes(json.dumps(resp).encode('utf8'))
                     except OSError:
                         pass  # Maybe the client hung up
                     if command == 'stop':
@@ -232,12 +232,9 @@ def serve(self) -> None:
             # simplify the logic and always remove the file, since
             # that could cause us to remove a future server's
             # status file.)
+            server.close()
             if command != 'stop':
                 os.unlink(STATUS_FILE)
-            try:
-                server.cleanup()  # try to remove the socket dir on Linux
-            except OSError:
-                pass
             exc_info = sys.exc_info()
             if exc_info[0] and exc_info[0] is not SystemExit:
                 traceback.print_exception(*exc_info)

mypy/dmypy_util.py

@@ -4,10 +4,14 @@
 """
 
 import json
+import sys
 
 from typing import Any
 
-from mypy.ipc import IPCBase
+if sys.platform == 'win32':
+    from multiprocessing.connection import PipeConnection as Connection
+else:
+    from multiprocessing.connection import Connection
 
 MYPY = False
 if MYPY:
@@ -16,13 +20,16 @@
 STATUS_FILE = '.dmypy.json'  # type: Final
 
 
-def receive(connection: IPCBase) -> Any:
+def receive(connection: Connection) -> Any:
     """Receive JSON data from a connection until EOF.
 
     Raise OSError if the data received is not valid JSON or if it is
     not a dict.
     """
-    bdata = connection.read()
+    try:
+        bdata = connection.recv_bytes()
+    except EOFError:
+        pass
     if not bdata:
         raise OSError("No data received")
     try:

mypy/ipc.py

@@ -2,14 +2,13 @@
 
 On Unix, this uses AF_UNIX sockets.
 On Windows, this uses NamedPipes.
+
+Portions of this file were modified from Python's multiprocess.connection module.
 """
 
 import base64
-import contextlib
 import os
-import shutil
 import sys
-import tempfile
 
 from typing import Iterator, Optional, Callable
 
@@ -19,199 +18,145 @@
 
 from types import TracebackType
 
+from multiprocessing.connection import address_type
+
 if sys.platform == 'win32':
     # This may be private, but it is needed for IPC on Windows, and is basically stable
     import _winapi
-    import ctypes
+    from multiprocessing.connection import PipeListener as Listener
+    from multiprocessing.connection import PipeClient as Client
+    from multiprocessing.connection import PipeConnection as Connection
+else:
+    import socket
+    from multiprocessing.connection import SocketListener as Listener
+    from multiprocessing.connection import SocketClient as Client
+    from multiprocessing.connection import Connection
 
-    _IPCHandle = int
 
-    kernel32 = ctypes.windll.kernel32
-    DisconnectNamedPipe = kernel32.DisconnectNamedPipe  # type: Callable[[_IPCHandle], int]
-    FlushFileBuffers = kernel32.FlushFileBuffers  # type: Callable[[_IPCHandle], int]
+if sys.platform == 'win32':
+    CONNECTION_FAMILY = 'AF_PIPE'
 else:
-    import socket
-    _IPCHandle = socket.socket
+    CONNECTION_FAMILY = 'AF_UNIX'
 
 
 class IPCException(Exception):
     """Exception for IPC issues."""
     pass
 
 
-class IPCBase:
-    """Base class for communication between the dmypy client and server.
-
-    This contains logic shared between the client and server, such as reading
-    and writing.
-    """
-
-    connection = None  # type: _IPCHandle
-
-    def __init__(self, name: str) -> None:
-        self.READ_SIZE = 100000
-        self.name = name
-
-    def read(self) -> bytes:
-        """Read bytes from an IPC connection until its empty."""
-        bdata = bytearray()
-        while True:
-            if sys.platform == 'win32':
-                more, _ = _winapi.ReadFile(self.connection, self.READ_SIZE)
-            else:
-                more = self.connection.recv(self.READ_SIZE)
-            if not more:
-                break
-            bdata.extend(more)
-        return bytes(bdata)
-
-    def write(self, data: bytes) -> None:
-        """Write bytes to an IPC connection."""
-        if sys.platform == 'win32':
-            try:
-                _winapi.WriteFile(self.connection, data)
-                # this empty write is to copy the behavior of socket.sendall,
-                # which also sends an empty message to signify it is done writing
-                _winapi.WriteFile(self.connection, b'')
-            except WindowsError as e:
-                raise IPCException("Failed to write with error: {}".format(e.winerror))
-        else:
-            self.connection.sendall(data)
-            self.connection.shutdown(socket.SHUT_WR)
-
-    def close(self) -> None:
-        if sys.platform == 'win32':
-            if self.connection != _winapi.NULL:
-                _winapi.CloseHandle(self.connection)
-        else:
-            self.connection.close()
-
-
-class IPCClient(IPCBase):
-    """The client side of an IPC connection."""
+class IPCClient:
+    """Client side of an IPC connection based on multiprocessing.connection.Client."""
 
-    def __init__(self, name: str, timeout: Optional[int]) -> None:
-        super().__init__(name)
+    def __init__(self, address: str, timeout: Optional[int] = None) -> None:
         if sys.platform == 'win32':
-            timeout = timeout or 0xFFFFFFFF  # NMPWAIT_WAIT_FOREVER
-            try:
-                _winapi.WaitNamedPipe(self.name, timeout)
-            except FileNotFoundError:
-                raise IPCException("The NamedPipe at {} was not found.".format(self.name))
-            except WindowsError as e:
-                if e.winerror == _winapi.ERROR_SEM_TIMEOUT:
-                    raise IPCException("Timed out waiting for connection.")
-                else:
-                    raise
-            try:
-                self.connection = _winapi.CreateFile(
-                    self.name,
-                    _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
-                    0,
-                    _winapi.NULL,
-                    _winapi.OPEN_EXISTING,
-                    0,
-                    _winapi.NULL,
-                )
-            except WindowsError as e:
-                if e.winerror == _winapi.ERROR_PIPE_BUSY:
-                    raise IPCException("The connection is busy.")
+            timeout = timeout * 1000 if timeout is not None else _winapi.NMPWAIT_WAIT_FOREVER
+            while 1:
+                try:
+                    _winapi.WaitNamedPipe(address, timeout)
+                    h = _winapi.CreateFile(
+                        address, _winapi.GENERIC_READ | _winapi.GENERIC_WRITE,
+                        0, _winapi.NULL, _winapi.OPEN_EXISTING,
+                        _winapi.FILE_FLAG_OVERLAPPED, _winapi.NULL
+                    )
+                except FileNotFoundError:
+                    raise IPCException("The NamedPipe at {} was not found.".format(address))
+                except WindowsError as e:
+                    if e.winerror == _winapi.ERROR_PIPE_BUSY:
+                        continue
+                    elif e.winerror == _winapi.ERROR_SEM_TIMEOUT:
+                        raise IPCException("Timed out waiting for connection.")
+                    else:
+                        raise
                 else:
-                    raise
-            _winapi.SetNamedPipeHandleState(self.connection,
-                                            _winapi.PIPE_READMODE_MESSAGE,
-                                            None,
-                                            None)
+                    break
+
+            _winapi.SetNamedPipeHandleState(
+                h, _winapi.PIPE_READMODE_MESSAGE, None, None
+            )
+            self.connection = Connection(h)
         else:
-            self.connection = socket.socket(socket.AF_UNIX)
-            self.connection.settimeout(timeout)
-            self.connection.connect(name)
+            family = address_type(address)
+            with socket.socket(getattr(socket, family)) as s:
+                if timeout is not None:
+                    s.settimeout(timeout)
+                s.connect(address)
+            self.connection = Connection(s.detach())
 
-    def __enter__(self) -> 'IPCClient':
-        return self
+    def __enter__(self) -> Connection:
+        return self.connection
 
     def __exit__(self,
-                 exc_ty: 'Optional[Type[BaseException]]' = None,
-                 exc_val: Optional[BaseException] = None,
-                 exc_tb: Optional[TracebackType] = None,
+                exc_ty: 'Optional[Type[BaseException]]' = None,
+                exc_val: Optional[BaseException] = None,
+                exc_tb: Optional[TracebackType] = None,
                  ) -> bool:
-        self.close()
+        self.connection.close()
         return False
 
 
-class IPCServer(IPCBase):
-
-    BUFFER_SIZE = 2**16
+class IPCServer(Listener):
+    """Server side of an IPC connection based on multiprocessing.connection.Listener."""
 
-    def __init__(self, name: str, timeout: Optional[int] = None) -> None:
-        if sys.platform == 'win32':
-            name = r'\\.\pipe\{}-{}.pipe'.format(name, base64.b64encode(os.urandom(6)))
-        super().__init__(name)
+    def __init__(self, name: str,
+                 timeout: Optional[int] = None
+                 ) -> None:
         if sys.platform == 'win32':
-            self.connection = _winapi.CreateNamedPipe(self.name,
-                _winapi.PIPE_ACCESS_DUPLEX
-                | _winapi.FILE_FLAG_FIRST_PIPE_INSTANCE,
-                _winapi.PIPE_READMODE_MESSAGE
-                | _winapi.PIPE_TYPE_MESSAGE
-                | _winapi.PIPE_WAIT
-                | 0x8,  # PIPE_REJECT_REMOTE_CLIENTS
-                1,  # one instance
-                self.BUFFER_SIZE,
-                self.BUFFER_SIZE,
-                1000,  # Default timeout in milis
-                0,  # Use default security descriptor
-                                                      )
-            if self.connection == -1:  # INVALID_HANDLE_VALUE
-                err = _winapi.GetLastError()
-                raise IPCException('Invalid handle to pipe: {err}'.format(err))
+            self._name = r'\\.\pipe\{}-{}.pipe'.format(name,
+                                                       base64.b64encode(os.urandom(6)).decode())
+            super().__init__(self._name, backlog=1)
+            self._timeout = timeout * 1000 if timeout else _winapi.INFINITE
         else:
-            self.sock_directory = tempfile.mkdtemp()
-            sockfile = os.path.join(self.sock_directory, self.name)
-            self.sock = socket.socket(socket.AF_UNIX)
-            self.sock.bind(sockfile)
-            self.sock.listen(1)
+            super().__init__(name, CONNECTION_FAMILY, backlog=1)
             if timeout is not None:
-                self.sock.settimeout(timeout)
+                self._sock.settimeout(timeout)
 
-    def __enter__(self) -> 'IPCServer':
+    def accept(self) -> Connection:
         if sys.platform == 'win32':
-            # NOTE: It is theoretically possible that this will hang forever if the
-            # client never connects, though this can be "solved" by killing the server
+            self._handle_queue.append(self._new_handle())
+            handle = self._handle_queue.pop(0)
             try:
-                _winapi.ConnectNamedPipe(self.connection, _winapi.NULL)
+                ov = _winapi.ConnectNamedPipe(handle, overlapped=True)
             except WindowsError as e:
-                if e.winerror == _winapi.ERROR_PIPE_CONNECTED:
-                    pass  # The client already exists, which is fine.
+                if e.winerror != _winapi.ERROR_NO_DATA:
+                    raise
+                # ERROR_NO_DATA can occur if a client has already connected,
+                # written data and then disconnected -- see Python Issue 14725.
+            else:
+                try:
+                    _ = _winapi.WaitForMultipleObjects(
+                        [ov.event], True, self._timeout)
+                    print(self._timeout)
+                except Exception:
+                    ov.cancel()
+                    _winapi.CloseHandle(handle)
+                    raise
+                finally:
+                    _, err = ov.GetOverlappedResult(False)
+                    if err != 0:
+                        raise IPCException('Timed out waiting for client to connect.')
+            self.connection = Connection(handle)
         else:
             try:
-                self.connection, _ = self.sock.accept()
+                s, self._last_accepted = self._socket.accept()
             except socket.timeout:
-                raise IPCException('The socket timed out')
-        return self
+                raise IPCException('Timed out waiting for client to connect.')
+            self.connection = Connection(s.detatch())
+        return self.connection
+
+    @property
+    def connection_name(self) -> str:
+        if sys.platform == 'win32':
+            return self._name
+        else:
+            return self._sock.getsockname()
+
+    def __enter__(self) -> Connection:
+        return self.accept()
 
     def __exit__(self,
                  exc_ty: 'Optional[Type[BaseException]]' = None,
                  exc_val: Optional[BaseException] = None,
                  exc_tb: Optional[TracebackType] = None,
                  ) -> bool:
-        if sys.platform == 'win32':
-            # Wait for the client to finish reading the last write before disconnecting
-            if not FlushFileBuffers(self.connection):
-                raise IPCException("Failed to flush NamedPipe buffer, maybe the client hung up?")
-            DisconnectNamedPipe(self.connection)
-        else:
-            self.close()
+        self.connection.close()
         return False
-
-    def cleanup(self) -> None:
-        if sys.platform == 'win32':
-            self.close()
-        else:
-            shutil.rmtree(self.sock_directory)
-
-    @property
-    def connection_name(self) -> str:
-        if sys.platform == 'win32':
-            return self.name
-        else:
-            return self.sock.getsockname()