bpo-29883: Add UDP to Windows Proactor Event Loop

Author: ameily

Doc/library/asyncio-eventloops.rst

@@ -109,7 +109,6 @@ Common limits of Windows event loops:
 
 :class:`ProactorEventLoop` specific limits:
 
-- :meth:`~AbstractEventLoop.create_datagram_endpoint` (UDP) is not supported
 - :meth:`~AbstractEventLoop.add_reader` and :meth:`~AbstractEventLoop.add_writer` are
   not supported
 

Lib/asyncio/proactor_events.py

@@ -8,6 +8,7 @@
 
 import socket
 import warnings
+import collections
 
 from . import base_events
 from . import constants
@@ -327,6 +328,102 @@ def _pipe_closed(self, fut):
             self.close()
 
 
+class _ProactorDatagramTransport(_ProactorBasePipeTransport):
+
+    def __init__(self, loop, sock, protocol, address=None,
+                 waiter=None, extra=None):
+        super(_ProactorDatagramTransport, self).__init__(loop, sock, protocol,
+                                                         waiter=waiter,
+                                                         extra=extra)
+        self._address = address
+        # We don't need to call _protocol.connection_made() since our base
+        # constructor does it for us.
+        self._buffer = collections.deque()
+        self._loop.call_soon(self._loop_reading)
+
+    def abort(self):
+        self._force_close(None)
+
+    def sendto(self, data, addr=None):
+        if not isinstance(data, (bytes, bytearray, memoryview)):
+            raise TypeError('data argument must be byte-ish (%r)',
+                            type(data))
+
+        if not data:
+            return
+
+        if self._conn_lost and self._address:
+            # close() or force_close() has been called on the bound endpoint
+            return
+
+        self._buffer.appendleft((data, addr))
+
+        if self._write_fut is None:
+            # No current write operations are active, kick one off
+            self._loop_writing()
+        else:
+            # A write operation is already kicked off
+            pass
+
+    def _loop_writing(self, fut=None):
+        if self._conn_lost:
+            return
+
+        assert fut is self._write_fut
+        if fut:
+            # We are in a _loop_writing() done callback, get the result
+            fut.result()
+
+        if not self._buffer or (self._conn_lost and self._address):
+            # The connection has been closed
+            self._write_fut = None
+            return
+
+        data, addr = self._buffer.pop()
+
+        self._write_fut = None
+        try:
+            if self._address:
+                self._write_fut = self._loop._proactor.send(self._sock, data)
+            else:
+                self._write_fut = self._loop._proactor.sendto(self._sock, data, addr=addr)
+        except OSError as exc:
+            self._protocol.error_received(exc)
+            self._fatal_error(exc, 'Fatal error sending UDP datagram')
+        else:
+            self._write_fut.add_done_callback(self._loop_writing)
+
+    def _loop_reading(self, fut=None):
+        if self._conn_lost:
+            return
+
+        assert self._read_fut is fut
+
+        if fut:
+            res = fut.result()
+
+            if self._address:
+                data, addr = res, self._address
+            else:
+                data, addr = res
+
+            self._protocol.datagram_received(data, addr)
+
+        if self._conn_lost:
+            return
+
+        try:
+            if self._address:
+                self._read_fut = self._loop._proactor.recv(self._sock, 4096)
+            else:
+                self._read_fut = self._loop._proactor.recvfrom(self._sock, 4096)
+        except OSError as exc:
+            self._protocol.error_received(exc)
+            self._fatal_error(exc, "Fatal error reading from UDP endpoint")
+        else:
+            self._read_fut.add_done_callback(self._loop_reading)
+
+
 class _ProactorDuplexPipeTransport(_ProactorReadPipeTransport,
                                    _ProactorBaseWritePipeTransport,
                                    transports.Transport):
@@ -390,6 +487,11 @@ def _make_socket_transport(self, sock, protocol, waiter=None,
         return _ProactorSocketTransport(self, sock, protocol, waiter,
                                         extra, server)
 
+    def _make_datagram_transport(self, sock, protocol,
+                                 address=None, waiter=None, extra=None):
+        return _ProactorDatagramTransport(self, sock, protocol, address,
+                                          waiter, extra)
+
     def _make_ssl_transport(self, rawsock, protocol, sslcontext, waiter=None,
                             *, server_side=False, server_hostname=None,
                             extra=None, server=None):

Lib/asyncio/windows_events.py

@@ -443,6 +443,14 @@ def recv_into(self, conn, buf, flags=0):
         except BrokenPipeError:
             return self._result(b'')
 
+    def recvfrom(self, conn, nbytes, flags=0):
+        self._register_with_iocp(conn)
+        ov = _overlapped.Overlapped(NULL)
+        try:
+            ov.WSARecvFrom(conn.fileno(), nbytes, flags)
+        except BrokenPipeError:
+            return self._result((b'', (None, None)))
+
         def finish_recv(trans, key, ov):
             try:
                 return ov.getresult()
@@ -454,6 +462,23 @@ def finish_recv(trans, key, ov):
 
         return self._register(ov, conn, finish_recv)
 
+    def sendto(self, conn, buf, flags=0, addr=None):
+        self._register_with_iocp(conn)
+        ov = _overlapped.Overlapped(NULL)
+
+        ov.WSASendTo(conn.fileno(), buf, flags, addr)
+
+        def finish_send(trans, key, ov):
+            try:
+                return ov.getresult()
+            except OSError as exc:
+                if exc.winerror == _overlapped.ERROR_NETNAME_DELETED:
+                    raise ConnectionResetError(*exc.args)
+                else:
+                    raise
+
+        return self._register(ov, conn, finish_send)
+
     def send(self, conn, buf, flags=0):
         self._register_with_iocp(conn)
         ov = _overlapped.Overlapped(NULL)
@@ -502,6 +527,15 @@ async def accept_coro(future, conn):
         return future
 
     def connect(self, conn, address):
+        if conn.type == socket.SOCK_DGRAM:
+            # WSAConnect will complete immediately for UDP sockets so we don't
+            # need to register any IOCP operation
+            _overlapped.WSAConnect(conn.fileno(), address)
+            fut = self._loop.create_future()
+
+            fut.set_result(None)
+            return fut
+
         self._register_with_iocp(conn)
         # The socket needs to be locally bound before we call ConnectEx().
         try:

Lib/test/test_asyncio/test_proactor_events.py

@@ -9,6 +9,7 @@
 from asyncio.proactor_events import _ProactorSocketTransport
 from asyncio.proactor_events import _ProactorWritePipeTransport
 from asyncio.proactor_events import _ProactorDuplexPipeTransport
+from asyncio.proactor_events import _ProactorDatagramTransport
 from test.test_asyncio import utils as test_utils
 
 
@@ -562,6 +563,81 @@ def test_stop_serving(self):
         self.assertTrue(sock.close.called)
         self.proactor._stop_serving.assert_called_with(sock)
 
+    def datagram_transport(self):
+        self.protocol = make_test_protocol(asyncio.DatagramProtocol)
+        return self.loop._make_datagram_transport(self.sock, self.protocol)
+
+    def test_make_datagram_transport(self):
+        tr = self.datagram_transport()
+        self.assertIsInstance(tr, _ProactorDatagramTransport)
+        close_transport(tr)
+
+    def test_datagram_loop_writing(self):
+        tr = self.datagram_transport()
+        tr._buffer.appendleft((b'data', ('127.0.0.1', 12068)))
+        tr._loop_writing()
+        self.loop._proactor.sendto.assert_called_with(self.sock, b'data', addr=('127.0.0.1', 12068))
+        self.loop._proactor.sendto.return_value.add_done_callback.\
+            assert_called_with(tr._loop_writing)
+
+        close_transport(tr)
+
+    def test_datagram_loop_reading(self):
+        tr = self.datagram_transport()
+        tr._loop_reading()
+        self.loop._proactor.recvfrom.assert_called_with(self.sock, 4096)
+        self.assertFalse(self.protocol.datagram_received.called)
+        self.assertFalse(self.protocol.error_received.called)
+
+    def test_datagram_loop_reading_data(self):
+        res = asyncio.Future(loop=self.loop)
+        res.set_result((b'data', ('127.0.0.1', 12068)))
+
+        tr = self.datagram_transport()
+        tr._read_fut = res
+        tr._loop_reading(res)
+        self.loop._proactor.recvfrom.assert_called_with(self.sock, 4096)
+        self.protocol.datagram_received.assert_called_with(b'data', ('127.0.0.1', 12068))
+
+    def test_datagram_loop_reading_no_data(self):
+        res = asyncio.Future(loop=self.loop)
+        res.set_result((b'', ('127.0.0.1', 12068)))
+
+        tr = self.datagram_transport()
+        self.assertRaises(AssertionError, tr._loop_reading, res)
+
+        tr.close = mock.Mock()
+        tr._read_fut = res
+        tr._loop_reading(res)
+        self.assertFalse(self.loop._proactor.recvfrom.called)
+        self.assertFalse(self.protocol.error_received.called)
+        self.assertFalse(tr.close.called)
+
+    def test_datagram_loop_reading_aborted(self):
+        err = self.loop._proactor.recvfrom.side_effect = ConnectionAbortedError()
+
+        tr = self.datagram_transport()
+        tr._fatal_error = mock.Mock()
+        tr._protocol.error_received = mock.Mock()
+        tr._loop_reading()
+        tr._protocol.error_received.assert_called_with(err)
+        tr._fatal_error.assert_called_with(
+                            err,
+                            'Fatal error reading from UDP endpoint')
+
+    def test_datagram_loop_writing_aborted(self):
+        err = self.loop._proactor.sendto.side_effect = ConnectionAbortedError()
+
+        tr = self.datagram_transport()
+        tr._fatal_error = mock.Mock()
+        tr._protocol.error_received = mock.Mock()
+        tr._protocol.error_received.assert_called_with(err)
+        tr._buffer.appendleft((b'Hello', ('127.0.0.1', 12068)))
+        tr._loop_writing()
+        tr._fatal_error.assert_called_with(
+                            err,
+                            'Fatal error sending UDP datagram')
+
 
 if __name__ == '__main__':
     unittest.main()