WIP: use context in transport/UVHandle directly

Author: fantix

tests/test_context.py

@@ -3,16 +3,26 @@
 import decimal
 import random
 import socket
+import unittest
 import weakref
 
 from uvloop import _testbase as tb
 
 
 class _Protocol(asyncio.Protocol):
-    def __init__(self, *, loop=None):
+    def __init__(self, cvar, *, loop=None):
+        self.cvar = cvar
+        self.transport = None
+        self.connection_made_fut = asyncio.Future(loop=loop)
+        self.connection_lost_ctx = None
         self.done = asyncio.Future(loop=loop)
 
+    def connection_made(self, transport):
+        self.transport = transport
+        self.connection_made_fut.set_result(self.cvar.get())
+
     def connection_lost(self, exc):
+        self.connection_lost_ctx = self.cvar.get()
         if exc is None:
             self.done.set_result(None)
         else:
@@ -141,7 +151,7 @@ async def main():
     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)
+        proto = _Protocol(cvar, loop=self.loop)
 
         def factory():
             try:
@@ -172,6 +182,71 @@ async def test():
 
         self.loop.run_until_complete(test())
 
+    def test_create_server_connection_made_lost(self):
+        cvar = contextvars.ContextVar('cvar', default='outer')
+        proto = _Protocol(cvar, loop=self.loop)
+
+        async def test():
+            cvar.set('inner')
+            port = tb.find_free_port()
+            srv = await self.loop.create_server(
+                lambda: proto, '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:
+                inner = await proto.connection_made_fut
+                self.assertEqual(inner, "inner")
+
+                await proto.done
+                self.assertEqual(proto.connection_lost_ctx, "inner")
+            finally:
+                srv.close()
+                await srv.wait_closed()
+
+        self.loop.run_until_complete(test())
+
+    def test_create_server_manual_connection_lost(self):
+        if self.implementation == 'asyncio':
+            raise unittest.SkipTest('this seems to be a bug in asyncio')
+
+        cvar = contextvars.ContextVar('cvar', default='outer')
+        proto = _Protocol(cvar, loop=self.loop)
+
+        async def close():
+            cvar.set('closing')
+            proto.transport.close()
+
+        async def test():
+            cvar.set('inner')
+            port = tb.find_free_port()
+            srv = await self.loop.create_server(
+                lambda: proto, '127.0.0.1', port,
+            )
+
+            s = socket.socket(socket.AF_INET)
+            s.setblocking(False)
+            await self.loop.sock_connect(s, ('127.0.0.1', port))
+
+            try:
+                inner = await proto.connection_made_fut
+                self.assertEqual(inner, "inner")
+
+                await asyncio.ensure_future(close())
+
+                await proto.done
+                self.assertEqual(proto.connection_lost_ctx, "inner")
+            finally:
+                s.close()
+                srv.close()
+                await srv.wait_closed()
+
+        self.loop.run_until_complete(test())
+
 
 class Test_UV_Context(_ContextBaseTests, tb.UVTestCase):
     pass

uvloop/handles/basetransport.pyx

@@ -26,15 +26,15 @@ cdef class UVBaseTransport(UVSocketHandle):
             new_MethodHandle(self._loop,
                              "UVTransport._call_connection_made",
                              <method_t>self._call_connection_made,
-                             None,
+                             self.context,
                              self))
 
     cdef inline _schedule_call_connection_lost(self, exc):
         self._loop._call_soon_handle(
             new_MethodHandle1(self._loop,
                               "UVTransport._call_connection_lost",
                               <method1_t>self._call_connection_lost,
-                              None,
+                              self.context,
                               self, exc))
 
     cdef _fatal_error(self, exc, throw, reason=None):

uvloop/handles/handle.pxd

@@ -5,6 +5,7 @@ cdef class UVHandle:
         readonly _source_traceback
         bint _closed
         bint _inited
+        object context
 
         # Added to enable current UDPTransport implementation,
         # which doesn't use libuv handles.

uvloop/handles/pipe.pyx

@@ -73,9 +73,10 @@ cdef class UnixServer(UVStreamServer):
 
         self._mark_as_open()
 
-    cdef UVStream _make_new_transport(self, object protocol, object waiter):
+    cdef UVStream _make_new_transport(self, object protocol, object waiter,
+                                      object context):
         cdef UnixTransport tr
-        tr = UnixTransport.new(self._loop, protocol, self._server, waiter)
+        tr = UnixTransport.new(self._loop, protocol, self._server, waiter)  # TODO: context
         return <UVStream>tr
 
 

uvloop/handles/streamserver.pxd

@@ -23,4 +23,5 @@ cdef class UVStreamServer(UVSocketHandle):
     cdef inline listen(self)
     cdef inline _on_listen(self)
 
-    cdef UVStream _make_new_transport(self, object protocol, object waiter)
+    cdef UVStream _make_new_transport(self, object protocol, object waiter,
+                                      object context)

uvloop/handles/streamserver.pyx

@@ -70,7 +70,8 @@ cdef class UVStreamServer(UVSocketHandle):
         protocol = self.listen_context.run(self.protocol_factory)
 
         if self.ssl is None:
-            client = self._make_new_transport(protocol, None)
+            client = self._make_new_transport(protocol, None,
+                                              self.listen_context)
 
         else:
             waiter = self._loop._new_future()
@@ -83,7 +84,8 @@ cdef class UVStreamServer(UVSocketHandle):
                 ssl_handshake_timeout=self.ssl_handshake_timeout,
                 ssl_shutdown_timeout=self.ssl_shutdown_timeout)
 
-            client = self._make_new_transport(ssl_protocol, None)
+            client = self._make_new_transport(ssl_protocol, None,
+                                              self.listen_context)
 
             waiter.add_done_callback(
                 ft_partial(self.__on_ssl_connected, client))
@@ -112,7 +114,8 @@ cdef class UVStreamServer(UVSocketHandle):
     cdef inline _mark_as_open(self):
         self.opened = 1
 
-    cdef UVStream _make_new_transport(self, object protocol, object waiter):
+    cdef UVStream _make_new_transport(self, object protocol, object waiter,
+                                      object context):
         raise NotImplementedError
 
     def __on_ssl_connected(self, transport, fut):

uvloop/handles/tcp.pxd

@@ -23,4 +23,4 @@ cdef class TCPTransport(UVStream):
 
     @staticmethod
     cdef TCPTransport new(Loop loop, object protocol, Server server,
-                          object waiter)
+                          object waiter, object context)

uvloop/handles/tcp.pyx

@@ -91,9 +91,11 @@ cdef class TCPServer(UVStreamServer):
         else:
             self._mark_as_open()
 
-    cdef UVStream _make_new_transport(self, object protocol, object waiter):
+    cdef UVStream _make_new_transport(self, object protocol, object waiter,
+                                      object context):
         cdef TCPTransport tr
-        tr = TCPTransport.new(self._loop, protocol, self._server, waiter)
+        tr = TCPTransport.new(self._loop, protocol, self._server, waiter,
+                              context)
         return <UVStream>tr
 
 
@@ -102,7 +104,7 @@ cdef class TCPTransport(UVStream):
 
     @staticmethod
     cdef TCPTransport new(Loop loop, object protocol, Server server,
-                          object waiter):
+                          object waiter, object context):
 
         cdef TCPTransport handle
         handle = TCPTransport.__new__(TCPTransport)
@@ -111,6 +113,7 @@ cdef class TCPTransport(UVStream):
         handle.__peername_set = 0
         handle.__sockname_set = 0
         handle._set_nodelay()
+        handle.context = context
         return handle
 
     cdef _set_nodelay(self):

uvloop/loop.pyx

@@ -1921,7 +1921,7 @@ cdef class Loop:
                 tr = None
                 try:
                     waiter = self._new_future()
-                    tr = TCPTransport.new(self, protocol, None, waiter)
+                    tr = TCPTransport.new(self, protocol, None, waiter, None)  # TODO: context
 
                     if lai is not NULL:
                         lai_iter = lai
@@ -1981,7 +1981,7 @@ cdef class Loop:
             sock.setblocking(False)
 
             waiter = self._new_future()
-            tr = TCPTransport.new(self, protocol, None, waiter)
+            tr = TCPTransport.new(self, protocol, None, waiter, None)  # TODO: context
             try:
                 # libuv will make socket non-blocking
                 tr._open(sock.fileno())
@@ -2595,7 +2595,7 @@ cdef class Loop:
                 self, protocol, None, transport_waiter)
         elif sock.family in (uv.AF_INET, uv.AF_INET6):
             transport = <UVStream>TCPTransport.new(
-                self, protocol, None, transport_waiter)
+                self, protocol, None, transport_waiter, None)  # TODO: context
 
         if transport is None:
             raise ValueError(