Add no_message_restart_interval to reconnect when no messages arrive for a while (#182)

Author: vrslev

README.md

@@ -41,6 +41,7 @@ async with stompman.Client(
     disconnect_confirmation_timeout=2,
     write_retry_attempts=3,
     check_server_alive_interval_factor=3,
+    no_message_restart_interval=datetime.timedelta(hours=1),  # None to disable
 ) as client:
     ...
 ```
@@ -149,6 +150,7 @@ stompman takes care of cleaning up resources automatically. When you leave the c
 - If multiple servers were provided, stompman will attempt to connect to each one simultaneously and will use the first that succeeds. If all servers fail to connect, an `stompman.FailedAllConnectAttemptsError` will be raised. In normal situation it doesn't need to be handled: tune retry and timeout parameters in `stompman.Client()` to your needs.
 
 - When connection is lost, stompman will attempt to handle it automatically. `stompman.FailedAllConnectAttemptsError` will be raised if all connection attempts fail. `stompman.FailedAllWriteAttemptsError` will be raised if connection succeeds but sending a frame or heartbeat lead to losing connection.
+- If no messages are received for `no_message_restart_interval` (defaults to 1 hour), stompman will force a reconnect. Set to `None` to disable.
 - To implement health checks, use `stompman.Client.is_alive()` — it will return `True` if everything is OK and `False` if server is not responding.
 - `stompman` will write log warnings when connection is lost, after successful reconnection or invalid state during ack/nack.
 

packages/stompman/stompman/client.py

@@ -1,7 +1,9 @@
 import asyncio
+import time
 from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
+from datetime import timedelta
 from functools import partial
 from ssl import SSLContext
 from types import TracebackType
@@ -45,6 +47,8 @@ class Client:
     disconnect_confirmation_timeout: int = 2
     check_server_alive_interval_factor: int = 3
     """Client will check if server alive `server heartbeat interval` times `interval factor`"""
+    no_message_restart_interval: timedelta | None = timedelta(hours=1)
+    """Force reconnect if no messages received within this interval. None to disable."""
 
     connection_class: type[AbstractConnection] = Connection
 
@@ -74,6 +78,7 @@ def __post_init__(self) -> None:
             read_max_chunk_size=self.read_max_chunk_size,
             write_retry_attempts=self.write_retry_attempts,
             check_server_alive_interval_factor=self.check_server_alive_interval_factor,
+            no_message_restart_interval=self.no_message_restart_interval,
             ssl=self.ssl,
         )
 
@@ -99,6 +104,7 @@ async def _listen_to_frames(self) -> None:
             async for frame in self._connection_manager.read_frames_reconnecting():
                 match frame:
                     case MessageFrame():
+                        self._connection_manager._last_message_received_time = time.time()
                         if subscription := self._active_subscriptions.get_by_id(frame.headers["subscription"]):
                             task_group.create_task(
                                 subscription._run_handler(frame=frame)

packages/stompman/stompman/connection_manager.py

@@ -2,6 +2,7 @@
 import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
+from datetime import timedelta
 from ssl import SSLContext
 from types import TracebackType
 from typing import TYPE_CHECKING, Literal, Self
@@ -50,12 +51,15 @@ class ConnectionManager:
     read_max_chunk_size: int
     write_retry_attempts: int
     check_server_alive_interval_factor: int
+    no_message_restart_interval: timedelta | None
 
     _active_connection_state: ActiveConnectionState | None = field(default=None, init=False)
     _reconnect_lock: asyncio.Lock = field(init=False, default_factory=asyncio.Lock)
     _task_group: asyncio.TaskGroup = field(init=False, default_factory=asyncio.TaskGroup)
     _send_heartbeat_task: asyncio.Task[None] = field(init=False, repr=False)
+    _monitor_no_message_task: asyncio.Task[None] | None = field(default=None, init=False, repr=False)
     _reconnection_count: int = field(default=0, init=False)
+    _last_message_received_time: float = field(init=False, default_factory=time.time)
 
     async def __aenter__(self) -> Self:
         await self._task_group.__aenter__()
@@ -67,7 +71,11 @@ async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
     ) -> None:
         self._send_heartbeat_task.cancel()
-        await asyncio.wait([self._send_heartbeat_task])
+        tasks = [self._send_heartbeat_task]
+        if self._monitor_no_message_task is not None:
+            self._monitor_no_message_task.cancel()
+            tasks.append(self._monitor_no_message_task)
+        await asyncio.wait(tasks)
         await self._task_group.__aexit__(exc_type, exc_value, traceback)
 
         if not self._active_connection_state:
@@ -78,18 +86,44 @@ async def __aexit__(
             return
         await self._active_connection_state.connection.close()
 
-    def _restart_heartbeat_tasks(self, server_heartbeat: Heartbeat) -> None:
+    def _restart_background_tasks(self, server_heartbeat: Heartbeat) -> None:
         self._send_heartbeat_task.cancel()
         self._send_heartbeat_task = self._task_group.create_task(
             self._send_heartbeats_forever(server_heartbeat.want_to_receive_interval_ms)
         )
 
+    def _restart_no_message_monitor(self) -> None:
+        if self._monitor_no_message_task is not None:
+            self._monitor_no_message_task.cancel()
+        if self.no_message_restart_interval is not None:
+            self._monitor_no_message_task = self._task_group.create_task(
+                self._monitor_no_message_timeout(self.no_message_restart_interval)
+            )
+
     async def _send_heartbeats_forever(self, send_heartbeat_interval_ms: int) -> None:
         send_heartbeat_interval_seconds = send_heartbeat_interval_ms / 1000
         while True:
             await self.write_heartbeat_reconnecting()
             await asyncio.sleep(send_heartbeat_interval_seconds)
 
+    async def _monitor_no_message_timeout(self, interval: timedelta) -> None:
+        interval_seconds = interval.total_seconds()
+        while True:
+            elapsed = time.time() - self._last_message_received_time
+            if (remaining := interval_seconds - elapsed) > 0:
+                await asyncio.sleep(remaining)
+            else:
+                if connection_state := self._active_connection_state:
+                    LOGGER.warning(
+                        "no messages received for %s seconds, forcing reconnect",
+                        interval_seconds,
+                    )
+                    self._clear_active_connection_state(
+                        ConnectionLostError(reason="no messages received within timeout")
+                    )
+                    await connection_state.connection.close()
+                await asyncio.sleep(interval_seconds)
+
     async def _create_connection_to_one_server(
         self, server: ConnectionParameters
     ) -> tuple[AbstractConnection, ConnectionParameters] | None:
@@ -121,7 +155,7 @@ async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionI
         lifespan = self.lifespan_factory(
             connection=connection,
             connection_parameters=connection_parameters,
-            set_heartbeat_interval=self._restart_heartbeat_tasks,
+            set_heartbeat_interval=self._restart_background_tasks,
         )
 
         try:
@@ -152,6 +186,8 @@ async def _get_active_connection_state(self, *, is_initial_call: bool = False) -
 
                 if isinstance(connection_result, ActiveConnectionState):
                     self._active_connection_state = connection_result
+                    self._last_message_received_time = time.time()
+                    self._restart_no_message_monitor()
                     if not is_initial_call:
                         LOGGER.warning(
                             "reconnected after connection failure. connection_parameters: %s",

packages/stompman/test_stompman/conftest.py

@@ -1,6 +1,7 @@
 import asyncio
 from collections.abc import AsyncGenerator, Callable
 from dataclasses import dataclass, field
+from datetime import timedelta
 from ssl import SSLContext
 from typing import Any, Literal, Self, TypeVar
 
@@ -53,6 +54,7 @@ class EnrichedClient(stompman.Client):
     servers: list[stompman.ConnectionParameters] = field(
         default_factory=lambda: [stompman.ConnectionParameters("localhost", 12345, "login", "passcode")], kw_only=False
     )
+    no_message_restart_interval: timedelta | None = None
 
 
 @dataclass(frozen=True, kw_only=True, slots=True)
@@ -80,6 +82,7 @@ class EnrichedConnectionManager(ConnectionManager):
     write_retry_attempts: int = 3
     ssl: Literal[True] | SSLContext | None = None
     check_server_alive_interval_factor: int = 3
+    no_message_restart_interval: timedelta | None = None
 
 
 DataclassType = TypeVar("DataclassType")

packages/stompman/test_stompman/test_connection_manager.py

@@ -1,5 +1,7 @@
 import asyncio
+import time
 from collections.abc import AsyncGenerator, AsyncIterable
+from datetime import timedelta
 from ssl import SSLContext
 from typing import Literal, Self
 from unittest import mock
@@ -147,12 +149,12 @@ async def test_get_active_connection_state_lifespan_flaky_ok() -> None:
         mock.call(
             connection=BaseMockConnection(),
             connection_parameters=manager.servers[0],
-            set_heartbeat_interval=manager._restart_heartbeat_tasks,
+            set_heartbeat_interval=manager._restart_background_tasks,
         ),
         mock.call(
             connection=BaseMockConnection(),
             connection_parameters=manager.servers[0],
-            set_heartbeat_interval=manager._restart_heartbeat_tasks,
+            set_heartbeat_interval=manager._restart_background_tasks,
         ),
     ]
 
@@ -368,3 +370,48 @@ class MockConnection(BaseMockConnection):
 async def test_maybe_write_frame_ok() -> None:
     async with EnrichedConnectionManager(connection_class=BaseMockConnection) as manager:
         assert await manager.maybe_write_frame(build_dataclass(ConnectFrame))
+
+
+async def test_no_message_restart_triggers_reconnect(monkeypatch: pytest.MonkeyPatch) -> None:
+    frozen_time = [time.time()]
+    monkeypatch.setattr("time.time", lambda: frozen_time[0])
+
+    async with EnrichedConnectionManager(
+        connection_class=BaseMockConnection, no_message_restart_interval=timedelta(seconds=10)
+    ) as manager:
+        frozen_time[0] += 11
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert manager._reconnection_count >= 1
+        old_monitor_task = manager._monitor_no_message_task
+        await manager._get_active_connection_state()
+        assert manager._monitor_no_message_task is not old_monitor_task
+
+
+async def test_no_message_restart_does_not_trigger_when_messages_flow(monkeypatch: pytest.MonkeyPatch) -> None:
+    frozen_time = [time.time()]
+    monkeypatch.setattr("time.time", lambda: frozen_time[0])
+
+    async with EnrichedConnectionManager(
+        connection_class=BaseMockConnection, no_message_restart_interval=timedelta(seconds=10)
+    ) as manager:
+        initial_reconnection_count = manager._reconnection_count
+        frozen_time[0] += 5
+        manager._last_message_received_time = frozen_time[0]
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert manager._reconnection_count == initial_reconnection_count
+
+
+async def test_no_message_restart_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
+    frozen_time = [time.time()]
+    monkeypatch.setattr("time.time", lambda: frozen_time[0])
+
+    async with EnrichedConnectionManager(
+        connection_class=BaseMockConnection, no_message_restart_interval=None
+    ) as manager:
+        assert manager._monitor_no_message_task is None
+        frozen_time[0] += 99999
+        for _ in range(10):
+            await asyncio.sleep(0)
+        assert manager._reconnection_count == 0