Add new max_cached_statement_use_count parameter to Connection.

The parameter allows asyncpg to refresh cached prepared statements
periodically.

See also issue #76.

Author: 1st1

asyncpg/connection.py

@@ -42,10 +42,12 @@ class Connection(metaclass=ConnectionMeta):
                  '_stmt_cache_max_size', '_stmt_cache', '_stmts_to_close',
                  '_addr', '_opts', '_command_timeout', '_listeners',
                  '_server_version', '_server_caps', '_intro_query',
-                 '_reset_query', '_proxy', '_stmt_exclusive_section')
+                 '_reset_query', '_proxy', '_stmt_exclusive_section',
+                 '_max_cached_statement_use_count')
 
     def __init__(self, protocol, transport, loop, addr, opts, *,
-                 statement_cache_size, command_timeout):
+                 statement_cache_size, command_timeout,
+                 max_cached_statement_use_count):
         self._protocol = protocol
         self._transport = transport
         self._loop = loop
@@ -61,6 +63,7 @@ def __init__(self, protocol, transport, loop, addr, opts, *,
         self._stmt_cache_max_size = statement_cache_size
         self._stmt_cache = collections.OrderedDict()
         self._stmts_to_close = set()
+        self._max_cached_statement_use_count = max_cached_statement_use_count
 
         if command_timeout is not None:
             try:
@@ -263,13 +266,20 @@ async def _get_statement(self, query, timeout, *, named: bool=False):
         use_cache = self._stmt_cache_max_size > 0
         if use_cache:
             try:
-                state = self._stmt_cache[query]
+                holder = self._stmt_cache[query]
             except KeyError:
                 pass
             else:
-                self._stmt_cache.move_to_end(query, last=True)
-                if not state.closed:
-                    return state
+                if holder.use_count < self._max_cached_statement_use_count:
+                    holder.use_count += 1
+
+                    if holder.statement.closed:
+                        self._stmt_cache.pop(query)
+                    else:
+                        self._stmt_cache.move_to_end(query, last=True)
+                        return holder.statement
+                else:
+                    self._stmt_cache.pop(query)
 
         protocol = self._protocol
 
@@ -278,9 +288,9 @@ async def _get_statement(self, query, timeout, *, named: bool=False):
         else:
             stmt_name = ''
 
-        state = await protocol.prepare(stmt_name, query, timeout)
+        statement = await protocol.prepare(stmt_name, query, timeout)
 
-        ready = state._init_types()
+        ready = statement._init_types()
         if ready is not True:
             if self._types_stmt is None:
                 self._types_stmt = await self.prepare(self._intro_query)
@@ -290,16 +300,16 @@ async def _get_statement(self, query, timeout, *, named: bool=False):
 
         if use_cache:
             if len(self._stmt_cache) > self._stmt_cache_max_size - 1:
-                old_query, old_state = self._stmt_cache.popitem(last=False)
-                self._maybe_gc_stmt(old_state)
-            self._stmt_cache[query] = state
+                old_query, old_holder = self._stmt_cache.popitem(last=False)
+                self._maybe_gc_stmt(old_holder.statement)
+            self._stmt_cache[query] = _StatementCacheHolder(statement)
 
         # If we've just created a new statement object, check if there
         # are any statements for GC.
         if self._stmts_to_close:
             await self._cleanup_stmts()
 
-        return state
+        return statement
 
     def cursor(self, query, *args, prefetch=None, timeout=None):
         """Return a *cursor factory* for the specified query.
@@ -465,8 +475,8 @@ def _get_unique_id(self, prefix):
         return '__asyncpg_{}_{}__'.format(prefix, self._uid)
 
     def _close_stmts(self):
-        for stmt in self._stmt_cache.values():
-            stmt.mark_closed()
+        for holder in self._stmt_cache.values():
+            holder.statement.mark_closed()
 
         for stmt in self._stmts_to_close:
             stmt.mark_closed()
@@ -680,6 +690,7 @@ async def connect(dsn=None, *,
                   loop=None,
                   timeout=60,
                   statement_cache_size=100,
+                  max_cached_statement_use_count=100,
                   command_timeout=None,
                   __connection_class__=Connection,
                   **opts):
@@ -715,6 +726,10 @@ async def connect(dsn=None, *,
     :param float timeout: connection timeout in seconds.
 
     :param int statement_cache_size: the size of prepared statement LRU cache.
+                                     Pass ``0`` to disable the cache.
+
+    :param int max_cached_statement_use_count: max number of uses for a cached
+                                               prepared statement.
 
     :param float command_timeout: the default timeout for operations on
                           this connection (the default is no timeout).
@@ -733,6 +748,9 @@ async def connect(dsn=None, *,
         ...     print(types)
         >>> asyncio.get_event_loop().run_until_complete(run())
         [<Record typname='bool' typnamespace=11 ...
+
+    .. versionchanged:: 0.10.0
+       Added ``max_cached_statement_use_count`` parameter.
     """
     if loop is None:
         loop = asyncio.get_event_loop()
@@ -776,13 +794,24 @@ async def connect(dsn=None, *,
         tr.close()
         raise
 
-    con = __connection_class__(pr, tr, loop, addr, opts,
-                               statement_cache_size=statement_cache_size,
-                               command_timeout=command_timeout)
+    con = __connection_class__(
+        pr, tr, loop, addr, opts,
+        statement_cache_size=statement_cache_size,
+        max_cached_statement_use_count=max_cached_statement_use_count,
+        command_timeout=command_timeout)
+
     pr.set_connection(con)
     return con
 
 
+class _StatementCacheHolder:
+    __slots__ = ('statement', 'use_count')
+
+    def __init__(self, statement):
+        self.use_count = 1
+        self.statement = statement
+
+
 class _Atomic:
     __slots__ = ('_acquired',)
 

tests/test_prepare.py

@@ -158,7 +158,7 @@ async def test_prepare_10_stmt_lru(self):
         # At this point our cache should be full.
         self.assertEqual(len(self.con._stmt_cache), cache_max)
         self.assertTrue(
-            all(not s.closed for s in self.con._stmt_cache.values()))
+            all(not s.statement.closed for s in self.con._stmt_cache.values()))
 
         # Since there are references to the statements (`stmts` list),
         # no statements are scheduled to be closed.
@@ -174,7 +174,7 @@ async def test_prepare_10_stmt_lru(self):
         self.assertEqual(len(self.con._stmts_to_close), iter_max - cache_max)
         self.assertTrue(all(s.closed for s in self.con._stmts_to_close))
         self.assertTrue(
-            all(not s.closed for s in self.con._stmt_cache.values()))
+            all(not s.statement.closed for s in self.con._stmt_cache.values()))
 
         zero = await self.con.prepare(query.format(0))
         # Hence, all stale statements should be closed now.
@@ -183,7 +183,7 @@ async def test_prepare_10_stmt_lru(self):
         # The number of cached statements will stay the same though.
         self.assertEqual(len(self.con._stmt_cache), cache_max)
         self.assertTrue(
-            all(not s.closed for s in self.con._stmt_cache.values()))
+            all(not s.statement.closed for s in self.con._stmt_cache.values()))
 
         # After closing all statements will be closed.
         await self.con.close()
@@ -456,3 +456,18 @@ async def check_simple():
         # Check that we can run queries after a failed cursor
         # operation.
         await check_simple()
+
+    async def test_prepare_24_max_use_count(self):
+        self.con._max_cached_statement_use_count = 3
+
+        s = await self.con.prepare('SELECT 1')
+        state = s._state
+
+        s = await self.con.prepare('SELECT 1')
+        self.assertIs(s._state, state)
+
+        s = await self.con.prepare('SELECT 1')
+        self.assertIs(s._state, state)
+
+        s = await self.con.prepare('SELECT 1')
+        self.assertIsNot(s._state, state)