From 93a6f79afb4251fa48edf6959aa87532eee91ef5 Mon Sep 17 00:00:00 2001
From: Andrew Geng <pteromys@gmail.com>
Date: Mon, 9 Oct 2023 15:51:22 -0400
Subject: [PATCH] Cut BaseProtocol circular reference on close. (#1049)

A bound method contains a reference to the instance it's bound to.
Most of the time, bound methods are created lazily at access time by
the descriptor protocol and discarded after calling. But saving a bound
method as another attribute on the instance creates a long-lived cycle,
here `.timeout_callback.__self__`, that needs to be explicitly broken
if we don't want to wake up python's garbage collector to do it.

Without this change, the new assertion in the tests would fail, and
`pytest --pdb` would show the bound methods `_on_timeout` and
`_on_waiter_completed` at the end of `p gc.get_referrers(protoref())`.

[Also, unset `transport` in `Protocol.abort()` to break another cycle]

Co-authored-by: Elvis Pranskevichus <elvis@edgedb.com>
---
 asyncpg/protocol/protocol.pxd |  2 --
 asyncpg/protocol/protocol.pyx |  7 +++----
 tests/test_connect.py         | 30 ++++++++++++++++++++++--------
 3 files changed, 25 insertions(+), 14 deletions(-)

diff --git a/asyncpg/protocol/protocol.pxd b/asyncpg/protocol/protocol.pxd
index 5f144e55..a9ac8d5f 100644
--- a/asyncpg/protocol/protocol.pxd
+++ b/asyncpg/protocol/protocol.pxd
@@ -39,8 +39,6 @@ cdef class BaseProtocol(CoreProtocol):
         bint return_extra
         object create_future
         object timeout_handle
-        object timeout_callback
-        object completed_callback
         object conref
         type record_class
         bint is_reading
diff --git a/asyncpg/protocol/protocol.pyx b/asyncpg/protocol/protocol.pyx
index 76c62dfc..b43b0e9c 100644
--- a/asyncpg/protocol/protocol.pyx
+++ b/asyncpg/protocol/protocol.pyx
@@ -98,8 +98,6 @@ cdef class BaseProtocol(CoreProtocol):
         self.writing_allowed.set()
 
         self.timeout_handle = None
-        self.timeout_callback = self._on_timeout
-        self.completed_callback = self._on_waiter_completed
 
         self.queries_count = 0
 
@@ -607,6 +605,7 @@ cdef class BaseProtocol(CoreProtocol):
         self._handle_waiter_on_connection_lost(None)
         self._terminate()
         self.transport.abort()
+        self.transport = None
 
     @cython.iterable_coroutine
     async def close(self, timeout):
@@ -777,8 +776,8 @@ cdef class BaseProtocol(CoreProtocol):
         self.waiter = self.create_future()
         if timeout is not None:
             self.timeout_handle = self.loop.call_later(
-                timeout, self.timeout_callback, self.waiter)
-        self.waiter.add_done_callback(self.completed_callback)
+                timeout, self._on_timeout, self.waiter)
+        self.waiter.add_done_callback(self._on_waiter_completed)
         return self.waiter
 
     cdef _on_result__connect(self, object waiter):
diff --git a/tests/test_connect.py b/tests/test_connect.py
index f61db61a..1af074f1 100644
--- a/tests/test_connect.py
+++ b/tests/test_connect.py
@@ -7,6 +7,7 @@
 
 import asyncio
 import contextlib
+import gc
 import ipaddress
 import os
 import pathlib
@@ -1846,14 +1847,27 @@ async def worker():
 class TestConnectionGC(tb.ClusterTestCase):
 
     async def _run_no_explicit_close_test(self):
-        con = await self.connect()
-        await con.fetchval("select 123")
-        proto = con._protocol
-        conref = weakref.ref(con)
-        del con
-
-        self.assertIsNone(conref())
-        self.assertTrue(proto.is_closed())
+        gc_was_enabled = gc.isenabled()
+        gc.disable()
+        try:
+            con = await self.connect()
+            await con.fetchval("select 123")
+            proto = con._protocol
+            conref = weakref.ref(con)
+            del con
+
+            self.assertIsNone(conref())
+            self.assertTrue(proto.is_closed())
+
+            # tick event loop; asyncio.selector_events._SelectorSocketTransport
+            # needs a chance to close itself and remove its reference to proto
+            await asyncio.sleep(0)
+            protoref = weakref.ref(proto)
+            del proto
+            self.assertIsNone(protoref())
+        finally:
+            if gc_was_enabled:
+                gc.enable()
 
     async def test_no_explicit_close_no_debug(self):
         olddebug = self.loop.get_debug()