From b01349356dab801d9c5ca87d2bc3d6fb4d7606dc Mon Sep 17 00:00:00 2001
From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Thu, 8 Jun 2023 15:56:22 +0200
Subject: [PATCH] fix(pool): check connections' max_lifetime on ``check()``

Close #482.
---
 docs/news_pool.rst                      |  7 +++++++
 psycopg_pool/psycopg_pool/pool.py       |  9 +++++++++
 psycopg_pool/psycopg_pool/pool_async.py |  9 +++++++++
 tests/pool/test_pool.py                 | 13 +++++++++++++
 tests/pool/test_pool_async.py           | 13 +++++++++++++
 5 files changed, 51 insertions(+)

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index 66a9a8e0b..6da2453ee 100644
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -17,6 +17,13 @@ psycopg_pool 3.2.0 (unreleased)
   (:ticket:`#520`).
 
 
+psycopg_pool 3.1.8 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Enforce connections' ``max_lifetime`` on `~ConnectionPool.check()`
+  (:ticket:`#482`).
+
+
 Current release
 ---------------
 
diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 67447cac5..023f07164 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -437,6 +437,15 @@ class ConnectionPool(BasePool[Connection[Any]]):
 
         while conns:
             conn = conns.pop()
+
+            # Check for expired connections
+            if conn._expire_at <= monotonic():
+                logger.info("discarding expired connection %s", conn)
+                conn.close()
+                self.run_task(AddConnection(self))
+                continue
+
+            # Check for broken connections
             try:
                 conn.execute("SELECT 1")
                 if conn.pgconn.transaction_status == TransactionStatus.INTRANS:
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index 01fdf6317..84e90e46d 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -376,6 +376,15 @@ class AsyncConnectionPool(BasePool[AsyncConnection[Any]]):
 
         while conns:
             conn = conns.pop()
+
+            # Check for expired connections
+            if conn._expire_at <= monotonic():
+                logger.info("discarding expired connection %s", conn)
+                await conn.close()
+                self.run_task(AddConnection(self))
+                continue
+
+            # Check for broken connections
             try:
                 await conn.execute("SELECT 1")
                 if conn.pgconn.transaction_status == TransactionStatus.INTRANS:
diff --git a/tests/pool/test_pool.py b/tests/pool/test_pool.py
index 29c4bf554..af75fe61e 100644
--- a/tests/pool/test_pool.py
+++ b/tests/pool/test_pool.py
@@ -1095,6 +1095,19 @@ def test_check_idle(dsn):
             assert conn.info.transaction_status == TransactionStatus.IDLE
 
 
+@pytest.mark.slow
+def test_check_max_lifetime(dsn):
+    with pool.ConnectionPool(dsn, min_size=1, max_lifetime=0.2) as p:
+        with p.connection() as conn:
+            pid = conn.info.backend_pid
+        with p.connection() as conn:
+            assert conn.info.backend_pid == pid
+        sleep(0.3)
+        p.check()
+        with p.connection() as conn:
+            assert conn.info.backend_pid != pid
+
+
 @pytest.mark.slow
 @pytest.mark.timing
 def test_stats_measures(dsn):
diff --git a/tests/pool/test_pool_async.py b/tests/pool/test_pool_async.py
index 5a0169822..238843e23 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1056,6 +1056,19 @@ async def test_check_idle(dsn):
             assert conn.info.transaction_status == TransactionStatus.IDLE
 
 
+@pytest.mark.slow
+async def test_check_max_lifetime(dsn):
+    async with pool.AsyncConnectionPool(dsn, min_size=1, max_lifetime=0.2) as p:
+        async with p.connection() as conn:
+            pid = conn.info.backend_pid
+        async with p.connection() as conn:
+            assert conn.info.backend_pid == pid
+        await asyncio.sleep(0.3)
+        await p.check()
+        async with p.connection() as conn:
+            assert conn.info.backend_pid != pid
+
+
 @pytest.mark.slow
 @pytest.mark.timing
 async def test_stats_measures(dsn):
-- 
2.47.3