fix(pool): manage CancelledError in some exception handling path

If a CancelledError was raised during check the connection would have
been lost. The exception would have bubbled up but likely users are
using some framework swallowing it because nobody reporting the "lost
connections" issue actually reported the CancelledError.

Close #1123
Close #1208

Author: dvarrazzo

docs/news_pool.rst

@@ -7,6 +7,16 @@
 ``psycopg_pool`` release notes
 ==============================
 
+Future releases
+---------------
+
+psycopg_pool 3.2.8 (unreleased)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+- Don't lose connections if a `~asyncio.CancelledError` is raised in a check
+  (:tickets:`#1123, #1208`)
+
+
 Current release
 ---------------
 

psycopg_pool/psycopg_pool/pool.py

@@ -15,6 +15,7 @@
 from time import monotonic
 from types import TracebackType
 from typing import Any, Generic, Iterator, cast
+from asyncio import CancelledError
 from weakref import ref
 from contextlib import contextmanager
 
@@ -209,7 +210,7 @@ def _getconn_with_check_loop(self, deadline: float) -> CT:
             conn = self._getconn_unchecked(deadline - monotonic())
             try:
                 self._check_connection(conn)
-            except Exception:
+            except (Exception, CancelledError):
                 self._putconn(conn, from_getconn=True)
             else:
                 logger.info("connection given by %r", self.name)
@@ -247,7 +248,7 @@ def _getconn_unchecked(self, timeout: float) -> CT:
         if not conn:
             try:
                 conn = pos.wait(timeout=timeout)
-            except Exception:
+            except BaseException:
                 self._stats[self._REQUESTS_ERRORS] += 1
                 raise
             finally:
@@ -283,7 +284,7 @@ def _check_connection(self, conn: CT) -> None:
             return
         try:
             self._check(conn)
-        except Exception as e:
+        except BaseException as e:
             logger.info("connection failed check: %s", e)
             raise
 
@@ -515,7 +516,7 @@ def check(self) -> None:
             # Check for broken connections
             try:
                 self.check_connection(conn)
-            except Exception:
+            except (Exception, CancelledError):
                 self._stats[self._CONNECTIONS_LOST] += 1
                 logger.warning("discarding broken connection: %s", conn)
                 self.run_task(AddConnection(self))
@@ -576,7 +577,7 @@ def worker(cls, q: Queue[MaintenanceTask]) -> None:
             # Run the task. Make sure don't die in the attempt.
             try:
                 task.run()
-            except Exception as ex:
+            except (Exception, CancelledError) as ex:
                 logger.warning(
                     "task run %s failed: %s: %s", task, ex.__class__.__name__, ex
                 )
@@ -591,7 +592,7 @@ def _connect(self, timeout: float | None = None) -> CT:
         t0 = monotonic()
         try:
             conn = self.connection_class.connect(self.conninfo, **kwargs)
-        except Exception:
+        except (Exception, CancelledError):
             self._stats[self._CONNECTIONS_ERRORS] += 1
             raise
         else:
@@ -628,7 +629,7 @@ def _add_connection(
 
         try:
             conn = self._connect()
-        except Exception as ex:
+        except (Exception, CancelledError) as ex:
             logger.warning("error connecting in %r: %s", self.name, ex)
             if attempt.time_to_give_up(now):
                 logger.warning(
@@ -774,7 +775,7 @@ def _reset_connection(self, conn: CT) -> None:
                     raise e.ProgrammingError(
                         f"connection left in status {sname} by reset function {self._reset}: discarded"
                     )
-            except Exception as ex:
+            except (Exception, CancelledError) as ex:
                 logger.warning("error resetting connection: %s", ex)
                 self._close_connection(conn)
 

psycopg_pool/psycopg_pool/pool_async.py

@@ -12,6 +12,7 @@
 from time import monotonic
 from types import TracebackType
 from typing import Any, AsyncIterator, Generic, cast
+from asyncio import CancelledError
 from weakref import ref
 from contextlib import asynccontextmanager
 
@@ -237,7 +238,7 @@ async def _getconn_with_check_loop(self, deadline: float) -> ACT:
             conn = await self._getconn_unchecked(deadline - monotonic())
             try:
                 await self._check_connection(conn)
-            except Exception:
+            except (Exception, CancelledError):
                 await self._putconn(conn, from_getconn=True)
             else:
                 logger.info("connection given by %r", self.name)
@@ -275,7 +276,7 @@ async def _getconn_unchecked(self, timeout: float) -> ACT:
         if not conn:
             try:
                 conn = await pos.wait(timeout=timeout)
-            except Exception:
+            except BaseException:
                 self._stats[self._REQUESTS_ERRORS] += 1
                 raise
             finally:
@@ -312,7 +313,7 @@ async def _check_connection(self, conn: ACT) -> None:
             return
         try:
             await self._check(conn)
-        except Exception as e:
+        except BaseException as e:
             logger.info("connection failed check: %s", e)
             raise
 
@@ -551,7 +552,7 @@ async def check(self) -> None:
             # Check for broken connections
             try:
                 await self.check_connection(conn)
-            except Exception:
+            except (Exception, CancelledError):
                 self._stats[self._CONNECTIONS_LOST] += 1
                 logger.warning("discarding broken connection: %s", conn)
                 self.run_task(AddConnection(self))
@@ -630,7 +631,7 @@ async def worker(cls, q: AQueue[MaintenanceTask]) -> None:
             # Run the task. Make sure don't die in the attempt.
             try:
                 await task.run()
-            except Exception as ex:
+            except (Exception, CancelledError) as ex:
                 logger.warning(
                     "task run %s failed: %s: %s", task, ex.__class__.__name__, ex
                 )
@@ -645,7 +646,7 @@ async def _connect(self, timeout: float | None = None) -> ACT:
         t0 = monotonic()
         try:
             conn = await self.connection_class.connect(self.conninfo, **kwargs)
-        except Exception:
+        except (Exception, CancelledError):
             self._stats[self._CONNECTIONS_ERRORS] += 1
             raise
         else:
@@ -683,7 +684,7 @@ async def _add_connection(
 
         try:
             conn = await self._connect()
-        except Exception as ex:
+        except (Exception, CancelledError) as ex:
             logger.warning("error connecting in %r: %s", self.name, ex)
             if attempt.time_to_give_up(now):
                 logger.warning(
@@ -835,7 +836,7 @@ async def _reset_connection(self, conn: ACT) -> None:
                         f"connection left in status {sname} by reset function"
                         f" {self._reset}: discarded"
                     )
-            except Exception as ex:
+            except (Exception, CancelledError) as ex:
                 logger.warning("error resetting connection: %s", ex)
                 await self._close_connection(conn)
 

tests/pool/test_pool_common.py

@@ -6,6 +6,7 @@
 import logging
 from time import time
 from typing import Any
+from asyncio import CancelledError
 
 import pytest
 
@@ -694,6 +695,29 @@ def worker(i):
             assert cur.fetchone() == (1,)
 
 
+@skip_sync
+def test_cancel_on_check(pool_cls, dsn):
+    do_cancel = True
+
+    def check(conn):
+        nonlocal do_cancel
+        if do_cancel:
+            do_cancel = False
+            raise CancelledError()
+
+        pool_cls.check_connection(conn)
+
+    with pool_cls(dsn, min_size=min_size(pool_cls, 1), check=check, timeout=1.0) as p:
+        try:
+            with p.connection() as conn:
+                conn.execute("select 1")
+        except CancelledError:
+            pass
+
+        with p.connection() as conn:
+            conn.execute("select 1")
+
+
 def min_size(pool_cls, num=1):
     """Return the minimum min_size supported by the pool class."""
     if pool_cls is pool.ConnectionPool:

tests/pool/test_pool_common_async.py

@@ -3,6 +3,7 @@
 import logging
 from time import time
 from typing import Any
+from asyncio import CancelledError
 
 import pytest
 
@@ -705,6 +706,31 @@ async def worker(i):
             assert await cur.fetchone() == (1,)
 
 
+@skip_sync
+async def test_cancel_on_check(pool_cls, dsn):
+    do_cancel = True
+
+    async def check(conn):
+        nonlocal do_cancel
+        if do_cancel:
+            do_cancel = False
+            raise CancelledError()
+
+        await pool_cls.check_connection(conn)
+
+    async with pool_cls(
+        dsn, min_size=min_size(pool_cls, 1), check=check, timeout=1.0
+    ) as p:
+        try:
+            async with p.connection() as conn:
+                await conn.execute("select 1")
+        except CancelledError:
+            pass
+
+        async with p.connection() as conn:
+            await conn.execute("select 1")
+
+
 def min_size(pool_cls, num=1):
     """Return the minimum min_size supported by the pool class."""
     if pool_cls is pool.AsyncConnectionPool: