Add support for connection termination listeners (#525)

The new `Connection.add_termination_listener()` method can be used
to register callbacks to be invoked when a connection has been terminated.

Co-authored-by: Elvis Pranskevichus <[email protected]>

Author: ioistired

asyncpg/_testbase/fuzzer.py

@@ -145,6 +145,10 @@ def _close_connection(self, connection):
         if conn_task is not None:
             conn_task.cancel()
 
+    def close_all_connections(self):
+        for conn in list(self.connections):
+            self.loop.call_soon_threadsafe(self._close_connection, conn)
+
 
 class Connection:
     def __init__(self, client_sock, backend_sock, proxy):
@@ -215,10 +219,11 @@ async def _read(self, sock, n):
             else:
                 return read_task.result()
         finally:
-            if not read_task.done():
-                read_task.cancel()
-            if not conn_event_task.done():
-                conn_event_task.cancel()
+            if not self.loop.is_closed():
+                if not read_task.done():
+                    read_task.cancel()
+                if not conn_event_task.done():
+                    conn_event_task.cancel()
 
     async def _write(self, sock, data):
         write_task = asyncio.ensure_future(
@@ -236,10 +241,11 @@ async def _write(self, sock, data):
             else:
                 return write_task.result()
         finally:
-            if not write_task.done():
-                write_task.cancel()
-            if not conn_event_task.done():
-                conn_event_task.cancel()
+            if not self.loop.is_closed():
+                if not write_task.done():
+                    write_task.cancel()
+                if not conn_event_task.done():
+                    conn_event_task.cancel()
 
     async def proxy_to_backend(self):
         buf = None
@@ -264,7 +270,8 @@ async def proxy_to_backend(self):
             pass
 
         finally:
-            self.loop.call_soon(self.close)
+            if not self.loop.is_closed():
+                self.loop.call_soon(self.close)
 
     async def proxy_from_backend(self):
         buf = None
@@ -289,4 +296,5 @@ async def proxy_from_backend(self):
             pass
 
         finally:
-            self.loop.call_soon(self.close)
+            if not self.loop.is_closed():
+                self.loop.call_soon(self.close)

asyncpg/connection.py

@@ -46,8 +46,8 @@ class Connection(metaclass=ConnectionMeta):
                  '_listeners', '_server_version', '_server_caps',
                  '_intro_query', '_reset_query', '_proxy',
                  '_stmt_exclusive_section', '_config', '_params', '_addr',
-                 '_log_listeners', '_cancellations', '_source_traceback',
-                 '__weakref__')
+                 '_log_listeners', '_termination_listeners', '_cancellations',
+                 '_source_traceback', '__weakref__')
 
     def __init__(self, protocol, transport, loop,
                  addr: (str, int) or str,
@@ -78,6 +78,7 @@ def __init__(self, protocol, transport, loop,
         self._listeners = {}
         self._log_listeners = set()
         self._cancellations = set()
+        self._termination_listeners = set()
 
         settings = self._protocol.get_settings()
         ver_string = settings.server_version
@@ -178,6 +179,28 @@ def remove_log_listener(self, callback):
         """
         self._log_listeners.discard(callback)
 
+    def add_termination_listener(self, callback):
+        """Add a listener that will be called when the connection is closed.
+
+        :param callable callback:
+            A callable receiving one argument:
+            **connection**: a Connection the callback is registered with.
+
+        .. versionadded:: 0.21.0
+        """
+        self._termination_listeners.add(callback)
+
+    def remove_termination_listener(self, callback):
+        """Remove a listening callback for connection termination.
+
+        :param callable callback:
+            The callable that was passed to
+            :meth:`Connection.add_termination_listener`.
+
+        .. versionadded:: 0.21.0
+        """
+        self._termination_listeners.discard(callback)
+
     def get_server_pid(self):
         """Return the PID of the Postgres server the connection is bound to."""
         return self._protocol.get_server_pid()
@@ -1120,6 +1143,7 @@ def _abort(self):
         self._protocol = None
 
     def _cleanup(self):
+        self._call_termination_listeners()
         # Free the resources associated with this connection.
         # This must be called when a connection is terminated.
 
@@ -1237,6 +1261,25 @@ def _call_log_listener(self, cb, con_ref, message):
                 'exception': ex
             })
 
+    def _call_termination_listeners(self):
+        if not self._termination_listeners:
+            return
+
+        con_ref = self._unwrap()
+        for cb in self._termination_listeners:
+            try:
+                cb(con_ref)
+            except Exception as ex:
+                self._loop.call_exception_handler({
+                    'message': (
+                        'Unhandled exception in asyncpg connection '
+                        'termination listener callback {!r}'.format(cb)
+                    ),
+                    'exception': ex
+                })
+
+        self._termination_listeners.clear()
+
     def _process_notification(self, pid, channel, payload):
         if channel not in self._listeners:
             return

tests/test_listeners.py

@@ -6,6 +6,10 @@
 
 
 import asyncio
+import os
+import platform
+import sys
+import unittest
 
 from asyncpg import _testbase as tb
 from asyncpg import exceptions
@@ -272,3 +276,41 @@ def listener1(*args):
                         pass
 
                     con.add_log_listener(listener1)
+
+
+@unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
+@unittest.skipIf(
+    platform.system() == 'Windows' and
+    sys.version_info >= (3, 8),
+    'not compatible with ProactorEventLoop which is default in Python 3.8')
+class TestConnectionTerminationListener(tb.ProxiedClusterTestCase):
+
+    async def test_connection_termination_callback_called_on_remote(self):
+
+        called = False
+
+        def close_cb(con):
+            nonlocal called
+            called = True
+
+        con = await self.connect()
+        con.add_termination_listener(close_cb)
+        self.proxy.close_all_connections()
+        try:
+            await con.fetchval('SELECT 1')
+        except Exception:
+            pass
+        self.assertTrue(called)
+
+    async def test_connection_termination_callback_called_on_local(self):
+
+        called = False
+
+        def close_cb(con):
+            nonlocal called
+            called = True
+
+        con = await self.connect()
+        con.add_termination_listener(close_cb)
+        await con.close()
+        self.assertTrue(called)