From 6f20f9af9014fcc34ff1f68dc0d232235ae5825e 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                      | 10 ++++++++++
 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, 54 insertions(+)

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index 8d4ec4567..2f94544e4 100644
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -7,6 +7,16 @@
 ``psycopg_pool`` release notes
 ==============================
 
+Future releases
+---------------
+
+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 498ef7020..1a1a85266 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -432,6 +432,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 74a30dd28..27c5dcba2 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -370,6 +370,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 83faf18da..df1348a35 100644
--- a/tests/pool/test_pool_async.py
+++ b/tests/pool/test_pool_async.py
@@ -1046,6 +1046,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.2