Support timeouts in Connection.close() and Pool.release()

Connection.close() and Pool.release() each gained the new timeout
parameter.  The pool.acquire() context manager now applies the
passed timeout to __aexit__() as well.

Connection.close() is now actually graceful.  Instead of simply dropping
the connection, it attempts to cancel the running query (if any), asks
the server to terminate the connection and waits for the connection to
terminate.

To test all this properly, implement a TCP proxy, which emulates sudden
connectivity loss (i.e. packets not reaching the server).

Closes: #220

Author: elprans

asyncpg/_testbase.py renamed to asyncpg/_testbase/__init__.py

@@ -21,6 +21,8 @@
 from asyncpg import connection as pg_connection
 from asyncpg import pool as pg_pool
 
+from . import fuzzer
+
 
 @contextlib.contextmanager
 def silence_asyncio_long_exec_warning():
@@ -36,6 +38,14 @@ def flt(log_record):
         logger.removeFilter(flt)
 
 
+def with_timeout(timeout):
+    def wrap(func):
+        func.__timeout__ = timeout
+        return func
+
+    return wrap
+
+
 class TestCaseMeta(type(unittest.TestCase)):
 
     @staticmethod
@@ -64,7 +74,11 @@ def __new__(mcls, name, bases, ns):
         for methname, meth in mcls._iter_methods(bases, ns):
             @functools.wraps(meth)
             def wrapper(self, *args, __meth__=meth, **kwargs):
-                self.loop.run_until_complete(__meth__(self, *args, **kwargs))
+                coro = __meth__(self, *args, **kwargs)
+                timeout = getattr(meth, '__timeout__', 10.0)
+                if timeout:
+                    coro = asyncio.wait_for(coro, timeout)
+                self.loop.run_until_complete(coro)
             ns[methname] = wrapper
 
         return super().__new__(mcls, name, bases, ns)
@@ -153,7 +167,8 @@ def _start_default_cluster(server_settings={}):
 
 
 def _shutdown_cluster(cluster):
-    cluster.stop()
+    if cluster.get_status() == 'running':
+        cluster.stop()
     cluster.destroy()
 
 
@@ -193,17 +208,70 @@ def setUpClass(cls):
         super().setUpClass()
         cls.setup_cluster()
 
+    @classmethod
+    def get_connection_spec(cls):
+        return cls.cluster.get_connection_spec()
+
     def create_pool(self, pool_class=pg_pool.Pool, **kwargs):
-        conn_spec = self.cluster.get_connection_spec()
+        conn_spec = self.get_connection_spec()
         conn_spec.update(kwargs)
         return create_pool(loop=self.loop, pool_class=pool_class, **conn_spec)
 
+    def connect(self, **kwargs):
+        conn_spec = self.get_connection_spec()
+        conn_spec.update(kwargs)
+        return pg_connection.connect(**conn_spec, loop=self.loop)
+
     @classmethod
     def start_cluster(cls, ClusterCls, *,
                       cluster_kwargs={}, server_settings={}):
         return _start_cluster(ClusterCls, cluster_kwargs, server_settings)
 
 
+class ProxiedClusterTestCase(ClusterTestCase):
+    @classmethod
+    def get_server_settings(cls):
+        settings = dict(super().get_server_settings())
+        settings['listen_addresses'] = '127.0.0.1'
+        return settings
+
+    @classmethod
+    def get_proxy_settings(cls):
+        return {'fuzzing-mode': None}
+
+    @classmethod
+    def setUpClass(cls):
+        super().setUpClass()
+        conn_spec = cls.cluster.get_connection_spec()
+        host = conn_spec.get('host')
+        if not host:
+            host = '127.0.0.1'
+        elif host.startswith('/'):
+            host = '127.0.0.1'
+        cls.proxy = fuzzer.TCPFuzzingProxy(
+            backend_host=host,
+            backend_port=conn_spec['port'],
+            loop=cls.loop
+        )
+        cls.loop.run_until_complete(cls.proxy.start())
+
+    @classmethod
+    def tearDownClass(cls):
+        cls.loop.run_until_complete(cls.proxy.stop())
+        super().tearDownClass()
+
+    @classmethod
+    def get_connection_spec(cls):
+        conn_spec = cls.cluster.get_connection_spec()
+        conn_spec['host'] = cls.proxy.listening_addr
+        conn_spec['port'] = cls.proxy.listening_port
+        return conn_spec
+
+    def tearDown(self):
+        self.proxy.reset()
+        super().tearDown()
+
+
 def with_connection_options(**options):
     if not options:
         raise ValueError('no connection options were specified')

asyncpg/_testbase/fuzzer.py

@@ -0,0 +1,193 @@
+# Copyright (C) 2016-present the asyncpg authors and contributors
+# <see AUTHORS file>
+#
+# This module is part of asyncpg and is released under
+# the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
+
+
+import asyncio
+import socket
+import typing
+
+from asyncpg import cluster
+
+
+class StopServer(Exception):
+    pass
+
+
+class TCPFuzzingProxy:
+    def __init__(self, *, listening_addr: str='127.0.0.1',
+                 listening_port: typing.Optional[int]=None,
+                 backend_host: str, backend_port: int,
+                 settings: typing.Optional[dict]=None,
+                 loop: typing.Optional[asyncio.AbstractEventLoop]) -> None:
+        self.listening_addr = listening_addr
+        self.listening_port = listening_port
+        self.backend_host = backend_host
+        self.backend_port = backend_port
+        self.settings = settings or {}
+        self.loop = loop or asyncio.get_event_loop()
+        self.connections = {}
+        self.connectivity = asyncio.Event(loop=self.loop)
+        self.connectivity.set()
+        self.connectivity_loss = asyncio.Event(loop=self.loop)
+        self.stop_event = asyncio.Event(loop=self.loop)
+        self.sock = None
+
+    async def _wait(self, work):
+        done, _ = await asyncio.wait([work, self.stop_event.wait()],
+                                     return_when=asyncio.FIRST_COMPLETED,
+                                     loop=self.loop)
+        if self.stop_event.is_set():
+            raise StopServer()
+        else:
+            return list(done)[0].result()
+
+    async def start(self):
+        if self.listening_port is None:
+            self.listening_port = cluster.find_available_port()
+
+        self.sock = socket.socket()
+        self.sock.bind((self.listening_addr, self.listening_port))
+        self.sock.listen(50)
+        self.sock.setblocking(False)
+        self.loop.create_task(self.listen())
+
+    async def stop(self):
+        self.stop_event.set()
+        for conn, conn_task in self.connections.items():
+            conn_task.cancel()
+            conn.close()
+        self.sock.close()
+        await asyncio.sleep(0.2, loop=self.loop)
+
+    async def listen(self):
+        while True:
+            try:
+                client_sock, _ = await self._wait(
+                    self.loop.sock_accept(self.sock))
+
+                backend_sock = socket.socket()
+                backend_sock.setblocking(False)
+
+                await self._wait(self.loop.sock_connect(
+                    backend_sock, (self.backend_host, self.backend_port)))
+            except StopServer:
+                break
+
+            conn = Connection(
+                client_sock, backend_sock,
+                connectivity=self.connectivity,
+                connectivity_loss=self.connectivity_loss, loop=self.loop)
+            conn_task = self.loop.create_task(conn.handle())
+            self.connections[conn] = conn_task
+
+    def trigger_connectivity_loss(self):
+        self.connectivity.clear()
+        self.connectivity_loss.set()
+
+    def restore_connectivity(self):
+        self.connectivity.set()
+        self.connectivity_loss.clear()
+
+    def reset(self):
+        self.restore_connectivity()
+
+
+class Connection:
+    def __init__(self, client_sock, backend_sock,
+                 connectivity, connectivity_loss, loop):
+        self.client_sock = client_sock
+        self.backend_sock = backend_sock
+        self.loop = loop
+        self.connectivity = connectivity
+        self.connectivity_loss = connectivity_loss
+        self.proxy_to_backend_task = None
+        self.proxy_from_backend_task = None
+
+    def close(self):
+        if self.proxy_to_backend_task is not None:
+            self.proxy_to_backend_task.cancel()
+
+        if self.proxy_from_backend_task is not None:
+            self.proxy_from_backend_task.cancel()
+
+    async def handle(self):
+        self.proxy_to_backend_task = self.loop.create_task(
+            self.proxy_to_backend())
+
+        self.proxy_from_backend_task = self.loop.create_task(
+            self.proxy_from_backend())
+
+        try:
+            await asyncio.gather(
+                self.proxy_to_backend_task, self.proxy_from_backend_task,
+                loop=self.loop)
+        finally:
+            self.client_sock.close()
+            self.backend_sock.close()
+
+    async def _read(self, sock, n):
+        done, _ = await asyncio.wait([
+            self.loop.sock_recv(sock, n), self.connectivity_loss.wait()],
+            return_when=asyncio.FIRST_COMPLETED,
+            loop=self.loop)
+        if self.connectivity_loss.is_set():
+            return None
+        else:
+            return list(done)[0].result()
+
+    async def _write(self, sock, data):
+        done, _ = await asyncio.wait([
+            self.loop.sock_sendall(sock, data), self.connectivity_loss.wait()],
+            return_when=asyncio.FIRST_COMPLETED,
+            loop=self.loop)
+        if self.connectivity_loss.is_set():
+            return
+        else:
+            return list(done)[0].result()
+
+    async def proxy_to_backend(self):
+        try:
+            buf = None
+
+            while True:
+                await self.connectivity.wait()
+                if buf is not None:
+                    data = buf
+                    buf = None
+                else:
+                    data = await self._read(self.client_sock, 4096)
+                if data == b'':
+                    self.close()
+                    break
+                if self.connectivity_loss.is_set():
+                    if data:
+                        buf = data
+                    continue
+                await self._write(self.backend_sock, data)
+        finally:
+            self.client_sock.shutdown()
+
+    async def proxy_from_backend(self):
+        try:
+            buf = None
+
+            while True:
+                await self.connectivity.wait()
+                if buf is not None:
+                    data = buf
+                    buf = None
+                else:
+                    data = await self._read(self.backend_sock, 4096)
+                if data == b'':
+                    self.close()
+                    break
+                if self.connectivity_loss.is_set():
+                    if data:
+                        buf = data
+                    continue
+                await self._write(self.client_sock, data)
+        finally:
+            self.backend_sock.shutdown()

asyncpg/connection.py

@@ -958,15 +958,22 @@ def is_closed(self):
         """
         return not self._protocol.is_connected() or self._aborted
 
-    async def close(self):
-        """Close the connection gracefully."""
+    async def close(self, *, timeout=None):
+        """Close the connection gracefully.
+
+        :param float timeout:
+            Optional timeout value in seconds.
+
+        .. versionchanged:: 0.14.0
+           Added the *timeout* parameter.
+        """
         if self.is_closed():
             return
         self._mark_stmts_as_closed()
         self._listeners.clear()
         self._log_listeners.clear()
         self._aborted = True
-        await self._protocol.close()
+        await self._protocol.close(timeout)
 
     def terminate(self):
         """Terminate the connection without waiting for pending data."""
@@ -976,13 +983,13 @@ def terminate(self):
         self._aborted = True
         self._protocol.abort()
 
-    async def reset(self):
+    async def reset(self, *, timeout=None):
         self._check_open()
         self._listeners.clear()
         self._log_listeners.clear()
         reset_query = self._get_reset_query()
         if reset_query:
-            await self.execute(reset_query)
+            await self.execute(reset_query, timeout=timeout)
 
     def _check_open(self):
         if self.is_closed():