Refactor heartbeat management by moving tasks to ConnectionManager and updating ConnectionLifespanFactory

Author: vrslev

packages/stompman/stompman/client.py

@@ -50,8 +50,6 @@ 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)
-    _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)
 
@@ -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_tasks,
             ),
             connection_class=self.connection_class,
             connect_retry_attempts=self.connect_retry_attempts,
@@ -81,9 +78,6 @@ def __post_init__(self) -> None:
 
     async def __aenter__(self) -> Self:
         self._task_group = await self._exit_stack.enter_async_context(asyncio.TaskGroup())
-        # TODO: move this tasks to lifespan
-        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,36 +90,9 @@ async def __aexit__(
                 await asyncio.Future()
         finally:
             self._listen_task.cancel()
-            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 asyncio.wait([self._listen_task])
             await self._exit_stack.aclose()
 
-    def _restart_heartbeat_tasks(self, server_heartbeat: Heartbeat) -> None:
-        self._send_heartbeat_task.cancel()
-        self._check_server_heartbeat_task.cancel()
-        print(server_heartbeat)
-        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, client_heartbeat_interval_ms: int) -> None:
-        while True:
-            await self._connection_manager.write_heartbeat_reconnecting()
-            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)
-            if not self._connection_manager._active_connection_state:
-                continue
-            if not self._connection_manager._active_connection_state.is_alive():
-                self._connection_manager._active_connection_state = None
-
     async def _listen_to_frames(self) -> None:
         async with asyncio.TaskGroup() as task_group:
             async for frame in self._connection_manager.read_frames_reconnecting():

packages/stompman/stompman/connection_lifespan.py

@@ -110,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

@@ -49,15 +49,26 @@ class ConnectionManager:
     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:
@@ -66,6 +77,31 @@ async def __aexit__(
             return
         await self._active_connection_state.connection.close()
 
+    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:
@@ -94,7 +130,11 @@ async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionI
         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)
+        lifespan = self.lifespan_factory(
+            connection=connection,
+            connection_parameters=connection_parameters,
+            set_heartbeat_interval=self._restart_heartbeat_tasks,
+        )
 
         try:
             connection_result = await lifespan.enter()

packages/stompman/test_stompman/conftest.py

@@ -1,12 +1,13 @@
 import asyncio
-from collections.abc import AsyncGenerator
+from collections.abc import AsyncGenerator, Callable
 from dataclasses import dataclass, field
 from ssl import SSLContext
 from typing import Any, Literal, Self, TypeVar
 
 import pytest
 import stompman
 from polyfactory.factories.dataclass_factory import DataclassFactory
+from stompman.config import Heartbeat
 from stompman.connection import AbstractConnection
 from stompman.connection_lifespan import AbstractConnectionLifespan, EstablishedConnectionResult
 from stompman.connection_manager import ConnectionManager
@@ -58,6 +59,7 @@ class EnrichedClient(stompman.Client):
 class NoopLifespan(AbstractConnectionLifespan):
     connection: AbstractConnection
     connection_parameters: stompman.ConnectionParameters
+    set_heartbeat_interval: Callable[[Heartbeat], Any]
 
     async def enter(self) -> EstablishedConnectionResult | stompman.StompProtocolConnectionIssue:
         return EstablishedConnectionResult(server_heartbeat=stompman.Heartbeat(1000, 1000))

packages/stompman/test_stompman/test_connection_manager.py

@@ -140,8 +140,16 @@ async def test_get_active_connection_state_lifespan_flaky_ok() -> None:
 
     assert enter.mock_calls == [mock.call(), mock.call()]
     assert lifespan_factory.mock_calls == [
-        mock.call(connection=BaseMockConnection(), connection_parameters=manager.servers[0]),
-        mock.call(connection=BaseMockConnection(), connection_parameters=manager.servers[0]),
+        mock.call(
+            connection=BaseMockConnection(),
+            connection_parameters=manager.servers[0],
+            set_heartbeat_interval=manager._restart_heartbeat_tasks,
+        ),
+        mock.call(
+            connection=BaseMockConnection(),
+            connection_parameters=manager.servers[0],
+            set_heartbeat_interval=manager._restart_heartbeat_tasks,
+        ),
     ]