Add max_inactive_connection_lifetime parameter to Pool.

Author: 1st1

asyncpg/_testbase.py

@@ -157,6 +157,7 @@ def create_pool(dsn=None, *,
  min_size=10,
  max_size=10,
  max_queries=50000,
+ max_inactive_connection_lifetime=60.0,
  setup=None,
  init=None,
  loop=None,
@@ -166,6 +167,7 @@ def create_pool(dsn=None, *,
  dsn,
  min_size=min_size, max_size=max_size,
  max_queries=max_queries, loop=loop, setup=setup, init=init,
+ max_inactive_connection_lifetime=max_inactive_connection_lifetime,
  **connect_kwargs)
 
 

asyncpg/pool.py

@@ -86,19 +86,24 @@ class PoolConnectionHolder:
 
  __slots__ = ('_con', '_pool', '_loop',
  '_connect_args', '_connect_kwargs',
- '_max_queries', '_setup', '_init')
+ '_max_queries', '_setup', '_init',
+ '_max_inactive_time', '_in_use',
+ '_inactive_callback')
 
  def __init__(self, pool, *, connect_args, connect_kwargs,
- max_queries, setup, init):
+ max_queries, setup, init, max_inactive_time):
 
  self._pool = pool
  self._con = None
 
  self._connect_args = connect_args
  self._connect_kwargs = connect_kwargs
  self._max_queries = max_queries
+ self._max_inactive_time = max_inactive_time
  self._setup = setup
  self._init = init
+ self._inactive_callback = None
+ self._in_use = False
 
  async def connect(self):
  assert self._con is None
@@ -134,6 +139,8 @@ async def acquire(self) -> PoolConnectionProxy:
  if self._con is None:
  await self.connect()
 
+ self._maybe_cancel_inactive_callback()
+
  proxy = PoolConnectionProxy(self, self._con)
 
  if self._setup is not None:
@@ -154,9 +161,12 @@ async def acquire(self) -> PoolConnectionProxy:
  self._con = None
  raise ex
 
+ self._in_use = True
  return proxy
 
  async def release(self):
+ self._in_use = False
+
  if self._con.is_closed():
  self._con = None
 
@@ -181,7 +191,13 @@ async def release(self):
  self._con = None
  raise ex
 
+ assert self._inactive_callback is None
+ if self._max_inactive_time:
+ self._inactive_callback = self._pool._loop.call_later(
+ self._max_inactive_time, self._deactivate_connection)
+
  async def close(self):
+ self._maybe_cancel_inactive_callback()
  if self._con is None:
  return
  if self._con.is_closed():
@@ -194,6 +210,7 @@ async def close(self):
  self._con = None
 
  def terminate(self):
+ self._maybe_cancel_inactive_callback()
  if self._con is None:
  return
  if self._con.is_closed():
@@ -205,6 +222,18 @@ def terminate(self):
  finally:
  self._con = None
 
+ def _maybe_cancel_inactive_callback(self):
+ if self._inactive_callback is not None:
+ self._inactive_callback.cancel()
+ self._inactive_callback = None
+
+ def _deactivate_connection(self):
+ assert not self._in_use
+ if self._con is None or self._con.is_closed():
+ return
+ self._con.terminate()
+ self._con = None
+
 
 class Pool:
  """A connection pool.
@@ -225,6 +254,7 @@ def __init__(self, *connect_args,
  min_size,
  max_size,
  max_queries,
+ max_inactive_connection_lifetime,
  setup,
  init,
  loop,
@@ -247,6 +277,11 @@ def __init__(self, *connect_args,
  if max_queries <= 0:
  raise ValueError('max_queries is expected to be greater than zero')
 
+ if max_inactive_connection_lifetime < 0:
+ raise ValueError(
+ 'max_inactive_connection_lifetime is expected to be greater '
+ 'or equal to zero')
+
  self._minsize = min_size
  self._maxsize = max_size
 
@@ -265,6 +300,7 @@ def __init__(self, *connect_args,
  connect_args=connect_args,
  connect_kwargs=connect_kwargs,
  max_queries=max_queries,
+ max_inactive_time=max_inactive_connection_lifetime,
  setup=setup,
  init=init)
 
@@ -511,6 +547,7 @@ def create_pool(dsn=None, *,
  min_size=10,
  max_size=10,
  max_queries=50000,
+ max_inactive_connection_lifetime=60.0,
  setup=None,
  init=None,
  loop=None,
@@ -548,6 +585,9 @@ def create_pool(dsn=None, *,
  :param int max_size: Max number of connections in the pool.
  :param int max_queries: Number of queries after a connection is closed
  and replaced with a new connection.
+ :param float max_inactive_connection_lifetime:
+ Number of seconds after which inactive connections in the
+ pool will be closed. Pass ``0`` to disable this mechanism.
  :param coroutine setup: A coroutine to prepare a connection right before
  it is returned from :meth:`~pool.Pool.acquire`.
  An example use case would be to automatically
@@ -567,7 +607,9 @@ def create_pool(dsn=None, *,
  An :exc:`~asyncpg.exceptions.InterfaceError` will be raised on any
  attempted operation on a released connection.
  """
- return Pool(dsn,
- min_size=min_size, max_size=max_size,
- max_queries=max_queries, loop=loop, setup=setup, init=init,
- **connect_kwargs)
+ return Pool(
+ dsn,
+ min_size=min_size, max_size=max_size,
+ max_queries=max_queries, loop=loop, setup=setup, init=init,
+ max_inactive_connection_lifetime=max_inactive_connection_lifetime,
+ **connect_kwargs)

tests/test_pool.py

@@ -11,6 +11,7 @@
 import os
 import platform
 import random
+import time
 import unittest
 
 from asyncpg import _testbase as tb
@@ -457,6 +458,99 @@ async def worker(pool):
  finally:
  await pool.execute('DROP TABLE exmany')
 
+ async def test_pool_max_inactive_time_01(self):
+ async with self.create_pool(
+ database='postgres', min_size=1, max_size=1,
+ max_inactive_connection_lifetime=0.1) as pool:
+
+ # Test that it's OK if a query takes longer time to execute
+ # than `max_inactive_connection_lifetime`.
+
+ con = pool._holders[0]._con
+
+ for _ in range(3):
+ await pool.execute('SELECT pg_sleep(0.5)')
+ self.assertIs(pool._holders[0]._con, con)
+
+ self.assertEqual(
+ await pool.execute('SELECT 1::int'),
+ 'SELECT 1')
+ self.assertIs(pool._holders[0]._con, con)
+
+ async def test_pool_max_inactive_time_02(self):
+ async with self.create_pool(
+ database='postgres', min_size=1, max_size=1,
+ max_inactive_connection_lifetime=0.5) as pool:
+
+ # Test that we have a new connection after pool not
+ # being used longer than `max_inactive_connection_lifetime`.
+
+ con = pool._holders[0]._con
+
+ self.assertEqual(
+ await pool.execute('SELECT 1::int'),
+ 'SELECT 1')
+ self.assertIs(pool._holders[0]._con, con)
+
+ await asyncio.sleep(1, loop=self.loop)
+ self.assertIs(pool._holders[0]._con, None)
+
+ self.assertEqual(
+ await pool.execute('SELECT 1::int'),
+ 'SELECT 1')
+ self.assertIsNot(pool._holders[0]._con, con)
+
+ async def test_pool_max_inactive_time_03(self):
+ async with self.create_pool(
+ database='postgres', min_size=1, max_size=1,
+ max_inactive_connection_lifetime=1) as pool:
+
+ # Test that we start counting inactive time *after*
+ # the connection is being released back to the pool.
+
+ con = pool._holders[0]._con
+
+ await pool.execute('SELECT pg_sleep(0.5)')
+ await asyncio.sleep(0.6, loop=self.loop)
+
+ self.assertIs(pool._holders[0]._con, con)
+
+ self.assertEqual(
+ await pool.execute('SELECT 1::int'),
+ 'SELECT 1')
+ self.assertIs(pool._holders[0]._con, con)
+
+ async def test_pool_max_inactive_time_04(self):
+ # Chaos test for max_inactive_connection_lifetime.
+ DURATION = 2.0
+ START = time.monotonic()
+ N = 0
+
+ async def worker(pool):
+ nonlocal N
+ await asyncio.sleep(random.random() / 10 + 0.1, loop=self.loop)
+ async with pool.acquire() as con:
+ if random.random() > 0.5:
+ await con.execute('SELECT pg_sleep({:.2f})'.format(
+ random.random() / 10))
+ self.assertEqual(
+ await con.fetchval('SELECT 42::int'),
+ 42)
+
+ if time.monotonic() - START < DURATION:
+ await worker(pool)
+
+ N += 1
+
+ async with self.create_pool(
+ database='postgres', min_size=10, max_size=30,
+ max_inactive_connection_lifetime=0.1) as pool:
+
+ workers = [worker(pool) for _ in range(50)]
+ await asyncio.gather(*workers, loop=self.loop)
+
+ self.assertGreater(N, 50)
+
 
 @unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
 class TestHostStandby(tb.ConnectedTestCase):