Restore context on listen in UVStreamServer. Fix #305

Author: versusvoid

tests/test_context.py

@@ -2,11 +2,23 @@
 import contextvars
 import decimal
 import random
+import socket
 import weakref
 
 from uvloop import _testbase as tb
 
 
+class _Protocol(asyncio.Protocol):
+    def __init__(self, *, loop=None):
+        self.done = asyncio.Future(loop=loop)
+
+    def connection_lost(self, exc):
+        if exc is None:
+            self.done.set_result(None)
+        else:
+            self.done.set_exception(exc)
+
+
 class _ContextBaseTests:
 
     def test_task_decimal_context(self):
@@ -126,6 +138,40 @@ async def main():
         del tracked
         self.assertIsNone(ref())
 
+    def test_create_server_protocol_factory_context(self):
+        cvar = contextvars.ContextVar('cvar', default='outer')
+        factory_called_future = self.loop.create_future()
+        proto = _Protocol(loop=self.loop)
+
+        def factory():
+            try:
+                self.assertEqual(cvar.get(), 'inner')
+            except Exception as e:
+                factory_called_future.set_exception(e)
+            else:
+                factory_called_future.set_result(None)
+
+            return proto
+
+        async def test():
+            cvar.set('inner')
+            port = tb.find_free_port()
+            srv = await self.loop.create_server(factory, '127.0.0.1', port)
+
+            s = socket.socket(socket.AF_INET)
+            with s:
+                s.setblocking(False)
+                await self.loop.sock_connect(s, ('127.0.0.1', port))
+
+            try:
+                await factory_called_future
+            finally:
+                srv.close()
+                await proto.done
+                await srv.wait_closed()
+
+        self.loop.run_until_complete(test())
+
 
 class Test_UV_Context(_ContextBaseTests, tb.UVTestCase):
     pass

uvloop/handles/streamserver.pxd

@@ -7,6 +7,7 @@ cdef class UVStreamServer(UVSocketHandle):
         object protocol_factory
         bint opened
         Server _server
+        object listen_context
 
     # All "inline" methods are final
 

uvloop/handles/streamserver.pyx

@@ -8,6 +8,7 @@ cdef class UVStreamServer(UVSocketHandle):
         self.ssl_handshake_timeout = None
         self.ssl_shutdown_timeout = None
         self.protocol_factory = None
+        self.listen_context = None
 
     cdef inline _init(self, Loop loop, object protocol_factory,
                       Server server,
@@ -53,6 +54,8 @@ cdef class UVStreamServer(UVSocketHandle):
         if self.opened != 1:
             raise RuntimeError('unopened TCPServer')
 
+        self.listen_context = Context_CopyCurrent()
+
         err = uv.uv_listen(<uv.uv_stream_t*> self._handle,
                            self.backlog,
                            __uv_streamserver_on_listen)
@@ -64,7 +67,7 @@ cdef class UVStreamServer(UVSocketHandle):
     cdef inline _on_listen(self):
         cdef UVStream client
 
-        protocol = self.protocol_factory()
+        protocol = self.listen_context.run(self.protocol_factory)
 
         if self.ssl is None:
             client = self._make_new_transport(protocol, None)