Add support for BufferedProtocol

Author: 1st1

tests/test_tcp.py

@@ -629,6 +629,154 @@ async def runner():
 
 class Test_UV_TCP(_TestTCP, tb.UVTestCase):
 
+    def test_create_server_buffered_1(self):
+        SIZE = 123123
+
+        class Proto(asyncio.BaseProtocol):
+            def connection_made(self, tr):
+                self.tr = tr
+                self.recvd = b''
+                self.data = bytearray(50)
+                self.buf = memoryview(self.data)
+
+            def get_buffer(self):
+                return self.buf
+
+            def buffer_updated(self, nbytes):
+                self.recvd += self.buf[:nbytes]
+                if self.recvd == b'a' * SIZE:
+                    self.tr.write(b'hello')
+
+            def eof_received(self):
+                pass
+
+        async def test():
+            port = tb.find_free_port()
+            srv = await self.loop.create_server(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))
+                await self.loop.sock_sendall(s, b'a' * SIZE)
+                d = await self.loop.sock_recv(s, 100)
+                self.assertEqual(d, b'hello')
+
+            srv.close()
+            await srv.wait_closed()
+
+        self.loop.run_until_complete(test())
+
+    def test_create_server_buffered_2(self):
+        class ProtoExc(asyncio.BaseProtocol):
+            def __init__(self):
+                self._lost_exc = None
+
+            def get_buffer(self):
+                1 / 0
+
+            def buffer_updated(self, nbytes):
+                pass
+
+            def connection_lost(self, exc):
+                self._lost_exc = exc
+
+            def eof_received(self):
+                pass
+
+        class ProtoZeroBuf1(asyncio.BaseProtocol):
+            def __init__(self):
+                self._lost_exc = None
+
+            def get_buffer(self):
+                return bytearray(0)
+
+            def buffer_updated(self, nbytes):
+                pass
+
+            def connection_lost(self, exc):
+                self._lost_exc = exc
+
+            def eof_received(self):
+                pass
+
+        class ProtoZeroBuf2(asyncio.BaseProtocol):
+            def __init__(self):
+                self._lost_exc = None
+
+            def get_buffer(self):
+                return memoryview(bytearray(0))
+
+            def buffer_updated(self, nbytes):
+                pass
+
+            def connection_lost(self, exc):
+                self._lost_exc = exc
+
+            def eof_received(self):
+                pass
+
+        class ProtoUpdatedError(asyncio.BaseProtocol):
+            def __init__(self):
+                self._lost_exc = None
+
+            def get_buffer(self):
+                return memoryview(bytearray(100))
+
+            def buffer_updated(self, nbytes):
+                raise RuntimeError('oups')
+
+            def connection_lost(self, exc):
+                self._lost_exc = exc
+
+            def eof_received(self):
+                pass
+
+        async def test(proto_factory, exc_type, exc_re):
+            port = tb.find_free_port()
+            proto = proto_factory()
+            srv = await self.loop.create_server(
+                lambda: proto, '127.0.0.1', port)
+
+            try:
+                s = socket.socket(socket.AF_INET)
+                with s:
+                    s.setblocking(False)
+                    await self.loop.sock_connect(s, ('127.0.0.1', port))
+                    await self.loop.sock_sendall(s, b'a')
+                    d = await self.loop.sock_recv(s, 100)
+                    if not d:
+                        raise ConnectionResetError
+            except ConnectionResetError:
+                pass
+            else:
+                self.fail("server didn't abort the connection")
+                return
+            finally:
+                srv.close()
+                await srv.wait_closed()
+
+            if proto._lost_exc is None:
+                self.fail("connection_lost() was not called")
+                return
+
+            with self.assertRaisesRegex(exc_type, exc_re):
+                raise proto._lost_exc
+
+        self.loop.set_exception_handler(lambda loop, ctx: None)
+
+        self.loop.run_until_complete(
+            test(ProtoExc, RuntimeError, 'unhandled error .* get_buffer'))
+
+        self.loop.run_until_complete(
+            test(ProtoZeroBuf1, RuntimeError, 'unhandled error .* get_buffer'))
+
+        self.loop.run_until_complete(
+            test(ProtoZeroBuf2, RuntimeError, 'unhandled error .* get_buffer'))
+
+        self.loop.run_until_complete(
+            test(ProtoUpdatedError, RuntimeError, r'^oups$'))
+
     def test_transport_get_extra_info(self):
         # This tests is only for uvloop.  asyncio should pass it
         # too in Python 3.6.

uvloop/handles/basetransport.pxd

@@ -38,7 +38,9 @@ cdef class UVBaseTransport(UVSocketHandle):
 
     cdef inline _set_server(self, Server server)
     cdef inline _set_waiter(self, object waiter)
-    cdef inline _set_protocol(self, object protocol)
+
+    cdef _set_protocol(self, object protocol)
+    cdef _clear_protocol(self)
 
     cdef inline _init_protocol(self)
     cdef inline _add_extra_info(self, str name, object obj)

uvloop/handles/basetransport.pyx

@@ -201,8 +201,7 @@ cdef class UVBaseTransport(UVSocketHandle):
             if self._protocol_connected:
                 self._protocol.connection_lost(exc)
         finally:
-            self._protocol = None
-            self._protocol_data_received = None
+            self._clear_protocol()
 
             self._close()
 
@@ -223,14 +222,18 @@ cdef class UVBaseTransport(UVSocketHandle):
 
         self._waiter = waiter
 
-    cdef inline _set_protocol(self, object protocol):
+    cdef _set_protocol(self, object protocol):
         self._protocol = protocol
         # Store a reference to the bound method directly
         try:
             self._protocol_data_received = protocol.data_received
         except AttributeError:
             pass
 
+    cdef _clear_protocol(self):
+        self._protocol = None
+        self._protocol_data_received = None
+
     cdef inline _init_protocol(self):
         self._loop._track_transport(self)
         if self._protocol is None:
@@ -263,6 +266,9 @@ cdef class UVBaseTransport(UVSocketHandle):
 
     def set_protocol(self, protocol):
         self._set_protocol(protocol)
+        if self._is_reading():
+            self._stop_reading()
+            self._start_reading()
 
     def _force_close(self, exc):
         # Used by SSLProto.  Might be removed in the future.

uvloop/handles/pipe.pyx

@@ -207,7 +207,7 @@ cdef void __pipe_connect_callback(uv.uv_connect_t* req, int status) with gil:
     try:
         transport._on_connect(exc)
     except BaseException as ex:
-        wrapper.transport._error(ex, False)
+        wrapper.transport._fatal_error(ex, False)
     finally:
         wrapper.on_done()
 

uvloop/handles/stream.pxd

@@ -4,10 +4,18 @@ cdef class UVStream(UVBaseTransport):
         bint __shutting_down
         bint __reading
         bint __read_error_close
+
+        bint __buffered
+        object _protocol_get_buffer
+        object _protocol_buffer_updated
+
         bint _eof
         list _buffer
         size_t _buffer_size
 
+        Py_buffer _read_pybuf
+        bint _read_pybuf_acquired
+
     # All "inline" methods are final
 
     cdef inline _init(self, Loop loop, object protocol, Server server,
@@ -29,7 +37,6 @@ cdef class UVStream(UVBaseTransport):
     cdef _close(self)
 
     cdef inline _on_accept(self)
-    cdef inline _on_read(self, bytes buf)
     cdef inline _on_eof(self)
     cdef inline _on_write(self)
     cdef inline _on_connect(self, object exc)