From: Daniele Varrazzo <daniele.varrazzo@gmail.com>
Date: Fri, 29 Dec 2023 13:15:07 +0000 (+0100)
Subject: fix(pool): respect timeout on getconn() when the check function fails
X-Git-Tag: pool-3.2.1~6^2~2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=4ce737dd97b84f34b6743a508356074dc246946a;p=thirdparty%2Fpsycopg.git

fix(pool): respect timeout on getconn() when the check function fails

Close #709.
---

diff --git a/docs/news_pool.rst b/docs/news_pool.rst
index ec758550d..ea969810e 100644
--- a/docs/news_pool.rst
+++ b/docs/news_pool.rst
@@ -13,6 +13,8 @@ Future releases
 psycopg_pool 3.2.1 (unreleased)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
+- Respect timeout on `~ConnectionPool.connection()` when `!check` fails
+  (:ticket:`#709`).
 - Use `typing.Self` as a more correct return value annotation of context
   managers and other self-returning methods (see :ticket:`708`).
 
diff --git a/psycopg_pool/psycopg_pool/null_pool.py b/psycopg_pool/psycopg_pool/null_pool.py
index 3e5651acb..a5408d08e 100644
--- a/psycopg_pool/psycopg_pool/null_pool.py
+++ b/psycopg_pool/psycopg_pool/null_pool.py
@@ -96,6 +96,9 @@ class NullConnectionPool(_BaseNullConnectionPool, ConnectionPool[CT]):
         logger.info("pool %r is ready to use", self.name)
 
     def _get_ready_connection(self, timeout: Optional[float]) -> Optional[CT]:
+        if timeout is not None and timeout <= 0.0:
+            raise PoolTimeout()
+
         conn: Optional[CT] = None
         if self.max_size == 0 or self._nconns < self.max_size:
             # Create a new connection for the client
diff --git a/psycopg_pool/psycopg_pool/null_pool_async.py b/psycopg_pool/psycopg_pool/null_pool_async.py
index 74725e87e..a9baa8e12 100644
--- a/psycopg_pool/psycopg_pool/null_pool_async.py
+++ b/psycopg_pool/psycopg_pool/null_pool_async.py
@@ -93,6 +93,9 @@ class AsyncNullConnectionPool(_BaseNullConnectionPool, AsyncConnectionPool[ACT])
         logger.info("pool %r is ready to use", self.name)
 
     async def _get_ready_connection(self, timeout: Optional[float]) -> Optional[ACT]:
+        if timeout is not None and timeout <= 0.0:
+            raise PoolTimeout()
+
         conn: Optional[ACT] = None
         if self.max_size == 0 or self._nconns < self.max_size:
             # Create a new connection for the client
diff --git a/psycopg_pool/psycopg_pool/pool.py b/psycopg_pool/psycopg_pool/pool.py
index 85765c96f..0b2d7ea33 100644
--- a/psycopg_pool/psycopg_pool/pool.py
+++ b/psycopg_pool/psycopg_pool/pool.py
@@ -187,10 +187,9 @@ class ConnectionPool(Generic[CT], BasePool):
         failing to do so will deplete the pool. A depleted pool is a sad pool:
         you don't want a depleted pool.
         """
-        t0 = monotonic()
         if timeout is None:
             timeout = self.timeout
-        deadline = t0 + timeout
+        deadline = monotonic() + timeout
 
         logger.info("connection requested from %r", self.name)
         self._stats[self._REQUESTS_NUM] += 1
@@ -249,6 +248,9 @@ class ConnectionPool(Generic[CT], BasePool):
 
     def _get_ready_connection(self, timeout: Optional[float]) -> Optional[CT]:
         """Return a connection, if the client deserves one."""
+        if timeout is not None and timeout <= 0.0:
+            raise PoolTimeout()
+
         conn: Optional[CT] = None
         if self._pool:
             # Take a connection ready out of the pool
diff --git a/psycopg_pool/psycopg_pool/pool_async.py b/psycopg_pool/psycopg_pool/pool_async.py
index 719b7221f..0a7e7e302 100644
--- a/psycopg_pool/psycopg_pool/pool_async.py
+++ b/psycopg_pool/psycopg_pool/pool_async.py
@@ -207,10 +207,9 @@ class AsyncConnectionPool(Generic[ACT], BasePool):
         failing to do so will deplete the pool. A depleted pool is a sad pool:
         you don't want a depleted pool.
         """
-        t0 = monotonic()
         if timeout is None:
             timeout = self.timeout
-        deadline = t0 + timeout
+        deadline = monotonic() + timeout
 
         logger.info("connection requested from %r", self.name)
         self._stats[self._REQUESTS_NUM] += 1
@@ -270,6 +269,9 @@ class AsyncConnectionPool(Generic[ACT], BasePool):
 
     async def _get_ready_connection(self, timeout: Optional[float]) -> Optional[ACT]:
         """Return a connection, if the client deserves one."""
+        if timeout is not None and timeout <= 0.0:
+            raise PoolTimeout()
+
         conn: Optional[ACT] = None
         if self._pool:
             # Take a connection ready out of the pool
diff --git a/tests/pool/test_pool_common.py b/tests/pool/test_pool_common.py
index ee8d03ffd..b37026715 100644
--- a/tests/pool/test_pool_common.py
+++ b/tests/pool/test_pool_common.py
@@ -610,6 +610,20 @@ def test_check_init(pool_cls, dsn):
     assert checked
 
 
+@pytest.mark.slow
+def test_check_timeout(pool_cls, dsn):
+    def check(conn):
+        raise Exception()
+
+    t0 = time()
+    with pytest.raises(pool.PoolTimeout):
+        with pool_cls(dsn, check=check, timeout=1.0) as p:
+            with p.connection():
+                assert False
+
+    assert time() - t0 <= 1.5
+
+
 @skip_sync
 def test_cancellation_in_queue(pool_cls, dsn):
     # https://github.com/psycopg/psycopg/issues/509
diff --git a/tests/pool/test_pool_common_async.py b/tests/pool/test_pool_common_async.py
index dbcfbbe46..567ad5b70 100644
--- a/tests/pool/test_pool_common_async.py
+++ b/tests/pool/test_pool_common_async.py
@@ -629,6 +629,20 @@ async def test_check_init(pool_cls, dsn):
     assert checked
 
 
+@pytest.mark.slow
+async def test_check_timeout(pool_cls, dsn):
+    async def check(conn):
+        raise Exception()
+
+    t0 = time()
+    with pytest.raises(pool.PoolTimeout):
+        async with pool_cls(dsn, check=check, timeout=1.0) as p:
+            async with p.connection():
+                assert False
+
+    assert time() - t0 <= 1.5
+
+
 @skip_sync
 async def test_cancellation_in_queue(pool_cls, dsn):
     # https://github.com/psycopg/psycopg/issues/509