Enhance connection handling by introducing EstablishedConnectionResult for better heartbeat management and connection lifecycle tracking

Author: vrslev

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: ...
 
 
@@ -39,7 +44,7 @@ class ConnectionLifespan(AbstractConnectionLifespan):
     active_transactions: ActiveTransactions
     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={
@@ -73,17 +78,18 @@ async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
                 given_version=connected_frame.headers["version"], supported_version=self.protocol_version
             )
 
-        self.set_heartbeat_interval(Heartbeat.from_header(connected_frame.headers["heart-beat"]))
-        return None
+        server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
+        self.set_heartbeat_interval(server_heartbeat)
+        return EstablishedConnectionResult(server_heartbeat=server_heartbeat)
 
-    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
+    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():

packages/stompman/stompman/connection_manager.py

@@ -6,8 +6,9 @@
 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.connection_lifespan import EstablishedConnectionResult
 from stompman.errors import (
     AllServersUnavailable,
     AnyConnectionIssue,
@@ -26,6 +27,7 @@
 class ActiveConnectionState:
     connection: AbstractConnection
     lifespan: "AbstractConnectionLifespan"
+    server_heartbeat: Heartbeat
 
 
 @dataclass(kw_only=True, slots=True)
@@ -60,7 +62,9 @@ async def __aexit__(
             return
         await self._active_connection_state.connection.close()
 
-    async def _create_connection_to_one_server(self, server: ConnectionParameters) -> ActiveConnectionState | 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,
@@ -69,31 +73,35 @@ 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()):
+        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)
 
         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: