Break strong ref between Connection and Protocol.

Connection objects currently stay alive and open if they are not
explicitly closed or terminated.  GC won't happen to them because event
loops have a strong reference to their underlying Transport object.

By replacing a strong Connection<->Protocol reference with a weak one,
we are able to implement Connection.__del__() method that:

* issues a warning if a Connection object is being GCed prior
  to be explicitly closed;

* terminates the underlying Protocol and Transport, effectively closing
  the open network connection to the Postgres server.

When in asyncio debug mode (enabled by PYTHONASYNCIODEBUG env variable
or explicitly with `loop.set_debug(True)`) Connection objects save the
traceback of their origin and later use it to make the GC warning
clarer.

Addresses #323.

Author: 1st1

asyncpg/connection.py

@@ -6,11 +6,14 @@
 
 
 import asyncio
+import asyncpg
 import collections
 import collections.abc
 import itertools
 import struct
+import sys
 import time
+import traceback
 import warnings
 
 from . import compat
@@ -44,7 +47,8 @@ class Connection(metaclass=ConnectionMeta):
                  '_listeners', '_server_version', '_server_caps',
                  '_intro_query', '_reset_query', '_proxy',
                  '_stmt_exclusive_section', '_config', '_params', '_addr',
-                 '_log_listeners', '_cancellations')
+                 '_log_listeners', '_cancellations', '_source_traceback',
+                 '__weakref__')
 
     def __init__(self, protocol, transport, loop,
                  addr: (str, int) or str,
@@ -98,6 +102,27 @@ def __init__(self, protocol, transport, loop,
         # `con.execute()`, and `con.executemany()`.
         self._stmt_exclusive_section = _Atomic()
 
+        if loop.get_debug():
+            self._source_traceback = _extract_stack()
+        else:
+            self._source_traceback = None
+
+    def __del__(self):
+        if not self.is_closed() and self._protocol is not None:
+            if self._source_traceback:
+                msg = "unclosed connection {!r}; created at:\n {}".format(
+                    self, self._source_traceback)
+            else:
+                msg = (
+                    "unclosed connection {!r}; run in asyncio debug "
+                    "mode to show the traceback of connection "
+                    "origin".format(self)
+                )
+
+            warnings.warn(msg, ResourceWarning)
+            if not self._loop.is_closed():
+                self.terminate()
+
     async def add_listener(self, channel, callback):
         """Add a listener for Postgres notifications.
 
@@ -1791,4 +1816,25 @@ def _detect_server_capabilities(server_version, connection_settings):
     )
 
 
+def _extract_stack(limit=10):
+    """Replacement for traceback.extract_stack() that only does the
+    necessary work for asyncio debug mode.
+    """
+    frame = sys._getframe().f_back
+    try:
+        stack = traceback.StackSummary.extract(
+            traceback.walk_stack(frame), lookup_lines=False)
+    finally:
+        del frame
+
+    apg_path = asyncpg.__path__[0]
+    i = 0
+    while i < len(stack) and stack[i][0].startswith(apg_path):
+        i += 1
+    stack = stack[i:i + limit]
+
+    stack.reverse()
+    return ''.join(traceback.format_list(stack))
+
+
 _uid = 0

asyncpg/pool.py

@@ -467,7 +467,21 @@ async def _get_new_connection(self):
                 connection_class=self._connection_class)
 
         if self._init is not None:
-            await self._init(con)
+            try:
+                await self._init(con)
+            except Exception as ex:
+                # If a user-defined `init` function fails, we don't
+                # know if the connection is safe for re-use, hence
+                # we close it.  A new connection will be created
+                # when `acquire` is called again.
+                try:
+                    # Use `close()` to close the connection gracefully.
+                    # An exception in `init` isn't necessarily caused
+                    # by an IO or a protocol error.  close() will
+                    # do the necessary cleanup via _release_on_close().
+                    await con.close()
+                finally:
+                    raise ex
 
         return con
 

asyncpg/protocol/prepared_stmt.pxd

@@ -17,7 +17,6 @@ cdef class PreparedStatementState:
         list         row_desc
         list         parameters_desc
 
-        BaseProtocol protocol
         ConnectionSettings settings
 
         int16_t      args_num

asyncpg/protocol/prepared_stmt.pyx

@@ -14,7 +14,6 @@ cdef class PreparedStatementState:
     def __cinit__(self, str name, str query, BaseProtocol protocol):
         self.name = name
         self.query = query
-        self.protocol = protocol
         self.settings = protocol.settings
         self.row_desc = self.parameters_desc = None
         self.args_codecs = self.rows_codecs = None

asyncpg/protocol/protocol.pxd

@@ -36,7 +36,7 @@ cdef class BaseProtocol(CoreProtocol):
         object timeout_handle
         object timeout_callback
         object completed_callback
-        object connection
+        object conref
         bint is_reading
 
         str last_query
@@ -48,6 +48,8 @@ cdef class BaseProtocol(CoreProtocol):
 
         PreparedStatementState statement
 
+    cdef get_connection(self)
+
     cdef _get_timeout_impl(self, timeout)
     cdef _check_state(self)
     cdef _new_waiter(self, timeout)

asyncpg/protocol/protocol.pyx

@@ -16,6 +16,7 @@ import codecs
 import collections
 import socket
 import time
+import weakref
 
 from libc.stdint cimport int8_t, uint8_t, int16_t, uint16_t, \
                          int32_t, uint32_t, int64_t, uint64_t, \
@@ -126,7 +127,10 @@ cdef class BaseProtocol(CoreProtocol):
             self.create_future = self._create_future_fallback
 
     def set_connection(self, connection):
-        self.connection = connection
+        self.conref = weakref.ref(connection)
+
+    cdef get_connection(self):
+        return self.conref()
 
     def get_server_pid(self):
         return self.backend_pid
@@ -585,7 +589,21 @@ cdef class BaseProtocol(CoreProtocol):
     def _request_cancel(self):
         self.cancel_waiter = self.create_future()
         self.cancel_sent_waiter = self.create_future()
-        self.connection._cancel_current_command(self.cancel_sent_waiter)
+
+        con = self.get_connection()
+        if con is not None:
+            # if 'con' is None it means that the connection object has been
+            # garbage collected and that the transport will soon be aborted.
+            con._cancel_current_command(self.cancel_sent_waiter)
+        else:
+            self.loop.call_exception_handler({
+                'message': 'asyncpg.Protocol has no reference to its '
+                           'Connection object and yet a cancellation '
+                           'was requested. Please report this at '
+                           'github.com/magicstack/asyncpg.'
+            })
+            self.abort()
+
         self._set_state(PROTOCOL_CANCELLED)
 
     def _on_timeout(self, fut):
@@ -636,7 +654,7 @@ cdef class BaseProtocol(CoreProtocol):
 
     cdef inline _get_timeout_impl(self, timeout):
         if timeout is None:
-            timeout = self.connection._config.command_timeout
+            timeout = self.get_connection()._config.command_timeout
         elif timeout is NO_TIMEOUT:
             timeout = None
         else:
@@ -688,7 +706,7 @@ cdef class BaseProtocol(CoreProtocol):
                 'cannot perform operation: another operation is in progress')
         self.waiter = self.create_future()
         if timeout is not None:
-            self.timeout_handle = self.connection._loop.call_later(
+            self.timeout_handle = self.loop.call_later(
                 timeout, self.timeout_callback, self.waiter)
         self.waiter.add_done_callback(self.completed_callback)
         return self.waiter
@@ -839,10 +857,14 @@ cdef class BaseProtocol(CoreProtocol):
             self.return_extra = False
 
     cdef _on_notice(self, parsed):
-        self.connection._process_log_message(parsed, self.last_query)
+        con = self.get_connection()
+        if con is not None:
+            con._process_log_message(parsed, self.last_query)
 
     cdef _on_notification(self, pid, channel, payload):
-        self.connection._process_notification(pid, channel, payload)
+        con = self.get_connection()
+        if con is not None:
+            con._process_notification(pid, channel, payload)
 
     cdef _on_connection_lost(self, exc):
         if self.closing:

tests/test_connect.py

@@ -7,6 +7,7 @@
 
 import asyncio
 import contextlib
+import gc
 import ipaddress
 import os
 import platform
@@ -15,6 +16,7 @@
 import tempfile
 import textwrap
 import unittest
+import weakref
 
 import asyncpg
 from asyncpg import _testbase as tb
@@ -851,3 +853,44 @@ async def worker():
         tasks = [worker() for _ in range(100)]
         await asyncio.gather(*tasks, loop=self.loop)
         await pool.close()
+
+
+class TestConnectionGC(tb.ClusterTestCase):
+
+    async def _run_no_explicit_close_test(self):
+        con = await self.connect()
+        proto = con._protocol
+        conref = weakref.ref(con)
+        del con
+
+        gc.collect()
+        gc.collect()
+        gc.collect()
+
+        self.assertIsNone(conref())
+        self.assertTrue(proto.is_closed())
+
+    async def test_no_explicit_close_no_debug(self):
+        olddebug = self.loop.get_debug()
+        self.loop.set_debug(False)
+        try:
+            with self.assertWarnsRegex(
+                    ResourceWarning,
+                    r'unclosed connection.*run in asyncio debug'):
+                await self._run_no_explicit_close_test()
+        finally:
+            self.loop.set_debug(olddebug)
+
+    async def test_no_explicit_close_with_debug(self):
+        olddebug = self.loop.get_debug()
+        self.loop.set_debug(True)
+        try:
+            with self.assertWarnsRegex(ResourceWarning,
+                                       r'unclosed connection') as rw:
+                await self._run_no_explicit_close_test()
+
+            msg = rw.warning.args[0]
+            self.assertIn(' created at:\n', msg)
+            self.assertIn('in test_no_explicit_close_with_debug', msg)
+        finally:
+            self.loop.set_debug(olddebug)

tests/test_pool.py

@@ -292,7 +292,8 @@ class Error(Exception):
             pass
 
         async def setup(con):
-            nonlocal setup_calls
+            nonlocal setup_calls, last_con
+            last_con = con
             setup_calls += 1
             if setup_calls > 1:
                 cons.append(con)
@@ -302,24 +303,28 @@ async def setup(con):
 
         with self.subTest(method='setup'):
             setup_calls = 0
+            last_con = None
             cons = []
             async with self.create_pool(database='postgres',
                                         min_size=1, max_size=1,
                                         setup=setup) as pool:
                 with self.assertRaises(Error):
                     await pool.acquire()
+                self.assertTrue(last_con.is_closed())
 
                 async with pool.acquire() as con:
                     self.assertEqual(cons, ['error', con])
 
         with self.subTest(method='init'):
             setup_calls = 0
+            last_con = None
             cons = []
             async with self.create_pool(database='postgres',
                                         min_size=0, max_size=1,
                                         init=setup) as pool:
                 with self.assertRaises(Error):
                     await pool.acquire()
+                self.assertTrue(last_con.is_closed())
 
                 async with pool.acquire() as con:
                     self.assertEqual(await con.fetchval('select 1::int'), 1)