Fix _StatementCache.clear() PS memory leak

un-def · elprans · commit f4aa9bf4a87e · 2019-03-19T12:27:47.000-04:00
See MagicStack#416
diff --git a/asyncpg/connection.py b/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)
diff --git a/tests/test_cache_invalidation.py b/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')

-Original file line number
+Diff line change
 classTestCacheInvalidation(tb.ConnectedTestCase):
++
 +def_get_cached_statements(self, connection=None):
 +ifconnectionisNone:
 +connection=self.con
 +returnlist(connection._stmt_cache.iter_statements())
++
 +def_check_statements_are_not_closed(self, statements):
 +self.assertGreater(len(statements), 0)
 +self.assertTrue(all(nots.closedforsinstatements))
++
 +def_check_statements_are_closed(self, statements):
 +self.assertGreater(len(statements), 0)
 +self.assertTrue(all(s.closedforsinstatements))
++
 asyncdeftest_prepare_cache_invalidation_silent(self):
 awaitself.con.execute('CREATE TABLE tab1(a int, b int)')
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, 2))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 awaitself.con.execute(
 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, '2'))
++
 +self._check_statements_are_closed(statements)
 finally:
 awaitself.con.execute('DROP TABLE tab1')
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, 2))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 awaitself.con.execute(
 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 asyncwithself.con.transaction():
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 +self._check_statements_are_closed(statements)
++
 # This is now OK,
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, '2'))
 result=awaitcon2.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)
++
 awaitself.con.execute(
 'ALTER TABLE tab1 ALTER COLUMN b SET DATA TYPE text')
 result=awaitcon1.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, '2'))
 +self._check_statements_are_closed(statements1)
 +self._check_statements_are_closed(statements2)
++
 asyncwithcon2.transaction():
 # This should work, as con1 should have invalidated
 # the plan cache.
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3)))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 asyncwithself.con.transaction():
 awaitself.con.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
 withself.assertRaisesRegex(
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitself.con.fetchrow('SELECT * FROM tab1')
++
 +self._check_statements_are_closed(statements)
++
 # The second request must be correct (cache was dropped):
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3, None)))
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3)))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 try:
 asyncwithself.con.transaction():
 awaitself.con.execute(
 'ALTER TYPE typ1 ADD ATTRIBUTE c text')
 withself.assertRaisesRegex(
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitself.con.fetchrow('SELECT * FROM tab1')
++
 +self._check_statements_are_closed(statements)
++
 # The second request must be correct (cache was dropped):
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3, None)))
 result=awaitprep.fetchrow()
 self.assertEqual(result, (1, (2, 3)))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 try:
 asyncwithself.con.transaction():
 awaitself.con.execute(
 'ALTER TYPE typ1 ADD ATTRIBUTE c text')
 withself.assertRaisesRegex(
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitprep.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.
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3, 'x')))
 +statements=self._get_cached_statements()
 +self._check_statements_are_not_closed(statements)
++
 awaitself.con.execute('ALTER TYPE typ1 DROP ATTRIBUTE x')
 withself.assertRaisesRegex(
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitself.con.fetchrow('SELECT * FROM tab1')
 +self._check_statements_are_closed(statements)
++
 # This is now OK, the cache is filled after being dropped.
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (3, 'x')))
 result=awaitself.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.
 asyncpg.OutdatedSchemaCacheError, ERRTYP):
 awaitself.con.fetchrow('SELECT * FROM tab1')
 +self._check_statements_are_closed(statements)
++
 # This is now OK, the cache is filled after being dropped.
 result=awaitself.con.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, None)))
 result=awaitcon1.fetchrow('SELECT * FROM tab1')
 self.assertEqual(result, (1, (2, 3)))
 +statements1=self._get_cached_statements(con1)
 +self._check_statements_are_not_closed(statements1)
++
 result=awaitcon2.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=awaitpool_chk.acquire()
 result=awaitcon_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.
 awaitself.con.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
 awaitcon_chk.execute('ALTER TYPE typ1 ADD ATTRIBUTE c text')
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitcon1.fetchrow('SELECT * FROM tab1')
 +self._check_statements_are_closed(statements1)
 +self._check_statements_are_closed(statements2)
++
 asyncwithcon2.transaction():
 # This should work, as con1 should have invalidated all caches.
 result=awaitcon2.fetchrow('SELECT * FROM tab1')
 # 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)
++
 withself.assertRaisesRegex(
 asyncpg.OutdatedSchemaCacheError, ERRNUM):
 awaitcon_chk.fetchrow('SELECT * FROM tab1')
 +self._check_statements_are_closed(statements_chk)
 finally:
 awaitself.con.execute('DROP TABLE tab1')
 awaitself.con.execute('DROP TYPE typ1')