SELECT 1 only if we haven't seen any events (#222)

* SELECT 1 only if we haven't seen any events

* CHANGELOG

* Update asyncpg_listen/listener.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Update asyncpg_listen/listener.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Fix

* Remove unused tasks

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: Pliner

CHANGELOG.md

@@ -1,3 +1,8 @@
+## unreleased
+
+# [SELECT 1 only if we haven't seen any events](https://github.com/anna-money/asyncpg-listen/pull/222)
+
+
 ## v0.0.8 (2025-05-25)
 
 * Support python 3.13

asyncpg_listen/listener.py

@@ -3,6 +3,7 @@
 import enum
 import logging
 import sys
+import time
 from typing import Any, Callable, Coroutine
 
 import asyncpg
@@ -31,6 +32,7 @@ class Notification:
 NotificationHandler = Callable[[NotificationOrTimeout], Coroutine]
 
 NO_TIMEOUT: float = -1
+MAX_SILENCED_FAILED_CONNECT_ATTEMPTS = 3
 
 
 def connect_func(*args: Any, **kwargs: Any) -> ConnectFunc:
@@ -63,7 +65,8 @@ async def run(
         async with asyncio.TaskGroup() as tg:
             tg.create_task(
                 self.__read_notifications(
-                    queue_per_channel=queue_per_channel, check_interval=max(1.0, notification_timeout / 3.0)
+                    queue_per_channel=queue_per_channel,
+                    notification_timeout=notification_timeout,
                 ),
                 name=__package__,
             )
@@ -128,31 +131,45 @@ async def run_coro(c: Coroutine) -> None:
                 logger.exception("Failed to handle %s", notification)
 
     async def __read_notifications(
-        self, queue_per_channel: dict[str, asyncio.Queue[Notification]], check_interval: float
+        self, queue_per_channel: dict[str, asyncio.Queue[Notification]], notification_timeout: float
     ) -> None:
         failed_connect_attempts = 0
+        per_attempt_keep_alive_budget = max(1.0, notification_timeout / 3.0)
         while True:
             try:
                 connection = await self.__connect()
                 failed_connect_attempts = 0
                 try:
+                    event = asyncio.Event()
                     for channel, queue in queue_per_channel.items():
-                        await connection.add_listener(channel, self.__get_push_callback(queue))
-
-                    while True:
-                        await asyncio.sleep(check_interval)
-                        await connection.execute("SELECT 1")
+                        await connection.add_listener(channel, self.__get_push_callback(queue, event))
+
+                    while not connection.is_closed():
+                        started_at = time.monotonic()
+                        if not event.is_set():
+                            await connection.execute("SELECT 1", timeout=per_attempt_keep_alive_budget)
+                        event.clear()
+                        finished_at = time.monotonic()
+                        elapsed = finished_at - started_at
+                        if elapsed < per_attempt_keep_alive_budget:
+                            await asyncio.sleep(per_attempt_keep_alive_budget - elapsed)
+                    logger.warning("Connection was lost")
                 finally:
                     await asyncio.shield(connection.close())
             except Exception:
-                logger.exception("Connection was lost or not established")
-
+                if failed_connect_attempts < MAX_SILENCED_FAILED_CONNECT_ATTEMPTS:
+                    logger.warning("Connection was lost or not established", exc_info=True)
+                else:
+                    logger.exception("Connection was lost or not established")
                 await asyncio.sleep(self.__reconnect_delay * failed_connect_attempts)
                 failed_connect_attempts += 1
 
     @staticmethod
-    def __get_push_callback(queue: asyncio.Queue[Notification]) -> Callable[[Any, Any, Any, Any], None]:
+    def __get_push_callback(
+        queue: asyncio.Queue[Notification], event: asyncio.Event
+    ) -> Callable[[Any, Any, Any, Any], None]:
         def _push(_: Any, __: Any, channel: Any, payload: Any) -> None:
             queue.put_nowait(Notification(channel, payload))
+            event.set()
 
         return _push

requirements-dev.txt

@@ -1,6 +1,6 @@
 pytest==8.3.4
 pytest-aiohttp==1.0.5
-pytest-pg==0.0.21
+pytest-pg==0.0.25
 pytest-runner==6.0.1
 pytest-asyncio==0.24.0
 isort==5.13.2

tests/test_listener.py

@@ -1,9 +1,11 @@
 import asyncio
 import contextlib
 import dataclasses
-from typing import Any, Awaitable, Callable, List
+import logging
+from typing import Awaitable, Callable
 
 import asyncpg
+import pytest
 import pytest_pg
 
 import asyncpg_listen
@@ -14,7 +16,7 @@
 class Handler:
     def __init__(self, delay: float = 0) -> None:
         self.delay = delay
-        self.notifications: List[asyncpg_listen.NotificationOrTimeout] = []
+        self.notifications: list[asyncpg_listen.NotificationOrTimeout] = []
 
     async def handle(self, notification: asyncpg_listen.NotificationOrTimeout) -> None:
         await asyncio.sleep(self.delay)
@@ -27,10 +29,10 @@ async def cancel_and_wait(future: "asyncio.Future[None]") -> None:
         await future
 
 
-async def test_two_inactive_channels(pg_11: pytest_pg.PG) -> None:
+async def test_two_inactive_channels(pg_14: pytest_pg.PG) -> None:
     handler_1 = Handler()
     handler_2 = Handler()
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(
         listener.run({"inactive_1": handler_1.handle, "inactive_2": handler_2.handle}, notification_timeout=1)
     )
@@ -42,14 +44,14 @@ async def test_two_inactive_channels(pg_11: pytest_pg.PG) -> None:
     assert handler_2.notifications == [asyncpg_listen.Timeout("inactive_2")]
 
 
-async def test_one_active_channel_and_one_passive_channel(pg_11: pytest_pg.PG) -> None:
+async def test_one_active_channel_and_one_passive_channel(pg_14: pytest_pg.PG) -> None:
     active_handler = Handler()
     inactive_handler = Handler()
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(
         listener.run({"active": active_handler.handle, "inactive": inactive_handler.handle}, notification_timeout=1)
     )
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         await asyncio.sleep(0.75)
         await connection.execute("NOTIFY active, '1'")
@@ -69,16 +71,16 @@ async def test_one_active_channel_and_one_passive_channel(pg_11: pytest_pg.PG) -
     ]
 
 
-async def test_two_active_channels(pg_11: pytest_pg.PG) -> None:
+async def test_two_active_channels(pg_14: pytest_pg.PG) -> None:
     handler_1 = Handler()
     handler_2 = Handler()
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(
         listener.run({"active_1": handler_1.handle, "active_2": handler_2.handle}, notification_timeout=1)
     )
     await asyncio.sleep(0.1)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         await connection.execute("NOTIFY active_1, '1'")
         await connection.execute("NOTIFY active_2, '2'")
@@ -100,15 +102,15 @@ async def test_two_active_channels(pg_11: pytest_pg.PG) -> None:
     ]
 
 
-async def test_listen_policy_last(pg_11: pytest_pg.PG) -> None:
+async def test_listen_policy_last(pg_14: pytest_pg.PG) -> None:
     handler = Handler(delay=0.1)
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(
         listener.run({"simple": handler.handle}, policy=asyncpg_listen.ListenPolicy.LAST, notification_timeout=1)
     )
     await asyncio.sleep(0.1)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         for i in range(10):
             await connection.execute(f"NOTIFY simple, '{i}'")
@@ -124,13 +126,13 @@ async def test_listen_policy_last(pg_11: pytest_pg.PG) -> None:
     ]
 
 
-async def test_listen_policy_all(pg_11: pytest_pg.PG) -> None:
+async def test_listen_policy_all(pg_14: pytest_pg.PG) -> None:
     handler = Handler(delay=0.05)
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(listener.run({"simple": handler.handle}, notification_timeout=1))
     await asyncio.sleep(0.1)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         for i in range(10):
             await connection.execute(f"NOTIFY simple, '{i}'")
@@ -171,16 +173,16 @@ async def connect() -> asyncpg.Connection:
     assert handler.notifications == []
 
 
-async def test_failing_handler(pg_11: pytest_pg.PG) -> None:
+async def test_failing_handler(pg_14: pytest_pg.PG) -> None:
     async def handle(_: asyncpg_listen.NotificationOrTimeout) -> None:
         raise RuntimeError("Oops")
 
-    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_11)))
+    listener = asyncpg_listen.NotificationListener(asyncpg_listen.connect_func(**dataclasses.asdict(pg_14)))
     listener_task = asyncio.create_task(listener.run({"simple": handle}, notification_timeout=1))
 
     await asyncio.sleep(0.1)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         await connection.execute("NOTIFY simple")
         await connection.execute("NOTIFY simple")
@@ -198,22 +200,23 @@ async def handle(_: asyncpg_listen.NotificationOrTimeout) -> None:
 async def test_reconnect(
     tcp_proxy: Callable[[int, int], Awaitable[TcpProxy]],
     unused_port: Callable[[], int],
-    pg_11: pytest_pg.PG,
-    caplog: Any,
+    pg_14: pytest_pg.PG,
+    caplog: pytest.LogCaptureFixture,
 ) -> None:
+    caplog.set_level(logging.WARNING)
     proxy_port = unused_port()
 
     handler = Handler()
-    proxy = await tcp_proxy(proxy_port, pg_11.port)
+    proxy = await tcp_proxy(proxy_port, pg_14.port)
     listener = asyncpg_listen.NotificationListener(
-        asyncpg_listen.connect_func(**{**(dataclasses.asdict(pg_11)), **{"port": proxy_port}})
+        asyncpg_listen.connect_func(**{**(dataclasses.asdict(pg_14)), **{"port": proxy_port}})
     )
 
     listener_task = asyncio.create_task(listener.run({"simple": handler.handle}, notification_timeout=1))
 
     await asyncio.sleep(0.5)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         await connection.execute("NOTIFY simple, 'before'")
     finally:
@@ -223,7 +226,7 @@ async def test_reconnect(
     await proxy.drop_connections()
     await asyncio.sleep(2)
 
-    connection = await asyncpg.connect(**dataclasses.asdict(pg_11))
+    connection = await asyncpg.connect(**dataclasses.asdict(pg_14))
     try:
         await connection.execute("NOTIFY simple, 'after'")
     finally:
@@ -235,4 +238,8 @@ async def test_reconnect(
     assert asyncpg_listen.Notification("simple", "before") in handler.notifications
     assert asyncpg_listen.Notification("simple", "after") in handler.notifications
 
-    assert any(record for record in caplog.records if "Connection was lost or not established" in record.message)
+    assert any(
+        record
+        for record in caplog.records
+        if "Connection was lost or not established" in record.message or "Connection was lost" in record.message
+    )