Fix weird retries (#54)

Author: vrslev

README.md

@@ -37,6 +37,7 @@ async with stompman.Client(
     connection_confirmation_timeout=2,
     disconnect_confirmation_timeout=2,
     read_timeout=2,
+    write_retry_attempts=3,
 ) as client:
     ...
 ```
@@ -107,7 +108,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` or `stompman.ConnectionLostDuringOperationError` will be raised if connection succeeds but operation (like sending a frame) leads to connection getting lost.
+- 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.
 
 ### ...and caveats
 

pyproject.toml

@@ -22,7 +22,7 @@ repository = "https://github.com/vrslev/stompman"
 [tool.uv]
 dev-dependencies = [
     "anyio~=4.4.0",
-    "mypy~=1.10.0",
+    "mypy~=1.11.1",
     "pytest-cov~=5.0.0",
     "pytest~=8.2.2",
     "ruff~=0.4.9",

stompman/__init__.py

@@ -8,12 +8,11 @@
     ConnectionManager,
 )
 from stompman.errors import (
-    ConnectionAttemptsFailedError,
     ConnectionConfirmationTimeout,
-    ConnectionLostDuringOperationError,
     ConnectionLostError,
     Error,
     FailedAllConnectAttemptsError,
+    FailedAllWriteAttemptsError,
     StompProtocolConnectionIssue,
     UnsupportedProtocolVersion,
 )
@@ -56,18 +55,17 @@
     "ConnectFrame",
     "ConnectedFrame",
     "Connection",
-    "ConnectionAttemptsFailedError",
     "ConnectionConfirmationTimeout",
     "ConnectionLifespan",
     "ConnectionLifespanFactory",
-    "ConnectionLostDuringOperationError",
     "ConnectionLostError",
     "ConnectionManager",
     "ConnectionParameters",
     "DisconnectFrame",
     "Error",
     "ErrorFrame",
     "FailedAllConnectAttemptsError",
+    "FailedAllWriteAttemptsError",
     "FrameParser",
     "Heartbeat",
     "HeartbeatFrame",

stompman/client.py

@@ -225,6 +225,7 @@ class Client:
     connect_timeout: int = 2
     read_timeout: int = 2
     read_max_chunk_size: int = 1024 * 1024
+    write_retry_attempts: int = 3
     connection_confirmation_timeout: int = 2
     disconnect_confirmation_timeout: int = 2
 
@@ -257,6 +258,7 @@ def __post_init__(self) -> None:
             connect_timeout=self.connect_timeout,
             read_timeout=self.read_timeout,
             read_max_chunk_size=self.read_max_chunk_size,
+            write_retry_attempts=self.write_retry_attempts,
         )
 
     async def __aenter__(self) -> Self:

stompman/connection_manager.py

@@ -7,11 +7,12 @@
 from stompman.config import ConnectionParameters
 from stompman.connection import AbstractConnection
 from stompman.errors import (
-    ConnectionAttemptsFailedError,
+    AllServersUnavailable,
+    AnyConnectionIssue,
     ConnectionLost,
-    ConnectionLostDuringOperationError,
     ConnectionLostError,
     FailedAllConnectAttemptsError,
+    FailedAllWriteAttemptsError,
     StompProtocolConnectionIssue,
 )
 from stompman.frames import AnyClientFrame, AnyServerFrame
@@ -44,6 +45,7 @@ class ConnectionManager:
     connect_timeout: int
     read_timeout: int
     read_max_chunk_size: int
+    write_retry_attempts: int
 
     _active_connection_state: ActiveConnectionState | None = field(default=None, init=False)
     _reconnect_lock: asyncio.Lock = field(default_factory=asyncio.Lock)
@@ -63,88 +65,84 @@ async def __aexit__(
             return
         await self._active_connection_state.connection.close()
 
-    async def _connect_to_one_server(
-        self, server: ConnectionParameters
-    ) -> tuple[AbstractConnection, ConnectionParameters] | None:
-        for attempt in range(self.connect_retry_attempts):
-            if connection := await self.connection_class.connect(
-                host=server.host,
-                port=server.port,
-                timeout=self.connect_timeout,
-                read_max_chunk_size=self.read_max_chunk_size,
-                read_timeout=self.read_timeout,
-            ):
-                return connection, server
-            await asyncio.sleep(self.connect_retry_interval * (attempt + 1))
+    async def _create_connection_to_one_server(self, server: ConnectionParameters) -> ActiveConnectionState | None:
+        if connection := await self.connection_class.connect(
+            host=server.host,
+            port=server.port,
+            timeout=self.connect_timeout,
+            read_max_chunk_size=self.read_max_chunk_size,
+            read_timeout=self.read_timeout,
+        ):
+            return ActiveConnectionState(
+                connection=connection,
+                lifespan=self.lifespan_factory(connection=connection, connection_parameters=server),
+            )
         return None
 
-    async def _connect_to_any_server(self) -> tuple[AbstractConnection, ConnectionParameters]:
+    async def _create_connection_to_any_server(self) -> ActiveConnectionState | None:
         for maybe_connection_future in asyncio.as_completed(
-            [self._connect_to_one_server(server) for server in self.servers]
+            [self._create_connection_to_one_server(server) for server in self.servers]
         ):
-            if maybe_result := await maybe_connection_future:
-                return maybe_result
-        raise FailedAllConnectAttemptsError(
-            servers=self.servers,
-            retry_attempts=self.connect_retry_attempts,
-            retry_interval=self.connect_retry_interval,
-            timeout=self.connect_timeout,
-        )
+            if connection_state := await maybe_connection_future:
+                return connection_state
+        return None
+
+    async def _connect_to_any_server(self) -> ActiveConnectionState | AnyConnectionIssue:
+        if not (active_connection_state := await self._create_connection_to_any_server()):
+            return AllServersUnavailable(servers=self.servers, timeout=self.connect_timeout)
+
+        try:
+            if connection_issue := await active_connection_state.lifespan.enter():
+                return connection_issue
+        except ConnectionLostError:
+            return ConnectionLost()
+
+        return active_connection_state
 
     async def _get_active_connection_state(self) -> ActiveConnectionState:
-        connection_issues: list[StompProtocolConnectionIssue | ConnectionLost] = []
+        if self._active_connection_state:
+            return self._active_connection_state
 
-        for _ in range(self.connect_retry_attempts):
+        connection_issues: list[AnyConnectionIssue] = []
+
+        async with self._reconnect_lock:
             if self._active_connection_state:
                 return self._active_connection_state
 
-            async with self._reconnect_lock:
-                if self._active_connection_state:
-                    return self._active_connection_state
-
-                connection, connection_parameters = await self._connect_to_any_server()
-                self._active_connection_state = ActiveConnectionState(
-                    connection=connection,
-                    lifespan=self.lifespan_factory(connection=connection, connection_parameters=connection_parameters),
-                )
-
-                try:
-                    lifespan_connection_issue: (
-                        StompProtocolConnectionIssue | ConnectionLost | None
-                    ) = await self._active_connection_state.lifespan.enter()
-                except ConnectionLostError:
-                    lifespan_connection_issue = ConnectionLost()
-                else:
-                    if lifespan_connection_issue is None:
-                        return self._active_connection_state
+            for attempt in range(self.connect_retry_attempts):
+                connection_result = await self._connect_to_any_server()
 
-                self._clear_active_connection_state()
-                connection_issues.append(lifespan_connection_issue)
+                if isinstance(connection_result, ActiveConnectionState):
+                    self._active_connection_state = connection_result
+                    return connection_result
+
+                connection_issues.append(connection_result)
+                await asyncio.sleep(self.connect_retry_interval * (attempt + 1))
 
-        raise ConnectionAttemptsFailedError(retry_attempts=self.connect_retry_attempts, issues=connection_issues)
+        raise FailedAllConnectAttemptsError(retry_attempts=self.connect_retry_attempts, issues=connection_issues)
 
     def _clear_active_connection_state(self) -> None:
         self._active_connection_state = None
 
     async def write_heartbeat_reconnecting(self) -> None:
-        for _ in range(self.connect_retry_attempts):
+        for _ in range(self.write_retry_attempts):
             connection_state = await self._get_active_connection_state()
             try:
                 return connection_state.connection.write_heartbeat()
             except ConnectionLostError:
                 self._clear_active_connection_state()
 
-        raise ConnectionLostDuringOperationError(retry_attempts=self.connect_retry_attempts)
+        raise FailedAllWriteAttemptsError(retry_attempts=self.write_retry_attempts)
 
     async def write_frame_reconnecting(self, frame: AnyClientFrame) -> None:
-        for _ in range(self.connect_retry_attempts):
+        for _ in range(self.write_retry_attempts):
             connection_state = await self._get_active_connection_state()
             try:
                 return await connection_state.connection.write_frame(frame)
             except ConnectionLostError:
                 self._clear_active_connection_state()
 
-        raise ConnectionLostDuringOperationError(retry_attempts=self.connect_retry_attempts)
+        raise FailedAllWriteAttemptsError(retry_attempts=self.write_retry_attempts)
 
     async def read_frames_reconnecting(self) -> AsyncGenerator[AnyServerFrame, None]:
         while True:

stompman/errors.py

@@ -18,14 +18,6 @@ class ConnectionLostError(Error):
     """Raised in stompman.AbstractConnection—and handled in stompman.ConnectionManager, therefore is private."""
 
 
-@dataclass(kw_only=True)
-class FailedAllConnectAttemptsError(Error):
-    servers: list["ConnectionParameters"]
-    retry_attempts: int
-    retry_interval: int
-    timeout: int
-
-
 @dataclass(frozen=True, kw_only=True, slots=True)
 class ConnectionConfirmationTimeout:
     timeout: int
@@ -42,15 +34,22 @@ class UnsupportedProtocolVersion:
 class ConnectionLost: ...
 
 
+@dataclass(frozen=True, kw_only=True, slots=True)
+class AllServersUnavailable:
+    servers: list["ConnectionParameters"]
+    timeout: int
+
+
 StompProtocolConnectionIssue = ConnectionConfirmationTimeout | UnsupportedProtocolVersion
+AnyConnectionIssue = StompProtocolConnectionIssue | ConnectionLost | AllServersUnavailable
 
 
 @dataclass(kw_only=True)
-class ConnectionAttemptsFailedError(Error):
+class FailedAllConnectAttemptsError(Error):
     retry_attempts: int
-    issues: list[StompProtocolConnectionIssue | ConnectionLost]
+    issues: list[AnyConnectionIssue]
 
 
 @dataclass(kw_only=True)
-class ConnectionLostDuringOperationError(Error):
+class FailedAllWriteAttemptsError(Error):
     retry_attempts: int

tests/conftest.py

@@ -2,7 +2,6 @@
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from typing import Any, Self, TypeVar
-from unittest import mock
 
 import pytest
 from polyfactory.factories.dataclass_factory import DataclassFactory
@@ -22,7 +21,8 @@ def anyio_backend(request: pytest.FixtureRequest) -> object:
 
 @pytest.fixture()
 def mock_sleep(monkeypatch: pytest.MonkeyPatch) -> None:  # noqa: PT004
-    monkeypatch.setattr("asyncio.sleep", mock.AsyncMock())
+    original_sleep = asyncio.sleep
+    monkeypatch.setattr("asyncio.sleep", lambda _: original_sleep(0))
 
 
 async def noop_message_handler(frame: stompman.MessageFrame) -> None: ...
@@ -74,6 +74,7 @@ class EnrichedConnectionManager(stompman.ConnectionManager):
     connect_timeout: int = 3
     read_timeout: int = 4
     read_max_chunk_size: int = 5
+    write_retry_attempts: int = 3
 
 
 DataclassType = TypeVar("DataclassType")

tests/test_client.py

@@ -21,7 +21,6 @@
     CommitFrame,
     ConnectedFrame,
     ConnectFrame,
-    ConnectionAttemptsFailedError,
     ConnectionConfirmationTimeout,
     ConnectionParameters,
     DisconnectFrame,
@@ -133,6 +132,7 @@ async def test_client_connection_lifespan_ok(monkeypatch: pytest.MonkeyPatch) ->
     assert collected_frames == [connect_frame, connected_frame, disconnect_frame, receipt_frame]
 
 
+@pytest.mark.usefixtures("mock_sleep")
 async def test_client_connection_lifespan_connection_not_confirmed(monkeypatch: pytest.MonkeyPatch) -> None:
     async def mock_wait_for(future: Coroutine[Any, Any, Any], timeout: float) -> object:
         assert timeout == connection_confirmation_timeout
@@ -151,29 +151,30 @@ async def read_frames() -> AsyncGenerator[AnyServerFrame, None]:
             yield error_frame
             await asyncio.sleep(0)
 
-    with pytest.raises(ConnectionAttemptsFailedError) as exc_info:
+    with pytest.raises(FailedAllConnectAttemptsError) as exc_info:
         await EnrichedClient(
             connection_class=MockConnection, connection_confirmation_timeout=connection_confirmation_timeout
         ).__aenter__()
 
-    assert exc_info.value == ConnectionAttemptsFailedError(
+    assert exc_info.value == FailedAllConnectAttemptsError(
         retry_attempts=3,
         issues=[ConnectionConfirmationTimeout(timeout=connection_confirmation_timeout, frames=[error_frame])] * 3,
     )
 
 
+@pytest.mark.usefixtures("mock_sleep")
 async def test_client_connection_lifespan_unsupported_protocol_version() -> None:
     given_version = FAKER.pystr()
 
-    with pytest.raises(ConnectionAttemptsFailedError) as exc_info:
+    with pytest.raises(FailedAllConnectAttemptsError) as exc_info:
         await EnrichedClient(
             connection_class=create_spying_connection(
                 [build_dataclass(ConnectedFrame, headers={"version": given_version})]
             )[0],
             connect_retry_attempts=1,
         ).__aenter__()
 
-    assert exc_info.value == ConnectionAttemptsFailedError(
+    assert exc_info.value == FailedAllConnectAttemptsError(
         retry_attempts=1,
         issues=[UnsupportedProtocolVersion(given_version=given_version, supported_version=Client.PROTOCOL_VERSION)],
     )