Split heartbeat tasks and improve server heartbeat monitoring

Author: vrslev

packages/stompman/stompman/client.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import time
 from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
@@ -52,7 +53,8 @@ class Client:
     _active_subscriptions: ActiveSubscriptions = field(default_factory=dict, init=False)
     _active_transactions: set[Transaction] = field(default_factory=set, init=False)
     _exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack, init=False)
-    _heartbeat_task: asyncio.Task[None] = field(init=False)
+    _send_heartbeat_task: asyncio.Task[None] = field(init=False)
+    _check_server_heartbeat_task: asyncio.Task[None] = field(init=False)
     _listen_task: asyncio.Task[None] = field(init=False)
     _task_group: asyncio.TaskGroup = field(init=False)
     _on_heartbeat_is_async: bool = field(init=False)
@@ -68,7 +70,7 @@ def __post_init__(self) -> None:
                 disconnect_confirmation_timeout=self.disconnect_confirmation_timeout,
                 active_subscriptions=self._active_subscriptions,
                 active_transactions=self._active_transactions,
-                set_heartbeat_interval=self._restart_heartbeat_task,
+                set_heartbeat_interval=self._restart_heartbeat_tasks,
             ),
             connection_class=self.connection_class,
             connect_retry_attempts=self.connect_retry_attempts,
@@ -83,7 +85,8 @@ def __post_init__(self) -> None:
 
     async def __aenter__(self) -> Self:
         self._task_group = await self._exit_stack.enter_async_context(asyncio.TaskGroup())
-        self._heartbeat_task = self._task_group.create_task(asyncio.sleep(0))
+        self._send_heartbeat_task = self._task_group.create_task(asyncio.sleep(0))
+        self._check_server_heartbeat_task = self._task_group.create_task(asyncio.sleep(0))
         await self._exit_stack.enter_async_context(self._connection_manager)
         self._listen_task = self._task_group.create_task(self._listen_to_frames())
         return self
@@ -96,18 +99,42 @@ async def __aexit__(
                 await asyncio.Future()
         finally:
             self._listen_task.cancel()
-            self._heartbeat_task.cancel()
-            await asyncio.wait([self._listen_task, self._heartbeat_task])
+            self._send_heartbeat_task.cancel()
+            self._check_server_heartbeat_task.cancel()
+            await asyncio.wait([self._listen_task, self._send_heartbeat_task, self._check_server_heartbeat_task])
             await self._exit_stack.aclose()
 
-    def _restart_heartbeat_task(self, interval: float) -> None:
-        self._heartbeat_task.cancel()
-        self._heartbeat_task = self._task_group.create_task(self._send_heartbeats_forever(interval))
+    def _restart_heartbeat_tasks(self, server_heartbeat: Heartbeat) -> None:
+        self._send_heartbeat_task.cancel()
+        self._check_server_heartbeat_task.cancel()
+        self._send_heartbeat_task = self._task_group.create_task(
+            self._send_heartbeats_forever(server_heartbeat.want_to_receive_interval_ms)
+        )
+        self._check_server_heartbeat_task = self._task_group.create_task(
+            self._monitor_server_aliveness(server_heartbeat.will_send_interval_ms)
+        )
 
-    async def _send_heartbeats_forever(self, interval: float) -> None:
+    async def _send_heartbeats_forever(self, client_heartbeat_interval_ms: int) -> None:
         while True:
             await self._connection_manager.write_heartbeat_reconnecting()
-            await asyncio.sleep(interval)
+            await asyncio.sleep(client_heartbeat_interval_ms / 1000)
+
+    async def _monitor_server_aliveness(self, server_heartbeat_interval_ms: int) -> None:
+        server_heartbeat_interval_seconds = server_heartbeat_interval_ms / 1000
+        while True:
+            await asyncio.sleep(server_heartbeat_interval_seconds * self.check_server_alive_interval_factor)
+            last_read_time_ms = (
+                self._connection_manager._active_connection_state
+                and self._connection_manager._active_connection_state.connection.last_read_time_ms
+            )
+            if not last_read_time_ms:
+                print("No last_Read_time_ms")
+                continue
+            if (t:=time.time() - last_read_time_ms) > server_heartbeat_interval_seconds:
+                print("reset time", t)
+                self._connection_manager._active_connection_state = None
+            else:
+                print("no reset")
 
     async def _listen_to_frames(self) -> None:
         async with asyncio.TaskGroup() as task_group:

packages/stompman/stompman/connection.py

@@ -14,7 +14,7 @@
 
 @dataclass(kw_only=True)
 class AbstractConnection(Protocol):
-    last_read_time: float | None = field(init=False, default=None)
+    last_read_time_ms: float | None = field(init=False, default=None)
 
     @classmethod
     async def connect(
@@ -103,7 +103,7 @@ async def read_frames(self) -> AsyncGenerator[AnyServerFrame, None]:
                 raw_frames = await asyncio.wait_for(
                     self._read_non_empty_bytes(self.read_max_chunk_size), timeout=self.read_timeout
                 )
-            self.last_read_time = time.time()
+            self.last_read_time_ms = time.time()
 
             for frame in cast("Iterator[AnyServerFrame]", parser.parse_frames_from_chunk(raw_frames)):
                 yield frame

packages/stompman/stompman/connection_lifespan.py

@@ -1,8 +1,7 @@
 import asyncio
-from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Protocol
+from typing import Any, Callable, Protocol
 from uuid import uuid4
 
 from stompman.config import ConnectionParameters, Heartbeat
@@ -27,6 +26,8 @@ async def enter(self) -> StompProtocolConnectionIssue | None: ...
     async def exit(self) -> None: ...
 
 
+
+
 @dataclass(kw_only=True, slots=True)
 class ConnectionLifespan(AbstractConnectionLifespan):
     connection: AbstractConnection
@@ -37,7 +38,7 @@ class ConnectionLifespan(AbstractConnectionLifespan):
     disconnect_confirmation_timeout: int
     active_subscriptions: ActiveSubscriptions
     active_transactions: ActiveTransactions
-    set_heartbeat_interval: Callable[[float], Any]
+    set_heartbeat_interval: Callable[[Heartbeat], Any]
 
     async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
         await self.connection.write_frame(
@@ -73,10 +74,7 @@ async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
                 given_version=connected_frame.headers["version"], supported_version=self.protocol_version
             )
 
-        server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
-        self.set_heartbeat_interval(
-            max(self.client_heartbeat.will_send_interval_ms, server_heartbeat.want_to_receive_interval_ms) / 1000
-        )
+        self.set_heartbeat_interval(Heartbeat.from_header(connected_frame.headers["heart-beat"]))
         return None
 
     async def enter(self) -> StompProtocolConnectionIssue | None: