diff --git a/asyncpg/_testbase/__init__.py b/asyncpg/_testbase/__init__.py index 16238ce1..9dd76f84 100644 --- a/asyncpg/_testbase/__init__.py +++ b/asyncpg/_testbase/__init__.py @@ -267,6 +267,7 @@ def _shutdown_cluster(cluster): def create_pool(dsn=None, *, + init_size=None, min_size=10, max_size=10, max_queries=50000, @@ -281,6 +282,7 @@ def create_pool(dsn=None, *, **connect_kwargs): return pool_class( dsn, + init_size=init_size, min_size=min_size, max_size=max_size, max_queries=max_queries, @@ -366,10 +368,16 @@ def setUp(self): self._pools = [] def tearDown(self): - super().tearDown() + maintenance_tasks = [] for pool in self._pools: pool.terminate() + if pool._maintenance_task is not None: + maintenance_tasks.append(pool._maintenance_task) + if maintenance_tasks: + self.loop.run_until_complete(asyncio.gather( + *maintenance_tasks, return_exceptions=True)) self._pools = [] + super().tearDown() def create_pool(self, pool_class=pg_pool.Pool, connection_class=pg_connection.Connection, **kwargs): diff --git a/asyncpg/connect_utils.py b/asyncpg/connect_utils.py index 528739d0..0c5801d7 100644 --- a/asyncpg/connect_utils.py +++ b/asyncpg/connect_utils.py @@ -1158,7 +1158,16 @@ async def __connect_addr( else: connector = loop.create_connection(proto_factory, *addr) - tr, pr = await connector + try: + tr, pr = await connector + except (Exception, asyncio.CancelledError): + # The protocol can exist before create_connection() returns. If + # that operation is cancelled, nobody will await authentication. + if not connected.done(): + connected.cancel() + elif not connected.cancelled(): + connected.exception() + raise try: await connected diff --git a/asyncpg/connection.py b/asyncpg/connection.py index 4a72307b..c4067fc7 100644 --- a/asyncpg/connection.py +++ b/asyncpg/connection.py @@ -54,6 +54,7 @@ class Connection(metaclass=ConnectionMeta): '_intro_query', '_reset_query', '_proxy', '_stmt_exclusive_section', '_config', '_params', '_addr', '_log_listeners', '_termination_listeners', '_cancellations', + '_pool_holder', '_source_traceback', '_query_loggers', '__weakref__') def __init__(self, protocol, transport, loop, @@ -105,6 +106,7 @@ def __init__(self, protocol, transport, loop, self._reset_query = None self._proxy = None + self._pool_holder = None # Used to serialize operations that might involve anonymous # statements. Specifically, we want to make the following @@ -134,6 +136,10 @@ def __del__(self): warnings.warn(msg, ResourceWarning) if not self._loop.is_closed(): + # A weak holder reference may still be live during GC. + # Finalization must not notify the pool and restart + # maintenance. + self._pool_holder = None self.terminate() async def add_listener(self, channel, callback): @@ -1599,10 +1605,11 @@ def _cleanup(self): # Free the resources associated with this connection. # This must be called when a connection is terminated. - if self._proxy is not None: - # Connection is a member of a pool, so let the pool - # know that this connection is dead. - self._proxy._holder._release_on_close() + if self._pool_holder is not None: + # Idle connections have no proxy, but still belong to a holder. + holder, self._pool_holder = self._pool_holder(), None + if holder is not None and holder._con is self: + holder._release_on_close() self._mark_stmts_as_closed() self._listeners.clear() diff --git a/asyncpg/pool.py b/asyncpg/pool.py index 81c2ad3b..29e9cb9d 100644 --- a/asyncpg/pool.py +++ b/asyncpg/pool.py @@ -14,6 +14,7 @@ from types import TracebackType from typing import Any, Optional, Type import warnings +import weakref from . import compat from . import connection @@ -24,6 +25,24 @@ logger = logging.getLogger(__name__) +class _PoolConnectionHolderQueue(asyncio.LifoQueue): + """Prefer live connections, with LIFO ordering within each group.""" + + def _get(self): + for i in range(len(self._queue) - 1, -1, -1): + if self._queue[i].is_connected(): + return self._queue.pop(i) + return super()._get() + + def get_disconnected_nowait(self): + for i in range(len(self._queue) - 1, -1, -1): + if not self._queue[i].is_connected(): + holder = self._queue.pop(i) + self._wakeup_next(self._putters) + return holder + raise asyncio.QueueEmpty + + class PoolConnectionProxyMeta(type): def __new__( @@ -114,7 +133,7 @@ class PoolConnectionHolder: '_max_queries', '_setup', '_max_inactive_time', '_in_use', '_inactive_callback', '_timeout', - '_generation') + '_generation', '__weakref__') def __init__( self, @@ -149,8 +168,18 @@ async def connect(self) -> None: 'PoolConnectionHolder.connect() called while another ' 'connection already exists') - self._con = await self._pool._get_new_connection() - self._generation = self._pool._generation + generation = self._pool._generation + con = await self._pool._get_new_connection() + if self._pool.is_closing(): + await con.close() + raise exceptions.InterfaceError('pool is closing') + if con.is_closed(): + raise exceptions.ConnectionDoesNotExistError( + 'connection was closed during pool initialization') + self._con = con + # Notify live holders without keeping an abandoned pool alive. + con._pool_holder = weakref.ref(self) + self._generation = generation self._maybe_cancel_inactive_callback() self._setup_inactive_callback() @@ -299,19 +328,23 @@ def _deactivate_inactive_connection(self) -> None: raise exceptions.InternalClientError( 'attempting to deactivate an acquired connection') + self._inactive_callback = None if self._con is not None: + if (self.is_connected() and not self._pool.is_closing() and + self._pool.get_size() <= self._pool.get_min_size()): + # A floor connection needs no further inactivity checks. + # Acquiring and releasing it will arm a new timer. + return + # The connection is idle and not in use, so it's fine to # use terminate() instead of close(). self._con.terminate() - # Must call clear_connection, because _deactivate_connection - # is called when the connection is *not* checked out, and - # so terminate() above will not call the below. - self._release_on_close() def _release_on_close(self) -> None: self._maybe_cancel_inactive_callback() self._release() self._con = None + self._pool._schedule_min_size_maintenance() def _release(self) -> None: """Release this connection holder.""" @@ -345,14 +378,16 @@ class Pool: """ __slots__ = ( - '_queue', '_loop', '_minsize', '_maxsize', + '_queue', '_loop', '_initsize', '_minsize', '_maxsize', '_init', '_connect', '_reset', '_connect_args', '_connect_kwargs', '_holders', '_initialized', '_initializing', '_closing', '_closed', '_connection_class', '_record_class', '_generation', - '_setup', '_max_queries', '_max_inactive_connection_lifetime' + '_setup', '_max_queries', '_max_inactive_connection_lifetime', + '_maintenance_task', ) def __init__(self, *connect_args, + init_size=None, min_size, max_size, max_queries, @@ -388,6 +423,19 @@ def __init__(self, *connect_args, if min_size > max_size: raise ValueError('min_size is greater than max_size') + if init_size is None: + init_size = min_size + + if init_size < 0: + raise ValueError( + 'init_size is expected to be greater or equal to zero') + + if init_size > max_size: + raise ValueError('init_size is greater than max_size') + + if init_size < min_size: + raise ValueError('init_size is smaller than min_size') + if max_queries <= 0: raise ValueError('max_queries is expected to be greater than zero') @@ -406,6 +454,7 @@ def __init__(self, *connect_args, 'record_class is expected to be a subclass of ' 'asyncpg.Record, got {!r}'.format(record_class)) + self._initsize = init_size self._minsize = min_size self._maxsize = max_size @@ -420,6 +469,7 @@ def __init__(self, *connect_args, self._closing = False self._closed = False self._generation = 0 + self._maintenance_task = None self._connect = connect if connect is not None else connection.connect self._connect_args = connect_args @@ -445,12 +495,19 @@ async def _async__init__(self): try: await self._initialize() return self + except (Exception, asyncio.CancelledError): + # Failed initialization must not leave warm connections behind. + self._closed = True + for holder in self._holders: + holder.terminate() + raise finally: self._initializing = False self._initialized = True + self._schedule_min_size_maintenance() async def _initialize(self): - self._queue = asyncio.LifoQueue(maxsize=self._maxsize) + self._queue = _PoolConnectionHolderQueue(maxsize=self._maxsize) for _ in range(self._maxsize): ch = PoolConnectionHolder( self, @@ -461,7 +518,7 @@ async def _initialize(self): self._holders.append(ch) self._queue.put_nowait(ch) - if self._minsize: + if self._initsize: # Since we use a LIFO queue, the first items in the queue will be # the last ones in `self._holders`. We want to pre-connect the # first few connections in the queue, therefore we want to walk @@ -472,16 +529,58 @@ async def _initialize(self): first_ch = self._holders[-1] # type: PoolConnectionHolder await first_ch.connect() - if self._minsize > 1: + if self._initsize > 1: connect_tasks = [] for i, ch in enumerate(reversed(self._holders[:-1])): - # `minsize - 1` because we already have first_ch - if i >= self._minsize - 1: + # `initsize - 1` because we already have first_ch + if i >= self._initsize - 1: break connect_tasks.append(ch.connect()) await asyncio.gather(*connect_tasks) + def _schedule_min_size_maintenance(self): + if (not self._initialized or self._initializing or self.is_closing() + or not self._minsize or self._maintenance_task is not None): + return + if self.get_size() < self._minsize: + self._maintenance_task = self._loop.create_task( + self._maintain_min_size()) + + async def _maintain_min_size(self): + retry_delay = 1.0 + try: + while not self.is_closing() and self.get_size() < self._minsize: + try: + holder = self._queue.get_disconnected_nowait() + except asyncio.QueueEmpty: + # Acquirers are already connecting the remaining holders. + return + + failed = False + try: + if holder._con is not None: + holder.terminate() + await holder.connect() + except asyncio.CancelledError: + raise + except Exception: + if self.is_closing(): + return + failed = True + logger.warning('Failed to restore the pool connection ' + 'floor; retrying', exc_info=True) + finally: + self._queue.put_nowait(holder) + + if failed: + await asyncio.sleep(retry_delay) + retry_delay = min(retry_delay * 2, 60.0) + else: + retry_delay = 1.0 + finally: + self._maintenance_task = None + def is_closing(self): """Return ``True`` if the pool is closing or is closed. @@ -496,10 +595,21 @@ def get_size(self): """ return sum(h.is_connected() for h in self._holders) + def get_init_size(self): + """Return the initial number of connections in this pool. + + .. versionadded:: 0.32.0 + """ + return self._initsize + def get_min_size(self): """Return the minimum number of connections in this pool. .. versionadded:: 0.25.0 + + .. versionchanged:: 0.32.0 + The parameter now controls the connection floor rather than the + initial pool size (see ``init_size``). """ return self._minsize @@ -888,6 +998,7 @@ async def _acquire_impl(): proxy = await ch.acquire() # type: PoolConnectionProxy except (Exception, asyncio.CancelledError): self._queue.put_nowait(ch) + self._schedule_min_size_maintenance() raise else: # Record the timeout, as we will apply it by default @@ -967,6 +1078,12 @@ async def close(self): warning_callback = None try: + if self._maintenance_task is not None: + self._maintenance_task.cancel() + await asyncio.gather( + self._maintenance_task, return_exceptions=True) + self._maintenance_task = None + warning_callback = self._loop.call_later( 60, self._warn_on_long_close) @@ -999,9 +1116,11 @@ def terminate(self): if self._closed: return self._check_init() + self._closed = True + if self._maintenance_task is not None: + self._maintenance_task.cancel() for ch in self._holders: ch.terminate() - self._closed = True async def expire_connections(self): """Expire all currently open connections. @@ -1080,6 +1199,7 @@ def __await__(self): def create_pool(dsn=None, *, + init_size=None, min_size=10, max_size=10, max_queries=50000, @@ -1154,8 +1274,15 @@ def create_pool(dsn=None, *, the connections in this pool. Must be a subclass of :class:`~asyncpg.Record`. + :param int init_size: + Number of connections the pool will be initialized with. Defaults + to *min_size*. Must be between *min_size* and *max_size*. + :param int min_size: - Number of connection the pool will be initialized with. + Minimum number of connections retained during idle periods. Closed + connections are replaced in the background to restore this floor. + The pool may temporarily fall below it while reconnecting or while + the server is unavailable. Pass ``0`` to allow the pool to drain. :param int max_size: Max number of connections in the pool. @@ -1166,7 +1293,8 @@ def create_pool(dsn=None, *, :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. + pool above *min_size* will be closed. Pass ``0`` to disable this + mechanism. :param coroutine connect: A coroutine that is called instead of @@ -1237,11 +1365,17 @@ def create_pool(dsn=None, *, .. versionchanged:: 0.30.0 Added the *connect* and *reset* parameters. + + .. versionchanged:: 0.32.0 + The *min_size* parameter now defines the connection floor. The former + role of *min_size* — setting the initial pool size — is now handled by + the new *init_size* parameter, which defaults to *min_size*. """ return Pool( dsn, connection_class=connection_class, record_class=record_class, + init_size=init_size, min_size=min_size, max_size=max_size, max_queries=max_queries, diff --git a/tests/test_adversity.py b/tests/test_adversity.py index 1d061ee6..0ef6c7e3 100644 --- a/tests/test_adversity.py +++ b/tests/test_adversity.py @@ -51,7 +51,7 @@ async def test_connection_close_timeout(self): @tb.with_timeout(30.0) async def test_pool_acquire_timeout(self): pool = await self.create_pool( - database='postgres', min_size=2, max_size=2) + database='postgres', init_size=2, min_size=0, max_size=2) try: self.proxy.trigger_connectivity_loss() for _ in range(2): @@ -68,7 +68,7 @@ async def test_pool_acquire_timeout(self): @tb.with_timeout(30.0) async def test_pool_release_timeout(self): pool = await self.create_pool( - database='postgres', min_size=2, max_size=2) + database='postgres', init_size=2, min_size=0, max_size=2) try: with self.assertRaises(asyncio.TimeoutError): async with pool.acquire(timeout=0.5): @@ -96,7 +96,8 @@ def kill_connectivity(): self.proxy.trigger_connectivity_loss() new_pool = self.create_pool( - database='postgres', min_size=pool_size, max_size=pool_size, + database='postgres', + init_size=pool_size, min_size=0, max_size=pool_size, timeout=cmd_timeout, command_timeout=cmd_timeout) with self.assertRunUnder(worst_runtime): diff --git a/tests/test_cache_invalidation.py b/tests/test_cache_invalidation.py index 5cab2d92..f2901c9e 100644 --- a/tests/test_cache_invalidation.py +++ b/tests/test_cache_invalidation.py @@ -77,7 +77,7 @@ async def test_prepare_cache_invalidation_in_transaction(self): async def test_prepare_cache_invalidation_in_pool(self): pool = await self.create_pool(database='postgres', - min_size=2, max_size=2) + init_size=2, min_size=0, max_size=2) await self.con.execute('CREATE TABLE tab1(a int, b int)') @@ -309,10 +309,10 @@ async def test_type_cache_invalidation_on_change_attr(self): async def test_type_cache_invalidation_in_pool(self): await self.con.execute('CREATE DATABASE testdb') pool = await self.create_pool(database='postgres', - min_size=2, max_size=2) + init_size=2, min_size=0, max_size=2) pool_chk = await self.create_pool(database='testdb', - min_size=2, max_size=2) + init_size=2, min_size=0, max_size=2) await self.con.execute('CREATE TYPE typ1 AS (x int, y int)') await self.con.execute('CREATE TABLE tab1(a int, b typ1)') diff --git a/tests/test_connect.py b/tests/test_connect.py index 66341f40..fbac0600 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -1671,6 +1671,28 @@ async def test_connect_args_validation(self): class TestConnection(tb.ConnectedTestCase): + async def test_connection_cancelled_during_transport_setup(self): + for connection_lost_first in (False, True): + with self.subTest(connection_lost_first=connection_lost_first): + async def interrupted_connector(factory, *args, **kwargs): + proto = factory() + if connection_lost_first: + proto.connection_lost(None) + else: + self.loop.call_soon(proto.connection_lost, None) + raise asyncio.CancelledError + + with unittest.mock.patch.object( + self.loop, 'create_connection', interrupted_connector, + ): + with self.assertRaises(asyncio.CancelledError): + await self.connect(host='127.0.0.1', ssl=False) + + # Authentication must not report an unobserved failure after + # cancellation, even if the transport closes afterwards. + await asyncio.sleep(0) + gc.collect() + async def test_connection_isinstance(self): self.assertTrue(isinstance(self.con, pg_connection.Connection)) self.assertTrue(isinstance(self.con, object)) diff --git a/tests/test_pool.py b/tests/test_pool.py index f8000a59..bc6180e0 100644 --- a/tests/test_pool.py +++ b/tests/test_pool.py @@ -6,6 +6,7 @@ import asyncio +import gc import inspect import os import pathlib @@ -14,6 +15,7 @@ import textwrap import time import unittest +import weakref from unittest import mock import asyncpg @@ -45,11 +47,20 @@ async def _cancel(self, waiter, cancel_waiter=None): class TestPool(tb.ConnectedTestCase): + async def wait_for_pool_size(self, pool, size): + async def wait(): + while pool.get_size() != size: + await asyncio.sleep(0.01) + + await asyncio.wait_for(wait(), 5) + async def test_pool_01(self): for n in {1, 5, 10, 20, 100}: with self.subTest(tasksnum=n): pool = await self.create_pool(database='postgres', - min_size=5, max_size=10) + init_size=5, + min_size=1, + max_size=10) async def worker(): con = await pool.acquire() @@ -64,7 +75,9 @@ async def test_pool_02(self): for n in {1, 3, 5, 10, 20, 100}: with self.subTest(tasksnum=n): async with self.create_pool(database='postgres', - min_size=5, max_size=5) as pool: + init_size=5, + min_size=1, + max_size=5) as pool: async def worker(): con = await pool.acquire(timeout=5) @@ -76,7 +89,7 @@ async def worker(): async def test_pool_03(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) con = await pool.acquire(timeout=1) with self.assertRaises(asyncio.TimeoutError): @@ -87,7 +100,7 @@ async def test_pool_03(self): async def test_pool_04(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) con = await pool.acquire(timeout=POOL_NOMINAL_TIMEOUT) @@ -112,7 +125,9 @@ async def test_pool_05(self): for n in {1, 3, 5, 10, 20, 100}: with self.subTest(tasksnum=n): pool = await self.create_pool(database='postgres', - min_size=5, max_size=10) + init_size=5, + min_size=1, + max_size=10) async def worker(): async with pool.acquire() as con: @@ -129,7 +144,7 @@ async def setup(con): fut.set_result(con) async with self.create_pool(database='postgres', - min_size=5, max_size=5, + init_size=5, min_size=1, max_size=5, setup=setup) as pool: async with pool.acquire() as con: pass @@ -171,8 +186,8 @@ async def user(pool): raise RuntimeError('init was not called') async with self.create_pool(database='postgres', - min_size=2, - max_size=5, + init_size=2, + min_size=1, max_size=5, connect=connect, init=init, setup=setup, @@ -198,7 +213,7 @@ async def bad_connect(*args, **kwargs): async def test_pool_08(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) con = await pool.acquire(timeout=POOL_NOMINAL_TIMEOUT) with self.assertRaisesRegex(asyncpg.InterfaceError, 'is not a member'): @@ -206,10 +221,10 @@ async def test_pool_08(self): async def test_pool_09(self): pool1 = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) pool2 = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) try: con = await pool1.acquire(timeout=POOL_NOMINAL_TIMEOUT) @@ -224,7 +239,7 @@ async def test_pool_09(self): async def test_pool_10(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) con = await pool.acquire() await pool.release(con) @@ -234,7 +249,7 @@ async def test_pool_10(self): async def test_pool_11(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) async with pool.acquire() as con: self.assertIn(repr(con._con), repr(con)) # Test __repr__. @@ -291,7 +306,7 @@ async def test_pool_11(self): async def test_pool_12(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) async with pool.acquire() as con: self.assertTrue(isinstance(con, pg_connection.Connection)) @@ -301,7 +316,7 @@ async def test_pool_12(self): async def test_pool_13(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) async with pool.acquire() as con: self.assertIn('Execute an SQL command', con.execute.__doc__) @@ -337,7 +352,7 @@ async def setup(con): last_con = None cons = [] async with self.create_pool(database='postgres', - min_size=1, max_size=1, + init_size=1, min_size=1, max_size=1, setup=setup) as pool: with self.assertRaises(Error): await pool.acquire() @@ -351,7 +366,7 @@ async def setup(con): last_con = None cons = [] async with self.create_pool(database='postgres', - min_size=0, max_size=1, + init_size=0, min_size=0, max_size=1, init=setup) as pool: with self.assertRaises(Error): await pool.acquire() @@ -393,7 +408,7 @@ async def test_pool_auth(self): pool = await self.create_pool(database='postgres', user='pooluser', password='poolpassword', - min_size=5, max_size=10) + init_size=5, min_size=1, max_size=10) async def worker(): con = await pool.acquire() @@ -414,7 +429,7 @@ async def worker(): async def test_pool_handles_task_cancel_in_acquire_with_timeout(self): # See https://github.com/MagicStack/asyncpg/issues/547 pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) async def worker(): async with pool.acquire(timeout=100): @@ -435,7 +450,7 @@ async def test_pool_handles_task_cancel_in_release(self): # Use SlowResetConnectionPool to simulate # the Task.cancel() and __aexit__ race. pool = await self.create_pool(database='postgres', - min_size=1, max_size=1, + init_size=1, min_size=1, max_size=1, connection_class=SlowResetConnection) async def worker(): @@ -456,7 +471,7 @@ async def test_pool_handles_query_cancel_in_release(self): # Use SlowResetConnectionPool to simulate # the Task.cancel() and __aexit__ race. pool = await self.create_pool(database='postgres', - min_size=1, max_size=1, + init_size=1, min_size=1, max_size=1, connection_class=SlowCancelConnection) async def worker(): @@ -516,7 +531,7 @@ async def cancel(**kwargs): async def test_pool_no_acquire_deadlock(self): async with self.create_pool(database='postgres', - min_size=1, max_size=1, + init_size=1, min_size=1, max_size=1, max_queries=1) as pool: async def sleep_and_release(): @@ -550,7 +565,7 @@ async def test(pool): cons.add(con) async with self.create_pool( - database='postgres', min_size=10, max_size=10, + database='postgres', init_size=10, min_size=1, max_size=10, max_queries=1, connection_class=MyConnection, statement_cache_size=3) as pool: @@ -561,7 +576,9 @@ async def test(pool): async def test_pool_release_in_xact(self): """Test that Connection.reset() closes any open transaction.""" async with self.create_pool(database='postgres', - min_size=1, max_size=1) as pool: + init_size=1, + min_size=1, + max_size=1) as pool: async def get_xact_id(con): return await con.fetchval('select txid_current()') @@ -624,7 +641,9 @@ async def test_execute_with_arg(pool): async def run(N, meth): async with self.create_pool(database='postgres', - min_size=5, max_size=10) as pool: + init_size=5, + min_size=0, + max_size=10) as pool: coros = [meth(pool) for _ in range(N)] res = await asyncio.gather(*coros) @@ -651,7 +670,9 @@ async def worker(pool): N = 200 async with self.create_pool(database='postgres', - min_size=5, max_size=10) as pool: + init_size=5, + min_size=0, + max_size=10) as pool: await pool.execute('CREATE TABLE exmany (a text, b int)') try: @@ -668,7 +689,7 @@ async def worker(pool): async def test_pool_max_inactive_time_01(self): async with self.create_pool( - database='postgres', min_size=1, max_size=1, + database='postgres', init_size=1, min_size=0, max_size=1, max_inactive_connection_lifetime=0.1) as pool: # Test that it's OK if a query takes longer time to execute @@ -687,7 +708,7 @@ async def test_pool_max_inactive_time_01(self): async def test_pool_max_inactive_time_02(self): async with self.create_pool( - database='postgres', min_size=1, max_size=1, + database='postgres', init_size=1, min_size=0, max_size=1, max_inactive_connection_lifetime=0.5) as pool: # Test that we have a new connection after pool not @@ -710,7 +731,7 @@ async def test_pool_max_inactive_time_02(self): async def test_pool_max_inactive_time_03(self): async with self.create_pool( - database='postgres', min_size=1, max_size=1, + database='postgres', init_size=1, min_size=0, max_size=1, max_inactive_connection_lifetime=1) as pool: # Test that we start counting inactive time *after* @@ -751,7 +772,7 @@ async def worker(pool): N += 1 async with self.create_pool( - database='postgres', min_size=10, max_size=30, + database='postgres', init_size=10, min_size=0, max_size=30, max_inactive_connection_lifetime=0.1) as pool: workers = [worker(pool) for _ in range(50)] @@ -770,7 +791,7 @@ async def init(con): con.add_termination_listener(lambda _: terminated.set()) async with self.create_pool( - database='postgres', min_size=2, max_size=2, + database='postgres', init_size=2, min_size=0, max_size=2, init=init, max_inactive_connection_lifetime=0.2) as pool: @@ -788,9 +809,386 @@ async def init(con): for holder in pool._holders: self.assertIsNone(holder._con) + async def test_pool_min_size_keeps_connections_alive(self): + # Test that min_size prevents idle connections from being closed. + async with self.create_pool( + database='postgres', init_size=2, min_size=2, max_size=2, + max_inactive_connection_lifetime=0.2) as pool: + + con0 = pool._holders[0]._con + con1 = pool._holders[1]._con + self.assertIsNotNone(con0) + self.assertIsNotNone(con1) + + await asyncio.sleep(0.5) + + # Connections should be kept alive because pool size == min_size. + self.assertIs(pool._holders[0]._con, con0) + self.assertIs(pool._holders[1]._con, con1) + + async def test_pool_min_size_partial_keep(self): + # When the pool has more connections than min_size, only the excess + # connections should be allowed to expire; the min_size ones are kept. + async with self.create_pool( + database='postgres', init_size=3, min_size=1, max_size=3, + max_inactive_connection_lifetime=0.2) as pool: + + # Exercise release timers for all three connections. + c1 = await pool.acquire() + c2 = await pool.acquire() + c3 = await pool.acquire() + await pool.release(c1) + await pool.release(c2) + await pool.release(c3) + + self.assertEqual(pool.get_size(), 3) + + await asyncio.sleep(0.5) + + # Only min_size (1) connection should survive. + self.assertEqual(pool.get_size(), 1) + + async def test_pool_implicit_init_size(self): + # Existing small, large, and lazy pool configurations still work. + for size in (0, 1, 20): + with self.subTest(size=size): + async with asyncpg.create_pool( + **self.get_connection_spec(), + min_size=size, max_size=max(1, size), + ) as pool: + self.assertEqual(pool.get_init_size(), size) + self.assertEqual(pool.get_size(), size) + + async def test_pool_min_size_reuses_retained_connections(self): + async with self.create_pool( + init_size=3, min_size=1, max_size=3, + max_inactive_connection_lifetime=0.05, + ) as pool: + await self.wait_for_pool_size(pool, 1) + retained = next(h._con for h in pool._holders if h.is_connected()) + + for _ in range(2): + async with pool.acquire() as con: + self.assertIs(con._con, retained) + self.assertEqual(await con.fetchval('SELECT 42'), 42) + await asyncio.sleep(0.1) + self.assertEqual(pool.get_size(), 1) + + async def test_pool_shutdown_cleans_idle_connections(self): + class WeakPool(pg_pool.Pool): + pass + + for action in ('close', 'terminate'): + for floor in (0, 1): + with self.subTest(action=action, floor=floor): + pool = await tb.create_pool( + **self.get_connection_spec(), pool_class=WeakPool, + init_size=1, min_size=floor, max_size=1, + max_inactive_connection_lifetime=0.05, + ) + try: + if action == 'close': + await pool.close() + else: + pool.terminate() + self.assertIsNone(pool._holders[0]._con) + self.assertIsNone(pool._holders[0]._inactive_callback) + finally: + pool.terminate() + + ref = weakref.ref(pool) + del pool + gc.collect() + self.assertIsNone(ref()) + + async def test_pool_gc_does_not_restore_min_size(self): + for inactive_timeout in (0, 0.05): + with self.subTest(inactive_timeout=inactive_timeout): + connections = [] + maintenance_pools = [] + + class WeakPool(pg_pool.Pool): + async def _maintain_min_size(self): + maintenance_pools.append(self) + await super()._maintain_min_size() + + async def connect(*args, **kwargs): + con = await pg_connection.connect(*args, **kwargs) + connections.append(weakref.ref(con)) + return con + + pool = await tb.create_pool( + **self.get_connection_spec(), pool_class=WeakPool, + init_size=2, min_size=2, max_size=2, connect=connect, + max_inactive_connection_lifetime=inactive_timeout, + ) + ref = weakref.ref(pool) + try: + if inactive_timeout: + await asyncio.sleep(inactive_timeout * 2) + self.assertTrue(all(h._inactive_callback is None + for h in pool._holders)) + del pool + with self.assertWarnsRegex(ResourceWarning, + 'unclosed connection'): + gc.collect() + + # Give any incorrectly scheduled maintenance task time + # to reconnect, then ensure no new connections appeared. + await asyncio.sleep(0.1) + self.assertFalse(maintenance_pools) + self.assertEqual(len(connections), 2) + self.assertIsNone(ref()) + self.assertTrue(all(con() is None for con in connections)) + finally: + remaining = ref() + if remaining is not None: + await remaining.close() + for remaining in maintenance_pools: + await remaining.close() + maintenance_pools.clear() + + async def test_pool_min_size_restored_after_recycling(self): + for cause in ('max_queries', 'expire', 'close', 'terminate'): + with self.subTest(cause=cause): + async with self.create_pool( + init_size=1, min_size=1, max_size=1, + max_queries=1 if cause == 'max_queries' else 50000, + max_inactive_connection_lifetime=0, + ) as pool: + async with pool.acquire() as con: + old_con = con._con + await con.fetchval('SELECT 42') + if cause == 'expire': + await pool.expire_connections() + elif cause == 'close': + await con.close() + elif cause == 'terminate': + con.terminate() + + await self.wait_for_pool_size(pool, 1) + self.assertIsNot(pool._holders[0]._con, old_con) + self.assertTrue(old_con.is_closed()) + + async def test_pool_min_size_restored_after_idle_connection_loss(self): + async with self.create_pool( + init_size=1, min_size=1, max_size=1, + max_inactive_connection_lifetime=0, + ) as pool: + old_con = pool._holders[0]._con + terminated = asyncio.Event() + old_con.add_termination_listener(lambda con: terminated.set()) + + await self.con.execute( + 'SELECT pg_terminate_backend($1)', old_con.get_server_pid()) + await asyncio.wait_for(terminated.wait(), 5) + await self.wait_for_pool_size(pool, 1) + self.assertIsNot(pool._holders[0]._con, old_con) + self.assertEqual(await pool.fetchval('SELECT 42'), 42) + + async def test_pool_min_size_retries_failed_reconnect(self): + offline = False + attempted = asyncio.Event() + + async def connect(*args, **kwargs): + if offline: + attempted.set() + raise OSError('server temporarily unavailable') + return await pg_connection.connect(*args, **kwargs) + + async with self.create_pool( + init_size=1, min_size=1, max_size=1, connect=connect, + ) as pool: + offline = True + with self.assertLogs('asyncpg.pool', level='WARNING'): + pool._holders[0]._con.terminate() + await asyncio.wait_for(attempted.wait(), 5) + self.assertEqual(pool.get_size(), 0) + + offline = False + await self.wait_for_pool_size(pool, 1) + self.assertEqual(await pool.fetchval('SELECT 42'), 42) + + async def test_pool_floor_reconnect_honors_expired_generation(self): + started = asyncio.Event() + resume = asyncio.Event() + attempts = 0 + + async def init(con): + nonlocal attempts + attempts += 1 + if attempts == 2: + started.set() + await resume.wait() + + async with self.create_pool( + init_size=1, min_size=1, max_size=1, init=init, + server_settings={'application_name': 'old_pool_args'}, + ) as pool: + pool._holders[0]._con.terminate() + await asyncio.wait_for(started.wait(), 5) + pool.set_connect_args(**self.get_connection_spec({ + 'server_settings': {'application_name': 'new_pool_args'}, + })) + await pool.expire_connections() + resume.set() + await self.wait_for_pool_size(pool, 1) + + async with pool.acquire() as con: + self.assertEqual(con.get_settings().application_name, + 'new_pool_args') + + async def test_pool_min_size_retries_closed_connection(self): + offline = False + attempted = asyncio.Event() + + async def init(con): + if offline: + await con.close() + attempted.set() + + async with self.create_pool( + init_size=1, min_size=1, max_size=1, init=init, + ) as pool: + offline = True + with self.assertLogs('asyncpg.pool', level='WARNING'): + pool._holders[0]._con.terminate() + await asyncio.wait_for(attempted.wait(), 5) + self.assertEqual(pool.get_size(), 0) + + offline = False + await self.wait_for_pool_size(pool, 1) + self.assertEqual(await pool.fetchval('SELECT 42'), 42) + + async def test_pool_min_size_restored_after_multiple_losses(self): + async with self.create_pool( + init_size=3, min_size=3, max_size=5, + ) as pool: + old_connections = { + h._con for h in pool._holders if h.is_connected() + } + for con in old_connections: + con.terminate() + + async def worker(): + async with pool.acquire() as con: + self.assertLessEqual(pool.get_size(), pool.get_max_size()) + self.assertEqual(await con.fetchval('SELECT 42'), 42) + + await asyncio.gather(*(worker() for _ in range(10))) + self.assertGreaterEqual(pool.get_size(), pool.get_min_size()) + self.assertTrue(all(h._con not in old_connections + for h in pool._holders if h.is_connected())) + + async def test_pool_shutdown_cancels_floor_reconnect(self): + for action in ('close', 'terminate'): + with self.subTest(action=action): + started = asyncio.Event() + cancelled = asyncio.Event() + + async def connect(*args, **kwargs): + if started.is_set(): + self.fail('duplicate background connection attempt') + if pool is not None: + started.set() + try: + await asyncio.Future() + finally: + cancelled.set() + return await pg_connection.connect(*args, **kwargs) + + pool = None + pool = await self.create_pool( + init_size=1, min_size=1, max_size=1, connect=connect, + ) + pool._holders[0]._con.terminate() + await asyncio.wait_for(started.wait(), 5) + + # The background connector reserves the holder, so an + # acquirer must wait rather than exceed max_size. + with self.assertRaises(asyncio.TimeoutError): + await pool.acquire(timeout=0.05) + + if action == 'close': + await pool.close() + else: + pool.terminate() + await asyncio.wait_for(cancelled.wait(), 5) + self.assertTrue(pool.is_closing()) + self.assertEqual(pool.get_size(), 0) + + async def test_pool_min_size_zero_allows_full_expiry(self): + # When min_size=0, all idle connections are allowed to expire. + async with self.create_pool( + database='postgres', init_size=2, min_size=0, max_size=2, + max_inactive_connection_lifetime=0.2) as pool: + + self.assertIsNotNone(pool._holders[0]._con) + self.assertIsNotNone(pool._holders[1]._con) + + await asyncio.sleep(0.5) + + self.assertIs(pool._holders[0]._con, None) + self.assertIs(pool._holders[1]._con, None) + + async def test_pool_min_size_validation(self): + # init_size < min_size should raise. + with self.assertRaisesRegex(ValueError, + 'init_size is smaller than min_size'): + await self.create_pool( + database='postgres', init_size=1, min_size=2, max_size=5) + + # min_size > max_size should raise. + with self.assertRaisesRegex(ValueError, + 'min_size is greater than max_size'): + await self.create_pool( + database='postgres', init_size=3, min_size=3, max_size=2) + + # init_size > max_size should raise. + with self.assertRaisesRegex(ValueError, + 'init_size is greater than max_size'): + await self.create_pool( + database='postgres', init_size=5, min_size=1, max_size=3) + + # init_size < 0 should raise. + with self.assertRaisesRegex( + ValueError, + 'init_size is expected to be greater or equal to zero'): + await self.create_pool( + database='postgres', init_size=-1, min_size=0, max_size=3) + + async def test_pool_init_size_and_min_size_getters(self): + async with self.create_pool( + database='postgres', + init_size=3, + min_size=2, + max_size=5) as pool: + self.assertEqual(pool.get_init_size(), 3) + self.assertEqual(pool.get_min_size(), 2) + self.assertEqual(pool.get_max_size(), 5) + self.assertEqual(pool.get_size(), 3) + + async def test_pool_min_size_reconnect_after_expiry(self): + # Connections kept alive by min_size should still be functional. + async with self.create_pool( + database='postgres', init_size=1, min_size=1, max_size=1, + max_inactive_connection_lifetime=0.2) as pool: + + con_before = pool._holders[0]._con + self.assertIsNotNone(con_before) + + await asyncio.sleep(0.5) + + # Connection must still be alive due to min_size=1. + self.assertIs(pool._holders[0]._con, con_before) + + # And it must still work. + result = await pool.fetchval('SELECT 42::int') + self.assertEqual(result, 42) + async def test_pool_handles_inactive_connection_errors(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=0, max_size=1) con = await pool.acquire(timeout=POOL_NOMINAL_TIMEOUT) @@ -811,10 +1209,12 @@ async def test_pool_handles_inactive_connection_errors(self): async def test_pool_size_and_capacity(self): async with self.create_pool( database='postgres', - min_size=2, + init_size=2, + min_size=1, max_size=3, ) as pool: - self.assertEqual(pool.get_min_size(), 2) + self.assertEqual(pool.get_init_size(), 2) + self.assertEqual(pool.get_min_size(), 1) self.assertEqual(pool.get_max_size(), 3) self.assertEqual(pool.get_size(), 2) self.assertEqual(pool.get_idle_size(), 2) @@ -842,7 +1242,7 @@ async def test_pool_closing(self): async def test_pool_handles_transaction_exit_in_asyncgen_1(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) locals_ = {} exec(textwrap.dedent('''\ @@ -863,7 +1263,7 @@ class MyException(Exception): async def test_pool_handles_transaction_exit_in_asyncgen_2(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) locals_ = {} exec(textwrap.dedent('''\ @@ -887,7 +1287,7 @@ class MyException(Exception): async def test_pool_handles_asyncgen_finalization(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) locals_ = {} exec(textwrap.dedent('''\ @@ -908,7 +1308,7 @@ class MyException(Exception): async def test_pool_close_waits_for_release(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) flag = self.loop.create_future() conn_released = False @@ -931,7 +1331,7 @@ async def worker(): async def test_pool_close_timeout(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) flag = self.loop.create_future() @@ -950,7 +1350,7 @@ async def worker(): async def test_pool_expire_connections(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) con = await pool.acquire() try: @@ -963,7 +1363,7 @@ async def test_pool_expire_connections(self): async def test_pool_set_connection_args(self): pool = await self.create_pool(database='postgres', - min_size=1, max_size=1) + init_size=1, min_size=1, max_size=1) # Test that connection is expired on release. con = await pool.acquire() @@ -1005,7 +1405,12 @@ async def test_pool_set_connection_args(self): await pool.close() async def test_pool_init_race(self): - pool = self.create_pool(database='postgres', min_size=1, max_size=1) + pool = self.create_pool( + database='postgres', + init_size=1, + min_size=1, + max_size=1, + ) t1 = asyncio.ensure_future(pool) t2 = asyncio.ensure_future(pool) @@ -1019,7 +1424,12 @@ async def test_pool_init_race(self): await pool.close() async def test_pool_init_and_use_race(self): - pool = self.create_pool(database='postgres', min_size=1, max_size=1) + pool = self.create_pool( + database='postgres', + init_size=1, + min_size=1, + max_size=1, + ) pool_task = asyncio.ensure_future(pool) await asyncio.sleep(0) @@ -1034,7 +1444,7 @@ async def test_pool_init_and_use_race(self): await pool.close() async def test_pool_remote_close(self): - pool = await self.create_pool(min_size=1, max_size=1) + pool = await self.create_pool(init_size=1, min_size=1, max_size=1) backend_pid_fut = self.loop.create_future() async def worker(): @@ -1059,7 +1469,7 @@ async def worker(): await pool.release(conn) async def test_pool_release_after_protocol_abort(self): - pool = await self.create_pool(min_size=1, max_size=1) + pool = await self.create_pool(init_size=1, min_size=0, max_size=1) conn = await pool.acquire() raw_conn = conn._con @@ -1125,6 +1535,7 @@ async def test_full_reconnect_on_node_change_role(self): return pool = await self.create_pool( + init_size=1, min_size=1, max_size=1, target_session_attrs='primary' @@ -1169,7 +1580,7 @@ async def test_standby_pool_01(self): with self.subTest(tasksnum=n): pool = await self.create_pool( database='postgres', user='postgres', - min_size=5, max_size=10) + init_size=5, min_size=0, max_size=10) async def worker(): con = await pool.acquire()