Remove Client(..., on_heartbeat=...), add Client.is_alive(), improve FastStream health check (#109)

Author: vrslev

Justfile

@@ -25,7 +25,7 @@ test *args:
 run-artemis:
     #!/bin/bash
     trap 'echo; docker compose down --remove-orphans' EXIT
-    docker compose run --service-ports artemis
+    docker compose run --service-ports activemq-artemis
 
 run-consumer:
     uv run examples/consumer.py

README.md

@@ -24,13 +24,13 @@ async with stompman.Client(
         stompman.ConnectionParameters(host="172.0.0.1", port=61616, login="user2", passcode="passcode2"),
     ],
 
-    # Handlers:
-    on_error_frame=lambda error_frame: print(error_frame.body),
-    on_heartbeat=lambda: print("Server sent a heartbeat"),  # also can be async
 
     # SSL — can be either `None` (default), `True`, or `ssl.SSLContext'
     ssl=None,
 
+    # Error frame handler:
+    on_error_frame=lambda error_frame: print(error_frame.body),
+
     # Optional parameters with sensible defaults:
     heartbeat=stompman.Heartbeat(will_send_interval_ms=1000, want_to_receive_interval_ms=1000),
     connect_retry_attempts=3,
@@ -40,6 +40,7 @@ async with stompman.Client(
     disconnect_confirmation_timeout=2,
     read_timeout=2,
     write_retry_attempts=3,
+    check_server_alive_interval_factor=3,
 ) as client:
     ...
 ```
@@ -131,6 +132,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.
+- To implement health checks, use `stompman.Client.is_alive()` — it will return `True` if everything is OK and `False` if server is not responding.
 
 ### ...and caveats
 

packages/faststream-stomp/faststream_stomp/broker.py

@@ -114,7 +114,7 @@ async def ping(self, timeout: float | None = None) -> bool:
                 if cancel_scope.cancel_called:
                     return False
 
-                if self._connection._connection_manager._active_connection_state:  # noqa: SLF001
+                if self._connection.is_alive():
                     return True
 
                 await anyio.sleep(sleep_time)  # pragma: no cover

packages/stompman/stompman/client.py

@@ -1,5 +1,4 @@
 import asyncio
-import inspect
 from collections.abc import AsyncGenerator, Awaitable, Callable, Coroutine
 from contextlib import AsyncExitStack, asynccontextmanager
 from dataclasses import dataclass, field
@@ -31,7 +30,6 @@ class Client:
 
     servers: list[ConnectionParameters] = field(kw_only=False)
     on_error_frame: Callable[[ErrorFrame], Any] | None = None
-    on_heartbeat: Callable[[], Any] | Callable[[], Awaitable[Any]] | None = None
 
     heartbeat: Heartbeat = field(default=Heartbeat(1000, 1000))
     ssl: Literal[True] | SSLContext | None = None
@@ -43,17 +41,17 @@ class Client:
     write_retry_attempts: int = 3
     connection_confirmation_timeout: int = 2
     disconnect_confirmation_timeout: int = 2
+    check_server_alive_interval_factor: int = 3
+    """Client will check if server alive `server heartbeat interval` times `interval factor`"""
 
     connection_class: type[AbstractConnection] = Connection
 
     _connection_manager: ConnectionManager = field(init=False)
     _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)
     _listen_task: asyncio.Task[None] = field(init=False)
     _task_group: asyncio.TaskGroup = field(init=False)
-    _on_heartbeat_is_async: bool = field(init=False)
 
     def __post_init__(self) -> None:
         self._connection_manager = ConnectionManager(
@@ -66,7 +64,6 @@ 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,
             ),
             connection_class=self.connection_class,
             connect_retry_attempts=self.connect_retry_attempts,
@@ -75,13 +72,12 @@ 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
 
     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))
         await self._exit_stack.enter_async_context(self._connection_manager)
         self._listen_task = self._task_group.create_task(self._listen_to_frames())
         return self
@@ -94,19 +90,9 @@ async def __aexit__(
                 await asyncio.Future()
         finally:
             self._listen_task.cancel()
-            self._heartbeat_task.cancel()
-            await asyncio.wait([self._listen_task, self._heartbeat_task])
+            await asyncio.wait([self._listen_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))
-
-    async def _send_heartbeats_forever(self, interval: float) -> None:
-        while True:
-            await self._connection_manager.write_heartbeat_reconnecting()
-            await asyncio.sleep(interval)
-
     async def _listen_to_frames(self) -> None:
         async with asyncio.TaskGroup() as task_group:
             async for frame in self._connection_manager.read_frames_reconnecting():
@@ -125,14 +111,7 @@ async def _listen_to_frames(self) -> None:
                     case ErrorFrame():
                         if self.on_error_frame:
                             self.on_error_frame(frame)
-                    case HeartbeatFrame():
-                        if self.on_heartbeat is None:
-                            pass
-                        elif self._on_heartbeat_is_async:
-                            task_group.create_task(self.on_heartbeat())  # type: ignore[arg-type]
-                        else:
-                            self.on_heartbeat()
-                    case ConnectedFrame() | ReceiptFrame():
+                    case HeartbeatFrame() | ConnectedFrame() | ReceiptFrame():
                         pass
 
     async def send(
@@ -192,3 +171,8 @@ async def subscribe_with_manual_ack(
         )
         await subscription._subscribe()  # noqa: SLF001
         return subscription
+
+    def is_alive(self) -> bool:
+        return (
+            self._connection_manager._active_connection_state or False  # noqa: SLF001
+        ) and self._connection_manager._active_connection_state.is_alive()  # noqa: SLF001

packages/stompman/stompman/connection.py

@@ -1,8 +1,9 @@
 import asyncio
 import socket
+import time
 from collections.abc import AsyncGenerator, Generator, Iterator
 from contextlib import contextmanager, suppress
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from ssl import SSLContext
 from typing import Literal, Protocol, Self, cast
 
@@ -13,6 +14,8 @@
 
 @dataclass(kw_only=True)
 class AbstractConnection(Protocol):
+    last_read_time: float | None = field(init=False, default=None)
+
     @classmethod
     async def connect(
         cls,
@@ -100,6 +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()
 
             for frame in cast("Iterator[AnyServerFrame]", parser.parse_frames_from_chunk(raw_frames)):
                 yield frame

packages/stompman/stompman/connection_lifespan.py

@@ -22,8 +22,13 @@
 from stompman.transaction import ActiveTransactions, commit_pending_transactions
 
 
+@dataclass(frozen=True, kw_only=True, slots=True)
+class EstablishedConnectionResult:
+    server_heartbeat: Heartbeat
+
+
 class AbstractConnectionLifespan(Protocol):
-    async def enter(self) -> StompProtocolConnectionIssue | None: ...
+    async def enter(self) -> EstablishedConnectionResult | StompProtocolConnectionIssue: ...
     async def exit(self) -> None: ...
 
 
@@ -37,9 +42,9 @@ 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:
+    async def _establish_connection(self) -> EstablishedConnectionResult | StompProtocolConnectionIssue:
         await self.connection.write_frame(
             ConnectFrame(
                 headers={
@@ -74,19 +79,17 @@ async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
             )
 
         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
-        )
-        return None
-
-    async def enter(self) -> StompProtocolConnectionIssue | None:
-        if connection_issue := await self._establish_connection():
-            return connection_issue
-        await resubscribe_to_active_subscriptions(
-            connection=self.connection, active_subscriptions=self.active_subscriptions
-        )
-        await commit_pending_transactions(connection=self.connection, active_transactions=self.active_transactions)
-        return None
+        self.set_heartbeat_interval(server_heartbeat)
+        return EstablishedConnectionResult(server_heartbeat=server_heartbeat)
+
+    async def enter(self) -> EstablishedConnectionResult | StompProtocolConnectionIssue:
+        connection_result = await self._establish_connection()
+        if isinstance(connection_result, EstablishedConnectionResult):
+            await resubscribe_to_active_subscriptions(
+                connection=self.connection, active_subscriptions=self.active_subscriptions
+            )
+            await commit_pending_transactions(connection=self.connection, active_transactions=self.active_transactions)
+        return connection_result
 
     async def _take_receipt_frame(self) -> None:
         async for frame in self.connection.read_frames():
@@ -107,5 +110,9 @@ def _make_receipt_id() -> str:
 
 class ConnectionLifespanFactory(Protocol):
     def __call__(
-        self, *, connection: AbstractConnection, connection_parameters: ConnectionParameters
+        self,
+        *,
+        connection: AbstractConnection,
+        connection_parameters: ConnectionParameters,
+        set_heartbeat_interval: Callable[[Heartbeat], Any],
     ) -> AbstractConnectionLifespan: ...

packages/stompman/stompman/connection_manager.py

@@ -1,11 +1,12 @@
 import asyncio
+import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from ssl import SSLContext
 from types import TracebackType
 from typing import TYPE_CHECKING, Literal, Self
 
-from stompman.config import ConnectionParameters
+from stompman.config import ConnectionParameters, Heartbeat
 from stompman.connection import AbstractConnection
 from stompman.errors import (
     AllServersUnavailable,
@@ -25,6 +26,12 @@
 class ActiveConnectionState:
     connection: AbstractConnection
     lifespan: "AbstractConnectionLifespan"
+    server_heartbeat: Heartbeat
+
+    def is_alive(self) -> bool:
+        if not (last_read_time := self.connection.last_read_time):
+            return True
+        return (self.server_heartbeat.will_send_interval_ms / 1000) > (time.time() - last_read_time)
 
 
 @dataclass(kw_only=True, slots=True)
@@ -39,17 +46,29 @@ 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)
+    _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)
+    _check_server_heartbeat_task: asyncio.Task[None] = field(init=False, repr=False)
 
     async def __aenter__(self) -> Self:
+        await self._task_group.__aenter__()
+        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))
         self._active_connection_state = await self._get_active_connection_state()
         return self
 
     async def __aexit__(
         self, exc_type: type[BaseException] | None, exc_value: BaseException | None, traceback: TracebackType | None
     ) -> None:
+        self._send_heartbeat_task.cancel()
+        self._check_server_heartbeat_task.cancel()
+        await asyncio.wait([self._send_heartbeat_task, self._check_server_heartbeat_task])
+        await self._task_group.__aexit__(exc_type, exc_value, traceback)
+
         if not self._active_connection_state:
             return
         try:
@@ -58,7 +77,34 @@ async def __aexit__(
             return
         await self._active_connection_state.connection.close()
 
-    async def _create_connection_to_one_server(self, server: ConnectionParameters) -> ActiveConnectionState | None:
+    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._check_server_heartbeat_forever(server_heartbeat.will_send_interval_ms)
+        )
+
+    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 _check_server_heartbeat_forever(self, receive_heartbeat_interval_ms: int) -> None:
+        receive_heartbeat_interval_seconds = receive_heartbeat_interval_ms / 1000
+        while True:
+            await asyncio.sleep(receive_heartbeat_interval_seconds * self.check_server_alive_interval_factor)
+            if not self._active_connection_state:
+                continue
+            if not self._active_connection_state.is_alive():
+                self._active_connection_state = None
+
+    async def _create_connection_to_one_server(
+        self, server: ConnectionParameters
+    ) -> tuple[AbstractConnection, ConnectionParameters] | None:
         if connection := await self.connection_class.connect(
             host=server.host,
             port=server.port,
@@ -67,31 +113,41 @@ async def _create_connection_to_one_server(self, server: ConnectionParameters) -
             read_timeout=self.read_timeout,
             ssl=self.ssl,
         ):
-            return ActiveConnectionState(
-                connection=connection,
-                lifespan=self.lifespan_factory(connection=connection, connection_parameters=server),
-            )
+            return (connection, server)
         return None
 
-    async def _create_connection_to_any_server(self) -> ActiveConnectionState | None:
+    async def _create_connection_to_any_server(self) -> tuple[AbstractConnection, ConnectionParameters] | None:
         for maybe_connection_future in asyncio.as_completed(
             [self._create_connection_to_one_server(server) for server in self.servers]
         ):
-            if connection_state := await maybe_connection_future:
-                return connection_state
+            if connection_and_server := await maybe_connection_future:
+                return connection_and_server
         return None
 
     async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionIssue:
-        if not (active_connection_state := await self._create_connection_to_any_server()):
+        from stompman.connection_lifespan import EstablishedConnectionResult  # noqa: PLC0415
+
+        if not (connection_and_server := await self._create_connection_to_any_server()):
             return AllServersUnavailable(servers=self.servers, timeout=self.connect_timeout)
+        connection, connection_parameters = connection_and_server
+        lifespan = self.lifespan_factory(
+            connection=connection,
+            connection_parameters=connection_parameters,
+            set_heartbeat_interval=self._restart_heartbeat_tasks,
+        )
 
         try:
-            if connection_issue := await active_connection_state.lifespan.enter():
-                return connection_issue
+            connection_result = await lifespan.enter()
         except ConnectionLostError:
             return ConnectionLost()
 
-        return active_connection_state
+        return (
+            ActiveConnectionState(
+                connection=connection, lifespan=lifespan, server_heartbeat=connection_result.server_heartbeat
+            )
+            if isinstance(connection_result, EstablishedConnectionResult)
+            else connection_result
+        )
 
     async def _get_active_connection_state(self) -> ActiveConnectionState:
         if self._active_connection_state: