fix udp

Author: fantix

tests/test_context.py

@@ -26,6 +26,7 @@ def __init__(self, cvar, *, loop=None):
         self.pipe_ctx = {0, 1, 2}
         self.pipe_connection_lost_fut = asyncio.Future(loop=loop)
         self.process_exited_fut = asyncio.Future(loop=loop)
+        self.error_received_fut = asyncio.Future(loop=loop)
         self.connection_lost_ctx = None
         self.done = asyncio.Future(loop=loop)
 
@@ -77,6 +78,14 @@ def buffer_updated(self, nbytes):
                 )
 
 
+class _DatagramProtocol(_BaseProtocol, asyncio.DatagramProtocol):
+    def datagram_received(self, data, addr):
+        self.data_received_fut.set_result(self.cvar.get())
+
+    def error_received(self, exc):
+        self.error_received_fut.set_result(self.cvar.get())
+
+
 class _SubprocessProtocol(_BaseProtocol, asyncio.SubprocessProtocol):
     def pipe_data_received(self, fd, data):
         self.data_received_fut.set_result(self.cvar.get())
@@ -703,6 +712,45 @@ async def test():
 
         self.loop.run_until_complete(test())
 
+    def test_datagram_protocol(self):
+        cvar = contextvars.ContextVar('cvar', default='outer')
+        proto = _DatagramProtocol(cvar, loop=self.loop)
+        server_addr = ('127.0.0.1', 8888)
+        client_addr = ('127.0.0.1', 0)
+
+        async def run():
+            self.assertEqual(cvar.get(), 'outer')
+            cvar.set('inner')
+
+            def close():
+                cvar.set('closing')
+                proto.transport.close()
+
+            try:
+                await self.loop.create_datagram_endpoint(
+                    lambda: proto, local_addr=server_addr)
+                inner = await proto.connection_made_fut
+                self.assertEqual(inner, "inner")
+
+                s = socket.socket(socket.AF_INET, type=socket.SOCK_DGRAM)
+                s.bind(client_addr)
+                s.sendto(b'data', server_addr)
+                inner = await proto.data_received_fut
+                self.assertEqual(inner, "inner")
+
+                self.loop.call_soon(close)
+                await proto.done
+                if self.implementation != 'asyncio':
+                    # bug in asyncio
+                    self.assertEqual(proto.connection_lost_ctx, "inner")
+            finally:
+                proto.transport.close()
+                s.close()
+                # let transports close
+                await asyncio.sleep(0.1)
+
+        self.loop.run_until_complete(run())
+
 
 class Test_UV_Context(_ContextBaseTests, tb.UVTestCase):
     pass

uvloop/handles/udp.pxd

@@ -19,4 +19,4 @@ cdef class UDPTransport(UVBaseTransport):
     cdef _send(self, object data, object addr)
 
     cdef _on_receive(self, bytes data, object exc, object addr)
-    cdef _on_sent(self, object exc)
+    cdef _on_sent(self, object exc, object context=*)

uvloop/handles/udp.pyx

@@ -56,6 +56,7 @@ cdef class UDPTransport(UVBaseTransport):
         self._family = uv.AF_UNSPEC
         self.__receiving = 0
         self._address = None
+        self.context = Context_CopyCurrent()
 
     cdef _init(self, Loop loop, unsigned int family):
         cdef int err
@@ -252,18 +253,20 @@ cdef class UDPTransport(UVBaseTransport):
                 exc = convert_error(err)
                 self._fatal_error(exc, True)
             else:
-                self._on_sent(None)
+                self._on_sent(None, self.context.copy())
 
     cdef _on_receive(self, bytes data, object exc, object addr):
         if exc is None:
-            self._protocol.datagram_received(data, addr)
+            self.context.run(self._protocol.datagram_received, data, addr)
         else:
-            self._protocol.error_received(exc)
+            self.context.run(self._protocol.error_received, exc)
 
-    cdef _on_sent(self, object exc):
+    cdef _on_sent(self, object exc, object context=None):
         if exc is not None:
             if isinstance(exc, OSError):
-                self._protocol.error_received(exc)
+                if context is None:
+                    context = self.context
+                context.run(self._protocol.error_received, exc)
             else:
                 self._fatal_error(
                     exc, False, 'Fatal write error on datagram transport')

uvloop/sslproto.pyx

@@ -739,7 +739,7 @@ cdef class SSLProtocol:
                         new_MethodHandle(self._loop,
                                          "SSLProtocol._do_read",
                                          <method_t>self._do_read,
-                                         None,
+                                         None,  # current context is good
                                          self))
         except ssl_SSLAgainErrors as exc:
             pass