Implement support for pool connection rotation

The new `Pool.cycle()` method expires all currently open connections,
so they would be replaced with fresh ones on the next `acquire()`
attempt.

The new `Pool.set_connect_args()` allows changing the connection
arguments for an existing pool instance.  Coupled with `cycle()`, it
allows adapting the pool to the new environment conditions without
having to replace the pool instance.

Fixes: #291

Author: elprans

asyncpg/_testbase/__init__.py

@@ -328,17 +328,29 @@ def get_connection_spec(cls, kwargs={}):
                 conn_spec['user'] = 'postgres'
         return conn_spec
 
-    def create_pool(self, pool_class=pg_pool.Pool,
-                    connection_class=pg_connection.Connection, **kwargs):
-        conn_spec = self.get_connection_spec(kwargs)
-        return create_pool(loop=self.loop, pool_class=pool_class,
-                           connection_class=connection_class, **conn_spec)
-
     @classmethod
     def connect(cls, **kwargs):
         conn_spec = cls.get_connection_spec(kwargs)
         return pg_connection.connect(**conn_spec, loop=cls.loop)
 
+    def setUp(self):
+        super().setUp()
+        self._pools = []
+
+    def tearDown(self):
+        super().tearDown()
+        for pool in self._pools:
+            pool.terminate()
+        self._pools = []
+
+    def create_pool(self, pool_class=pg_pool.Pool,
+                    connection_class=pg_connection.Connection, **kwargs):
+        conn_spec = self.get_connection_spec(kwargs)
+        pool = create_pool(loop=self.loop, pool_class=pool_class,
+                           connection_class=connection_class, **conn_spec)
+        self._pools.append(pool)
+        return pool
+
 
 class ProxiedClusterTestCase(ClusterTestCase):
     @classmethod

asyncpg/pool.py

@@ -92,67 +92,46 @@ def __repr__(self):
 class PoolConnectionHolder:
 
     __slots__ = ('_con', '_pool', '_loop', '_proxy',
-                 '_connect_args', '_connect_kwargs',
-                 '_max_queries', '_setup', '_init',
+                 '_max_queries', '_setup',
                  '_max_inactive_time', '_in_use',
-                 '_inactive_callback', '_timeout')
+                 '_inactive_callback', '_timeout',
+                 '_generation')
 
-    def __init__(self, pool, *, connect_args, connect_kwargs,
-                 max_queries, setup, init, max_inactive_time):
+    def __init__(self, pool, *, max_queries, setup, max_inactive_time):
 
         self._pool = pool
         self._con = None
         self._proxy = 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 = None  # type: asyncio.Future
         self._timeout = None
+        self._generation = None
 
     async def connect(self):
         if self._con is not None:
             raise exceptions.InternalClientError(
                 'PoolConnectionHolder.connect() called while another '
                 'connection already exists')
 
-        if self._pool._working_addr is None:
-            # First connection attempt on this pool.
-            con = await connection.connect(
-                *self._connect_args,
-                loop=self._pool._loop,
-                connection_class=self._pool._connection_class,
-                **self._connect_kwargs)
-
-            self._pool._working_addr = con._addr
-            self._pool._working_config = con._config
-            self._pool._working_params = con._params
-
-        else:
-            # We've connected before and have a resolved address,
-            # and parsed options and config.
-            con = await connect_utils._connect_addr(
-                loop=self._pool._loop,
-                addr=self._pool._working_addr,
-                timeout=self._pool._working_params.connect_timeout,
-                config=self._pool._working_config,
-                params=self._pool._working_params,
-                connection_class=self._pool._connection_class)
-
-        if self._init is not None:
-            await self._init(con)
-
-        self._con = con
+        self._con = await self._pool._get_new_connection()
+        self._generation = self._pool._generation
 
     async def acquire(self) -> PoolConnectionProxy:
         if self._con is None or self._con.is_closed():
             self._con = None
             await self.connect()
 
+        elif self._generation != self._pool._generation:
+            # Connections have been expired, re-connect the holder.
+            self._pool._loop.create_task(
+                self._con.close(timeout=self._timeout))
+            self._con = None
+            await self.connect()
+
         self._maybe_cancel_inactive_callback()
 
         self._proxy = proxy = PoolConnectionProxy(self, self._con)
@@ -197,6 +176,13 @@ async def release(self, timeout):
             await self._con.close(timeout=timeout)
             return
 
+        if self._generation != self._pool._generation:
+            # The connection has expired because it belongs to
+            # an older generation (Pool.expire_connections() has
+            # been called.)
+            await self._con.close(timeout=timeout)
+            return
+
         try:
             budget = timeout
 
@@ -312,9 +298,10 @@ class Pool:
     """
 
     __slots__ = ('_queue', '_loop', '_minsize', '_maxsize',
+                 '_init', '_connect_args', '_connect_kwargs',
                  '_working_addr', '_working_config', '_working_params',
                  '_holders', '_initialized', '_closing', '_closed',
-                 '_connection_class')
+                 '_connection_class', '_generation')
 
     def __init__(self, *connect_args,
                  min_size,
@@ -349,6 +336,11 @@ def __init__(self, *connect_args,
                 'max_inactive_connection_lifetime is expected to be greater '
                 'or equal to zero')
 
+        if not issubclass(connection_class, connection.Connection):
+            raise TypeError(
+                'connection_class is expected to be a subclass of '
+                'asyncpg.Connection, got {!r}'.format(connection_class))
+
         self._minsize = min_size
         self._maxsize = max_size
 
@@ -364,16 +356,17 @@ def __init__(self, *connect_args,
 
         self._closing = False
         self._closed = False
+        self._generation = 0
+        self._init = init
+        self._connect_args = connect_args
+        self._connect_kwargs = connect_kwargs
 
         for _ in range(max_size):
             ch = PoolConnectionHolder(
                 self,
-                connect_args=connect_args,
-                connect_kwargs=connect_kwargs,
                 max_queries=max_queries,
                 max_inactive_time=max_inactive_connection_lifetime,
-                setup=setup,
-                init=init)
+                setup=setup)
 
             self._holders.append(ch)
             self._queue.put_nowait(ch)
@@ -409,6 +402,73 @@ async def _async__init__(self):
         self._initialized = True
         return self
 
+    def set_connect_args(self, dsn=None, *,
+                         connection_class=connection.Connection, **kwargs):
+        """Set the new connection arguments for this pool.
+
+        The new connection arguments will be used for all subsequent
+        new connection attempts.  Existing connections will remain until
+        they expire. Use :meth:`Pool.expire_connections()
+        <asyncpg.pool.Pool.expire_connections>` to expedite the connection
+        expiry.
+
+        :param str dsn:
+            Connection arguments specified using as a single string in
+            the following format:
+            ``postgres://user:pass@host:port/database?option=value``.
+
+        :param \*\*connect_kwargs:
+            Keyword arguments for the :func:`~asyncpg.connection.connect`
+            function.
+
+        :param Connection connection_class:
+            The class to use for connections.  Must be a subclass of
+            :class:`~asyncpg.connection.Connection`.
+
+        .. versionadded:: 0.16.0
+        """
+
+        if not issubclass(connection_class, connection.Connection):
+            raise TypeError(
+                'connection_class is expected to be a subclass of '
+                'asyncpg.Connection, got {!r}'.format(connection_class))
+
+        self._connect_args = [dsn]
+        self._connect_kwargs = kwargs
+        self._connection_class = connection_class
+        self._working_addr = None
+        self._working_config = None
+        self._working_params = None
+
+    async def _get_new_connection(self):
+        if self._working_addr is None:
+            # First connection attempt on this pool.
+            con = await connection.connect(
+                *self._connect_args,
+                loop=self._loop,
+                connection_class=self._connection_class,
+                **self._connect_kwargs)
+
+            self._working_addr = con._addr
+            self._working_config = con._config
+            self._working_params = con._params
+
+        else:
+            # We've connected before and have a resolved address,
+            # and parsed options and config.
+            con = await connect_utils._connect_addr(
+                loop=self._loop,
+                addr=self._working_addr,
+                timeout=self._working_params.connect_timeout,
+                config=self._working_config,
+                params=self._working_params,
+                connection_class=self._connection_class)
+
+        if self._init is not None:
+            await self._init(con)
+
+        return con
+
     async def execute(self, query: str, *args, timeout: float=None) -> str:
         """Execute an SQL command (or commands).
 
@@ -602,6 +662,16 @@ def terminate(self):
             ch.terminate()
         self._closed = True
 
+    async def expire_connections(self):
+        """Expire all currently open connections.
+
+        Cause all currently open connections to get replaced on the
+        next :meth:`~asyncpg.pool.Pool.acquire()` call.
+
+        .. versionadded:: 0.16.0
+        """
+        self._generation += 1
+
     def _check_init(self):
         if not self._initialized:
             raise exceptions.InterfaceError('pool is not initialized')
@@ -708,6 +778,10 @@ def create_pool(dsn=None, *,
         Keyword arguments for the :func:`~asyncpg.connection.connect`
         function.
 
+    :param Connection connection_class:
+        The class to use for connections.  Must be a subclass of
+        :class:`~asyncpg.connection.Connection`.
+
     :param int min_size:
         Number of connection the pool will be initialized with.
 
@@ -759,11 +833,6 @@ def create_pool(dsn=None, *,
        <connection.Connection.add_log_listener>`) present on the connection
        at the moment of its release to the pool.
     """
-    if not issubclass(connection_class, connection.Connection):
-        raise TypeError(
-            'connection_class is expected to be a subclass of '
-            'asyncpg.Connection, got {!r}'.format(connection_class))
-
     return Pool(
         dsn,
         connection_class=connection_class,

tests/test_pool.py

@@ -794,6 +794,43 @@ async def worker():
 
         await task
 
+    async def test_pool_expire_connections(self):
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1)
+
+        con = await pool.acquire()
+        try:
+            await pool.expire_connections()
+        finally:
+            await pool.release(con)
+
+        self.assertIsNone(pool._holders[0]._con)
+
+    async def test_pool_set_connection_args(self):
+        pool = await self.create_pool(database='postgres',
+                                      min_size=1, max_size=1)
+
+        # Test that connection is expired on release.
+        con = await pool.acquire()
+        pool.set_connect_args(
+            server_settings={'application_name': 'set_conn_args_test'})
+        await pool.expire_connections()
+        await pool.release(con)
+
+        con = await pool.acquire()
+        self.assertEqual(con.get_settings().application_name,
+                         'set_conn_args_test')
+        await pool.release(con)
+
+        # Test that connection is expired before acquire.
+        pool.set_connect_args(
+            server_settings={'application_name': 'set_conn_args_test_2'})
+        await pool.expire_connections()
+
+        con = await pool.acquire()
+        self.assertEqual(con.get_settings().application_name,
+                         'set_conn_args_test_2')
+
 
 @unittest.skipIf(os.environ.get('PGHOST'), 'using remote cluster for testing')
 class TestHotStandby(tb.ClusterTestCase):