Fix _StatementCache.clear() PS memory leak

See #416

Author: un-def

asyncpg/connection.py

@@ -1792,13 +1792,18 @@ def iter_statements(self):
         return (e._statement for e in self._entries.values())
 
     def clear(self):
-        # First, make sure that we cancel all scheduled callbacks.
-        for entry in self._entries.values():
-            self._clear_entry_callback(entry)
+        # Store entries for later.
+        entries = tuple(self._entries.values())
 
         # Clear the entries dict.
         self._entries.clear()
 
+        # Make sure that we cancel all scheduled callbacks
+        # and call on_remove callback for each entry.
+        for entry in entries:
+            self._clear_entry_callback(entry)
+            self._on_remove(entry._statement)
+
     def _set_entry_timeout(self, entry):
         # Clear the existing timeout.
         self._clear_entry_callback(entry)

tests/test_cache_invalidation.py

@@ -13,6 +13,20 @@
 
 
 class TestCacheInvalidation(tb.ConnectedTestCase):
+
+    def _get_cached_statements(self, connection=None):
+        if connection is None:
+            connection = self.con
+        return list(connection._stmt_cache.iter_statements())
+
+    def _check_statements_are_not_closed(self, statements):
+        self.assertGreater(len(statements), 0)
+        self.assertTrue(all(not s.closed for s in statements))
+
+    def _check_statements_are_closed(self, statements):
+        self.assertGreater(len(statements), 0)
+        self.assertTrue(all(s.closed for s in statements))
+
     async def test_prepare_cache_invalidation_silent(self):
         await self.con.execute('CREATE TABLE tab1(a int, b int)')
 
@@ -21,11 +35,16 @@ async def test_prepare_cache_invalidation_silent(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, 2))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             await self.con.execute(
                 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, '2'))
+
+            self._check_statements_are_closed(statements)
         finally:
             await self.con.execute('DROP TABLE tab1')
 
@@ -37,6 +56,9 @@ async def test_prepare_cache_invalidation_in_transaction(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, 2))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             await self.con.execute(
                 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 
@@ -45,6 +67,8 @@ async def test_prepare_cache_invalidation_in_transaction(self):
                 async with self.con.transaction():
                     result = await self.con.fetchrow('SELECT * FROM tab1')
 
+            self._check_statements_are_closed(statements)
+
             # This is now OK,
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, '2'))
@@ -69,6 +93,12 @@ async def test_prepare_cache_invalidation_in_pool(self):
             result = await con2.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, 2))
 
+            statements1 = self._get_cached_statements(con1)
+            self._check_statements_are_not_closed(statements1)
+
+            statements2 = self._get_cached_statements(con2)
+            self._check_statements_are_not_closed(statements2)
+
             await self.con.execute(
                 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 
@@ -77,6 +107,9 @@ async def test_prepare_cache_invalidation_in_pool(self):
             result = await con1.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, '2'))
 
+            self._check_statements_are_closed(statements1)
+            self._check_statements_are_closed(statements2)
+
             async with con2.transaction():
                 # This should work, as con1 should have invalidated
                 # the plan cache.
@@ -98,11 +131,17 @@ async def test_type_cache_invalidation_in_transaction(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             async with self.con.transaction():
                 await self.con.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
                 with self.assertRaisesRegex(
                         asyncpg.OutdatedSchemaCacheError, ERRNUM):
                     await self.con.fetchrow('SELECT * FROM tab1')
+
+                self._check_statements_are_closed(statements)
+
                 # The second request must be correct (cache was dropped):
                 result = await self.con.fetchrow('SELECT * FROM tab1')
                 self.assertEqual(result, (1, (2, 3, None)))
@@ -123,13 +162,19 @@ async def test_type_cache_invalidation_in_cancelled_transaction(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             try:
                 async with self.con.transaction():
                     await self.con.execute(
                         'ALTER TYPE typ1 ADD ATTRIBUTE c text')
                     with self.assertRaisesRegex(
                             asyncpg.OutdatedSchemaCacheError, ERRNUM):
                         await self.con.fetchrow('SELECT * FROM tab1')
+
+                    self._check_statements_are_closed(statements)
+
                     # The second request must be correct (cache was dropped):
                     result = await self.con.fetchrow('SELECT * FROM tab1')
                     self.assertEqual(result, (1, (2, 3, None)))
@@ -158,13 +203,19 @@ async def test_prepared_type_cache_invalidation(self):
             result = await prep.fetchrow()
             self.assertEqual(result, (1, (2, 3)))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             try:
                 async with self.con.transaction():
                     await self.con.execute(
                         'ALTER TYPE typ1 ADD ATTRIBUTE c text')
                     with self.assertRaisesRegex(
                             asyncpg.OutdatedSchemaCacheError, ERRNUM):
                         await prep.fetchrow()
+
+                    self._check_statements_are_closed(statements)
+
                     # PS has its local cache for types codecs, even after the
                     # cache cleanup it is not possible to use it.
                     # That's why it is marked as closed.
@@ -206,11 +257,16 @@ async def test_type_cache_invalidation_on_drop_type_attr(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3, 'x')))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             await self.con.execute('ALTER TYPE typ1 DROP ATTRIBUTE x')
             with self.assertRaisesRegex(
                     asyncpg.OutdatedSchemaCacheError, ERRNUM):
                 await self.con.fetchrow('SELECT * FROM tab1')
 
+            self._check_statements_are_closed(statements)
+
             # This is now OK, the cache is filled after being dropped.
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (3, 'x')))
@@ -228,6 +284,9 @@ async def test_type_cache_invalidation_on_change_attr(self):
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements = self._get_cached_statements()
+            self._check_statements_are_not_closed(statements)
+
             # It is slightly artificial, but can take place in transactional
             # schema changing. Nevertheless, if the code checks and raises it
             # the most probable reason is a difference with the cache type.
@@ -237,6 +296,8 @@ async def test_type_cache_invalidation_on_change_attr(self):
                     asyncpg.OutdatedSchemaCacheError, ERRTYP):
                 await self.con.fetchrow('SELECT * FROM tab1')
 
+            self._check_statements_are_closed(statements)
+
             # This is now OK, the cache is filled after being dropped.
             result = await self.con.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, None)))
@@ -265,9 +326,15 @@ async def test_type_cache_invalidation_in_pool(self):
             result = await con1.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements1 = self._get_cached_statements(con1)
+            self._check_statements_are_not_closed(statements1)
+
             result = await con2.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements2 = self._get_cached_statements(con2)
+            self._check_statements_are_not_closed(statements2)
+
             # Create the same schema in the "testdb", fetch data which caches
             # type info.
             con_chk = await pool_chk.acquire()
@@ -277,6 +344,9 @@ async def test_type_cache_invalidation_in_pool(self):
             result = await con_chk.fetchrow('SELECT * FROM tab1')
             self.assertEqual(result, (1, (2, 3)))
 
+            statements_chk = self._get_cached_statements(con_chk)
+            self._check_statements_are_not_closed(statements_chk)
+
             # Change schema in the databases.
             await self.con.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
             await con_chk.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
@@ -287,6 +357,9 @@ async def test_type_cache_invalidation_in_pool(self):
                     asyncpg.OutdatedSchemaCacheError, ERRNUM):
                 await con1.fetchrow('SELECT * FROM tab1')
 
+            self._check_statements_are_closed(statements1)
+            self._check_statements_are_closed(statements2)
+
             async with con2.transaction():
                 # This should work, as con1 should have invalidated all caches.
                 result = await con2.fetchrow('SELECT * FROM tab1')
@@ -298,10 +371,14 @@ async def test_type_cache_invalidation_in_pool(self):
 
             # Check the invalidation is database-specific, i.e. cache entries
             # for pool_chk/con_chk was not dropped via pool/con1.
+
+            self._check_statements_are_not_closed(statements_chk)
+
             with self.assertRaisesRegex(
                     asyncpg.OutdatedSchemaCacheError, ERRNUM):
                 await con_chk.fetchrow('SELECT * FROM tab1')
 
+            self._check_statements_are_closed(statements_chk)
         finally:
             await self.con.execute('DROP TABLE tab1')
             await self.con.execute('DROP TYPE typ1')