Refactor lifespan and errors (#53)

Author: vrslev

README.md

@@ -107,7 +107,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 handle it automatically. `stompman.FailedAllConnectAttemptsError` will be raised if all connection attempts fail. `stompman.RepeatedConnectionLostError` will be raised if connection succeeds but operation (like sending a frame) leads to connection getting lost.
+- When connection is lost, stompman will handle it automatically. `stompman.FailedAllConnectAttemptsError` will be raised if all connection attempts fail. `stompman.RepeatedConnectionLostError` or `stompman.ConnectionLostDuringOperationError` will be raised if connection succeeds but operation (like sending a frame) leads to connection getting lost.
 
 ### ...and caveats
 

stompman/__init__.py

@@ -1,14 +1,21 @@
-from stompman.client import Client, Subscription, Transaction
+from stompman.client import Client, ConnectionLifespan, Subscription, Transaction
 from stompman.config import ConnectionParameters, Heartbeat
 from stompman.connection import AbstractConnection, Connection
-from stompman.connection_manager import ActiveConnectionState, ConnectionManager
+from stompman.connection_manager import (
+    AbstractConnectionLifespan,
+    ActiveConnectionState,
+    ConnectionLifespanFactory,
+    ConnectionManager,
+)
 from stompman.errors import (
-    ConnectionConfirmationTimeoutError,
+    ConnectionAttemptsFailedError,
+    ConnectionConfirmationTimeout,
+    ConnectionLostDuringOperationError,
     ConnectionLostError,
     Error,
     FailedAllConnectAttemptsError,
-    RepeatedConnectionLostError,
-    UnsupportedProtocolVersionError,
+    StompProtocolConnectionIssue,
+    UnsupportedProtocolVersion,
 )
 from stompman.frames import (
     AbortFrame,
@@ -31,10 +38,12 @@
     SubscribeFrame,
     UnsubscribeFrame,
 )
+from stompman.serde import FrameParser, dump_frame
 
 __all__ = [
     "AbortFrame",
     "AbstractConnection",
+    "AbstractConnectionLifespan",
     "AckFrame",
     "AckMode",
     "ActiveConnectionState",
@@ -47,24 +56,30 @@
     "ConnectFrame",
     "ConnectedFrame",
     "Connection",
-    "ConnectionConfirmationTimeoutError",
+    "ConnectionAttemptsFailedError",
+    "ConnectionConfirmationTimeout",
+    "ConnectionLifespan",
+    "ConnectionLifespanFactory",
+    "ConnectionLostDuringOperationError",
     "ConnectionLostError",
     "ConnectionManager",
     "ConnectionParameters",
     "DisconnectFrame",
     "Error",
     "ErrorFrame",
     "FailedAllConnectAttemptsError",
+    "FrameParser",
     "Heartbeat",
     "HeartbeatFrame",
     "MessageFrame",
     "NackFrame",
     "ReceiptFrame",
-    "RepeatedConnectionLostError",
     "SendFrame",
+    "StompProtocolConnectionIssue",
     "SubscribeFrame",
     "Subscription",
     "Transaction",
     "UnsubscribeFrame",
-    "UnsupportedProtocolVersionError",
+    "UnsupportedProtocolVersion",
+    "dump_frame",
 ]

stompman/client.py

@@ -1,15 +1,16 @@
 import asyncio
-from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
+from collections.abc import AsyncGenerator, Callable, Coroutine
 from contextlib import AsyncExitStack, asynccontextmanager, suppress
 from dataclasses import dataclass, field
+from functools import partial
 from types import TracebackType
 from typing import ClassVar, Self
 from uuid import uuid4
 
 from stompman.config import ConnectionParameters, Heartbeat
 from stompman.connection import AbstractConnection, Connection
-from stompman.connection_manager import ConnectionManager
-from stompman.errors import ConnectionConfirmationTimeoutError, UnsupportedProtocolVersionError
+from stompman.connection_manager import AbstractConnectionLifespan, ConnectionManager
+from stompman.errors import ConnectionConfirmationTimeout, StompProtocolConnectionIssue, UnsupportedProtocolVersion
 from stompman.frames import (
     AbortFrame,
     AckFrame,
@@ -30,66 +31,93 @@
 )
 
 
-@asynccontextmanager
-async def connection_lifespan(
-    *,
-    connection: AbstractConnection,
-    connection_parameters: ConnectionParameters,
-    protocol_version: str,
-    client_heartbeat: Heartbeat,
-    connection_confirmation_timeout: int,
-    disconnect_confirmation_timeout: int,
-) -> AsyncIterator[float]:
-    await connection.write_frame(
-        ConnectFrame(
-            headers={
-                "accept-version": protocol_version,
-                "heart-beat": client_heartbeat.to_header(),
-                "host": connection_parameters.host,
-                "login": connection_parameters.login,
-                "passcode": connection_parameters.unescaped_passcode,
-            },
-        )
-    )
-    collected_frames = []
-
-    async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
-        async for frame in connection.read_frames():
-            if isinstance(frame, ConnectedFrame):
-                return frame
-            collected_frames.append(frame)
-        msg = "unreachable"  # pragma: no cover
-        raise AssertionError(msg)  # pragma: no cover
-
-    try:
-        connected_frame = await asyncio.wait_for(
-            take_connected_frame_and_collect_other_frames(), timeout=connection_confirmation_timeout
+@dataclass(kw_only=True, slots=True)
+class ConnectionLifespan(AbstractConnectionLifespan):
+    connection: AbstractConnection
+    connection_parameters: ConnectionParameters
+    protocol_version: str
+    client_heartbeat: Heartbeat
+    connection_confirmation_timeout: int
+    disconnect_confirmation_timeout: int
+    active_subscriptions: dict[str, "Subscription"]
+    active_transactions: set["Transaction"]
+    set_heartbeat_interval: Callable[[float], None]
+
+    async def _establish_connection(self) -> StompProtocolConnectionIssue | None:
+        await self.connection.write_frame(
+            ConnectFrame(
+                headers={
+                    "accept-version": self.protocol_version,
+                    "heart-beat": self.client_heartbeat.to_header(),
+                    "host": self.connection_parameters.host,
+                    "login": self.connection_parameters.login,
+                    "passcode": self.connection_parameters.unescaped_passcode,
+                },
+            )
         )
-    except TimeoutError as exception:
-        raise ConnectionConfirmationTimeoutError(
-            timeout=connection_confirmation_timeout, frames=collected_frames
-        ) from exception
-
-    if connected_frame.headers["version"] != protocol_version:
-        raise UnsupportedProtocolVersionError(
-            given_version=connected_frame.headers["version"], supported_version=protocol_version
+        collected_frames = []
+
+        async def take_connected_frame_and_collect_other_frames() -> ConnectedFrame:
+            async for frame in self.connection.read_frames():
+                if isinstance(frame, ConnectedFrame):
+                    return frame
+                collected_frames.append(frame)
+            msg = "unreachable"  # pragma: no cover
+            raise AssertionError(msg)  # pragma: no cover
+
+        try:
+            connected_frame = await asyncio.wait_for(
+                take_connected_frame_and_collect_other_frames(), timeout=self.connection_confirmation_timeout
+            )
+        except TimeoutError:
+            return ConnectionConfirmationTimeout(timeout=self.connection_confirmation_timeout, frames=collected_frames)
+
+        if connected_frame.headers["version"] != self.protocol_version:
+            return UnsupportedProtocolVersion(
+                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
         )
+        return None
 
-    server_heartbeat = Heartbeat.from_header(connected_frame.headers["heart-beat"])
-    heartbeat_interval = (
-        max(client_heartbeat.will_send_interval_ms, server_heartbeat.want_to_receive_interval_ms) / 1000
-    )
-    yield heartbeat_interval
+    async def _resubscribe(self) -> None:
+        for subscription in self.active_subscriptions.values():
+            await self.connection.write_frame(
+                SubscribeFrame(
+                    headers={"id": subscription.id, "destination": subscription.destination, "ack": subscription.ack}
+                )
+            )
 
-    await connection.write_frame(DisconnectFrame(headers={"receipt": _make_receipt_id()}))
+    async def _commit_pending_transactions(self) -> None:
+        for transaction in self.active_transactions:
+            for frame in transaction.sent_frames:
+                await self.connection.write_frame(frame)
+            await self.connection.write_frame(CommitFrame(headers={"transaction": transaction.id}))
+        self.active_transactions.clear()
 
-    async def take_receipt_frame() -> None:
-        async for frame in connection.read_frames():
-            if isinstance(frame, ReceiptFrame):
-                break
+    async def enter(self) -> StompProtocolConnectionIssue | None:
+        if connection_issue := await self._establish_connection():
+            return connection_issue
+        await self._resubscribe()
+        await self._commit_pending_transactions()
+        return None
 
-    with suppress(TimeoutError):
-        await asyncio.wait_for(take_receipt_frame(), timeout=disconnect_confirmation_timeout)
+    async def exit(self) -> None:
+        for subscription in self.active_subscriptions.copy().values():
+            await subscription.unsubscribe()
+
+        await self.connection.write_frame(DisconnectFrame(headers={"receipt": _make_receipt_id()}))
+
+        async def take_receipt_frame() -> None:
+            async for frame in self.connection.read_frames():
+                if isinstance(frame, ReceiptFrame):
+                    break
+
+        with suppress(TimeoutError):
+            await asyncio.wait_for(take_receipt_frame(), timeout=self.disconnect_confirmation_timeout)
 
 
 def _make_receipt_id() -> str:
@@ -146,21 +174,6 @@ def _make_subscription_id() -> str:
     return str(uuid4())
 
 
-@asynccontextmanager
-async def subscriptions_lifespan(
-    *, connection: AbstractConnection, active_subscriptions: dict[str, Subscription]
-) -> AsyncIterator[None]:
-    for subscription in active_subscriptions.values():
-        await connection.write_frame(
-            SubscribeFrame(
-                headers={"id": subscription.id, "destination": subscription.destination, "ack": subscription.ack}
-            )
-        )
-    yield
-    for subscription in active_subscriptions.copy().values():
-        await subscription.unsubscribe()
-
-
 @dataclass(kw_only=True, slots=True, unsafe_hash=True)
 class Transaction:
     id: str = field(default_factory=lambda: _make_transaction_id(), init=False)  # noqa: PLW0108
@@ -198,14 +211,6 @@ def _make_transaction_id() -> str:
     return str(uuid4())
 
 
-async def commit_pending_transactions(*, connection: AbstractConnection, active_transactions: set[Transaction]) -> None:
-    for transaction in active_transactions:
-        for frame in transaction.sent_frames:
-            await connection.write_frame(frame)
-        await connection.write_frame(CommitFrame(headers={"transaction": transaction.id}))
-    active_transactions.clear()
-
-
 @dataclass(kw_only=True, slots=True)
 class Client:
     PROTOCOL_VERSION: ClassVar = "1.2"  # https://stomp.github.io/stomp-specification-1.2.html
@@ -236,7 +241,16 @@ class Client:
     def __post_init__(self) -> None:
         self._connection_manager = ConnectionManager(
             servers=self.servers,
-            lifespan=self._lifespan,
+            lifespan_factory=partial(
+                ConnectionLifespan,
+                protocol_version=self.PROTOCOL_VERSION,
+                client_heartbeat=self.heartbeat,
+                connection_confirmation_timeout=self.connection_confirmation_timeout,
+                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,
             connect_retry_interval=self.connect_retry_interval,
@@ -264,23 +278,6 @@ async def __aexit__(
             await asyncio.wait([self._listen_task, self._heartbeat_task])
             await self._exit_stack.aclose()
 
-    @asynccontextmanager
-    async def _lifespan(
-        self, connection: AbstractConnection, connection_parameters: ConnectionParameters
-    ) -> AsyncGenerator[None, None]:
-        async with connection_lifespan(
-            connection=connection,
-            connection_parameters=connection_parameters,
-            protocol_version=self.PROTOCOL_VERSION,
-            client_heartbeat=self.heartbeat,
-            connection_confirmation_timeout=self.connection_confirmation_timeout,
-            disconnect_confirmation_timeout=self.disconnect_confirmation_timeout,
-        ) as heartbeat_interval:
-            self._restart_heartbeat_task(heartbeat_interval)
-            async with subscriptions_lifespan(connection=connection, active_subscriptions=self._active_subscriptions):
-                await commit_pending_transactions(connection=connection, active_transactions=self._active_transactions)
-                yield
-
     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))