Enhance connection management with server alive check and factor in heartbeat calculation

Author: vrslev

packages/stompman/stompman/client.py

@@ -79,6 +79,7 @@ def __post_init__(self) -> None:
             read_timeout=self.read_timeout,
             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,
             ssl=self.ssl,
         )
         self._on_heartbeat_is_async = inspect.iscoroutinefunction(self.on_heartbeat) if self.on_heartbeat else False
@@ -130,7 +131,7 @@ async def _monitor_server_aliveness(self, server_heartbeat_interval_ms: int) ->
             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:
+            if (t := time.time() - last_read_time_ms) > server_heartbeat_interval_seconds:
                 print("reset time", t)
                 self._connection_manager._active_connection_state = None
             else:

packages/stompman/stompman/connection_lifespan.py

@@ -1,7 +1,8 @@
 import asyncio
+from collections.abc import Callable
 from contextlib import suppress
 from dataclasses import dataclass
-from typing import Any, Callable, Protocol
+from typing import Any, Protocol
 from uuid import uuid4
 
 from stompman.config import ConnectionParameters, Heartbeat
@@ -26,8 +27,6 @@ async def enter(self) -> StompProtocolConnectionIssue | None: ...
     async def exit(self) -> None: ...
 
 
-
-
 @dataclass(kw_only=True, slots=True)
 class ConnectionLifespan(AbstractConnectionLifespan):
     connection: AbstractConnection

packages/stompman/stompman/connection_manager.py

@@ -1,4 +1,5 @@
 import asyncio
+import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from ssl import SSLContext
@@ -39,6 +40,7 @@ class ConnectionManager:
     read_timeout: int
     read_max_chunk_size: int
     write_retry_attempts: int
+    check_server_alive_interval_factor: int
 
     _active_connection_state: ActiveConnectionState | None = field(default=None, init=False)
     _reconnect_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
@@ -118,6 +120,15 @@ async def _get_active_connection_state(self) -> ActiveConnectionState:
     def _clear_active_connection_state(self) -> None:
         self._active_connection_state = None
 
+    def is_alive(self) -> bool:
+        if not self._connection_manager._active_connection_state:
+            return False
+        if not (last_read_time_ms := self._connection_manager._active_connection_state.connection.last_read_time_ms):
+            return True
+        if (time.time() - last_read_time_ms) > server_heartbeat_interval_seconds:
+            ...
+        return None
+
     async def write_heartbeat_reconnecting(self) -> None:
         for _ in range(self.write_retry_attempts):
             connection_state = await self._get_active_connection_state()