diff --git a/README.md b/README.md index 288bf51..f0102ea 100644 --- a/README.md +++ b/README.md @@ -24,14 +24,10 @@ async with stompman.Client( stompman.ConnectionParameters(host="171.0.0.1", port=61616, login="user1", passcode="passcode1"), stompman.ConnectionParameters(host="172.0.0.1", port=61616, login="user2", passcode="passcode2"), ], - - # 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, @@ -81,6 +77,41 @@ async with client.begin() as transaction: await asyncio.sleep(0.1) ``` +#### Confirming sends + +By default `client.send()` returns once the frame has been written to the socket. That says nothing about whether the +broker accepted it. Pass `receipt_timeout` to wait for the broker's +[STOMP receipt](https://stomp.github.io/stomp-specification-1.2.html#RECEIPT) instead: + +```python +await client.send(b"hi there!", destination="DLQ", headers={"persistent": "true"}, receipt_timeout=3.0) +``` + +The default `receipt_timeout=None` preserves the existing write-only behavior. A timeout must be finite and positive. +It covers writing `SEND` and waiting for its receipt, after a connection is available. The client generates the +`receipt` header itself, so passing your own `receipt` header with `receipt_timeout` is an error. + +A failure raises `stompman.SendReceiptError`, with `reason` equal to `rejected`, `timeout`, or `connection_lost`. +Raw broker error frames are available through `error.frame`, but are excluded from the exception's representation. + +Only `rejected` is a definitive answer. On `timeout` and `connection_lost` the outcome is **ambiguous**: the frame is +fully buffered before the socket is drained, so the broker may well have accepted the message even though the +confirmation was lost. A confirmed send still waits for a connection through `connect_retry_attempts`, but the write +itself is attempted exactly once: `write_retry_attempts` does not apply, because a connection can fail after the frame +reached the broker and no connection type can tell the two apart. Nothing is silently replayed on a new connection — +retrying is your decision. Give each application event a stable ID and deduplicate in the broker or downstream +consumer; `JMSCorrelationID` alone does not deduplicate anything. + +A receipt means the broker took ownership of the message. It does not mean any consumer has received or processed it. + +Concurrent confirmed sends never consume each other's receipts, and a receipt is bound to the connection it was issued +on, so a late receipt from a previous connection cannot confirm a newer send. An `ERROR` without `receipt-id` fails all +pending confirmations on that connection, the same policy as subscriptions. Broker errors still reach the +`on_error_frame` callback. + +Confirmation is not available for `transaction.send()`: a receipt for a transactional `SEND` says nothing about whether +the `COMMIT` succeeded. + ### Listening for Messages Now, let's subscribe to a destination and listen for messages: @@ -125,7 +156,13 @@ await client.subscribe("DLQ", handle_message_from_dlq, ack="auto", on_suppressed You can pass custom headers to `client.subscribe()`: ```python -await client.subscribe("DLQ", handle_message_from_dlq, ack="client", headers={"selector": "location = 'Europe'"}, on_suppressed_exception=print) +await client.subscribe( + "DLQ", + handle_message_from_dlq, + ack="client", + headers={"selector": "location = 'Europe'"}, + on_suppressed_exception=print, +) ``` #### Handling ACK/NACKs yourself @@ -137,6 +174,7 @@ async def handle_message_from_dlq(message_frame: stompman.AckableMessageFrame) - print(message_frame.body) await message_frame.ack() + await client.subscribe_with_manual_ack("DLQ", handle_message_from_dlq, ack="client") ``` diff --git a/packages/faststream-stomp/README.md b/packages/faststream-stomp/README.md index 8c25db0..41a8dc8 100644 --- a/packages/faststream-stomp/README.md +++ b/packages/faststream-stomp/README.md @@ -64,6 +64,32 @@ Use `reply_add_content_length` on `subscriber()` when automatic replies need a d are available for batch publishing and delayed `StompRoute`/`StompRoutePublisher` registrations. Per-call settings override publisher or subscriber reply defaults, which in turn override the broker default. +### Confirming publication + +By default `publish()` returns once the frame has been written to the socket, which does not mean the broker +accepted the message. Pass `receipt_timeout` to wait for the broker's STOMP `RECEIPT` instead. It follows the same precedence as +`add_content_length`: per-call, then publisher, then broker default. + +```python +broker = faststream_stomp.StompBroker(stompman.Client([server]), receipt_timeout=3.0) +publisher = broker.publisher("events", receipt_timeout=5.0) + +await broker.publish("text message", "events", receipt_timeout=1.0) +await publisher.publish("another message") +``` + +The default is `None` everywhere, which preserves the existing behavior. A failed confirmation raises +`stompman.SendReceiptError` from `publish()`, with `reason` equal to `rejected`, `timeout`, or `connection_lost`. + +`timeout` and `connection_lost` are ambiguous outcomes: the broker may have accepted the message even though its +receipt was lost. Nothing is replayed automatically: give each application event a stable ID and deduplicate in the +broker or downstream consumer if you retry. `JMSCorrelationID` alone does not deduplicate anything. And a receipt only +means the broker took ownership: it does not mean a consumer has processed the message. + +`receipt_timeout` is not available for batch publishing, which publishes inside a transaction: a receipt for a +transactional `SEND` says nothing about whether the `COMMIT` succeeded. A broker or publisher default is ignored there +rather than silently claiming confirmation. + ### Caveats - When exception is raised in consumer handler, the message will be nacked (FastStream doesn't do this by default) diff --git a/packages/faststream-stomp/faststream_stomp/broker.py b/packages/faststream-stomp/faststream_stomp/broker.py index 48e0e2a..58e2879 100644 --- a/packages/faststream-stomp/faststream_stomp/broker.py +++ b/packages/faststream-stomp/faststream_stomp/broker.py @@ -102,6 +102,7 @@ def __init__( graceful_timeout: float | None = 15.0, routers: Sequence[Registrator[stompman.MessageFrame]] = (), add_content_length: bool = True, + receipt_timeout: float | None = None, # Logging args logger: LoggerProto | None = EMPTY, log_level: int = logging.INFO, @@ -129,6 +130,7 @@ def __init__( client=client, serializer=fd_config._serializer, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ), client=client, ) @@ -199,6 +201,7 @@ async def publish( correlation_id: str | None = None, headers: dict[str, str] | None = None, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> None: publish_command = StompPublishCommand( message, @@ -207,6 +210,7 @@ async def publish( correlation_id=correlation_id, headers=headers, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ) return typing.cast("None", await self._basic_publish(publish_command, producer=self.config.producer)) @@ -218,6 +222,7 @@ async def request( # type: ignore[override] correlation_id: str | None = None, headers: dict[str, str] | None = None, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> Any: # noqa: ANN401 publish_command = StompPublishCommand( message, @@ -226,6 +231,7 @@ async def request( # type: ignore[override] correlation_id=correlation_id, headers=headers, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ) return await self._basic_request(publish_command, producer=self.config.producer) diff --git a/packages/faststream-stomp/faststream_stomp/models.py b/packages/faststream-stomp/faststream_stomp/models.py index 5bf50da..26da5e5 100644 --- a/packages/faststream-stomp/faststream_stomp/models.py +++ b/packages/faststream-stomp/faststream_stomp/models.py @@ -56,6 +56,7 @@ def __init__( correlation_id: str | None = None, headers: dict[str, Any] | None = None, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> None: super().__init__( body, @@ -67,6 +68,7 @@ def __init__( headers=headers, ) self.add_content_length = add_content_length + self.receipt_timeout = receipt_timeout @classmethod def from_cmd( @@ -75,10 +77,14 @@ def from_cmd( *, batch: bool = False, # noqa: ARG003 add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> Self: messages = cmd.batch_bodies - if isinstance(cmd, StompPublishCommand) and cmd.add_content_length is not None: - add_content_length = cmd.add_content_length + if isinstance(cmd, StompPublishCommand): + if cmd.add_content_length is not None: + add_content_length = cmd.add_content_length + if cmd.receipt_timeout is not None: + receipt_timeout = cmd.receipt_timeout return cls( *messages, _publish_type=cmd.publish_type, @@ -87,6 +93,7 @@ def from_cmd( correlation_id=cmd.correlation_id, headers=cmd.headers, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ) @@ -137,6 +144,7 @@ class StompPublisherSpecificationConfig(_StompBasePublisherConfig, PublisherSpec class StompPublisherUsecaseConfig(_StompBasePublisherConfig, PublisherUsecaseConfig): _outer_config: BrokerConfigWithStompClient add_content_length: bool | None + receipt_timeout: float | None @property def full_destination(self) -> str: diff --git a/packages/faststream-stomp/faststream_stomp/publisher.py b/packages/faststream-stomp/faststream_stomp/publisher.py index e8a4f70..73d75ec 100644 --- a/packages/faststream-stomp/faststream_stomp/publisher.py +++ b/packages/faststream-stomp/faststream_stomp/publisher.py @@ -31,10 +31,12 @@ def __init__( client: stompman.Client, serializer: SerializerProto | None, add_content_length: bool = True, + receipt_timeout: float | None = None, ) -> None: self.client = client self.serializer = serializer self.add_content_length = add_content_length + self.receipt_timeout = receipt_timeout self.codec: CodecProto = DefaultCodec() async def publish(self, cmd: StompPublishCommand) -> None: @@ -45,6 +47,7 @@ async def publish(self, cmd: StompPublishCommand) -> None: content_type=content_type, add_content_length=self._resolve_add_content_length(cmd), headers=_make_headers_for_publish(cmd), + receipt_timeout=self._resolve_receipt_timeout(cmd), ) async def request(self, cmd: StompPublishCommand) -> NoReturn: @@ -66,6 +69,9 @@ async def publish_batch(self, cmd: StompPublishCommand) -> None: def _resolve_add_content_length(self, cmd: StompPublishCommand) -> bool: return self.add_content_length if cmd.add_content_length is None else cmd.add_content_length + def _resolve_receipt_timeout(self, cmd: StompPublishCommand) -> float | None: + return self.receipt_timeout if cmd.receipt_timeout is None else cmd.receipt_timeout + def _make_headers_for_publish(cmd: StompPublishCommand) -> dict[str, str]: all_headers = cmd.headers.copy() @@ -102,7 +108,9 @@ def __init__(self, config: StompPublisherUsecaseConfig, specification: StompPubl async def _publish( self, cmd: PublishCommand, *, _extra_middlewares: typing.Iterable[PublisherMiddleware[PublishCommand]] ) -> None: - publish_command = StompPublishCommand.from_cmd(cmd, add_content_length=self.config.add_content_length) + publish_command = StompPublishCommand.from_cmd( + cmd, add_content_length=self.config.add_content_length, receipt_timeout=self.config.receipt_timeout + ) publish_command.destination = self.config.full_destination return typing.cast( "None", @@ -118,6 +126,7 @@ async def publish( correlation_id: str | None = None, headers: dict[str, str] | None = None, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> None: publish_command = StompPublishCommand( message, @@ -126,6 +135,7 @@ async def publish( correlation_id=correlation_id, headers=headers, add_content_length=self.config.add_content_length if add_content_length is None else add_content_length, + receipt_timeout=self.config.receipt_timeout if receipt_timeout is None else receipt_timeout, ) return typing.cast( "None", @@ -141,6 +151,7 @@ async def request( correlation_id: str | None = None, headers: dict[str, str] | None = None, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> Any: # noqa: ANN401 publish_command = StompPublishCommand( message, @@ -149,6 +160,7 @@ async def request( correlation_id=correlation_id, headers=headers, add_content_length=self.config.add_content_length if add_content_length is None else add_content_length, + receipt_timeout=self.config.receipt_timeout if receipt_timeout is None else receipt_timeout, ) return await self._basic_request(publish_command, producer=self.config._outer_config.producer) diff --git a/packages/faststream-stomp/faststream_stomp/registrator.py b/packages/faststream-stomp/faststream_stomp/registrator.py index b9d3759..a2d3777 100644 --- a/packages/faststream-stomp/faststream_stomp/registrator.py +++ b/packages/faststream-stomp/faststream_stomp/registrator.py @@ -74,11 +74,13 @@ def publisher( # type: ignore[override] description_: str | None = None, include_in_schema: bool = True, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> StompPublisher: usecase_config = StompPublisherUsecaseConfig( _outer_config=self.config, # type: ignore[arg-type] destination_without_prefix=destination, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ) specification = StompPublisherSpecification( _outer_config=self.config, # type: ignore[arg-type] diff --git a/packages/faststream-stomp/faststream_stomp/router.py b/packages/faststream-stomp/faststream_stomp/router.py index 7ca6984..bfe30df 100644 --- a/packages/faststream-stomp/faststream_stomp/router.py +++ b/packages/faststream-stomp/faststream_stomp/router.py @@ -26,6 +26,7 @@ def __init__( description_: str | None = None, include_in_schema: bool = True, add_content_length: bool | None = None, + receipt_timeout: float | None = None, ) -> None: super().__init__( destination=destination, @@ -34,6 +35,7 @@ def __init__( description_=description_, include_in_schema=include_in_schema, add_content_length=add_content_length, + receipt_timeout=receipt_timeout, ) diff --git a/packages/faststream-stomp/faststream_stomp/testing.py b/packages/faststream-stomp/faststream_stomp/testing.py index 2e34b6d..c70b5d2 100644 --- a/packages/faststream-stomp/faststream_stomp/testing.py +++ b/packages/faststream-stomp/faststream_stomp/testing.py @@ -66,7 +66,9 @@ async def nack(self) -> None: ... class FakeStompProducer(StompProducer): def __init__(self, broker: StompBroker) -> None: self.broker = broker - self.add_content_length = typing.cast("StompProducer", broker.config.producer).add_content_length + real_producer = typing.cast("StompProducer", broker.config.producer) + self.add_content_length = real_producer.add_content_length + self.receipt_timeout = real_producer.receipt_timeout self.codec = DefaultCodec() async def publish(self, cmd: StompPublishCommand) -> None: diff --git a/packages/faststream-stomp/pyproject.toml b/packages/faststream-stomp/pyproject.toml index 12f9409..992c739 100644 --- a/packages/faststream-stomp/pyproject.toml +++ b/packages/faststream-stomp/pyproject.toml @@ -2,7 +2,7 @@ name = "faststream-stomp" description = "FastStream STOMP broker" authors = [{ name = "Lev Vereshchagin", email = "mail@vrslev.com" }] -dependencies = ["faststream~=0.7", "stompman>=3.11.1"] +dependencies = ["faststream~=0.7", "stompman>=3.16.0"] requires-python = ">=3.11" readme = "README.md" license = { text = "MIT" } diff --git a/packages/faststream-stomp/test_faststream_stomp/test_main.py b/packages/faststream-stomp/test_faststream_stomp/test_main.py index 910bec1..82f97be 100644 --- a/packages/faststream-stomp/test_faststream_stomp/test_main.py +++ b/packages/faststream-stomp/test_faststream_stomp/test_main.py @@ -242,6 +242,116 @@ def test_subscriber_reply_default(self, broker: faststream_stomp.StompBroker, fa assert command.add_content_length is False +class TestReceiptTimeout: + @pytest.mark.parametrize( + ("broker_default", "call_override", "expected"), + [(None, None, None), (3.0, None, 3.0), (None, 1.0, 1.0), (3.0, 1.0, 1.0)], + ) + async def test_broker_publish( + self, + fake_connection_params: stompman.ConnectionParameters, + faker: faker.Faker, + *, + broker_default: float | None, + call_override: float | None, + expected: float | None, + ) -> None: + client, client_mock = make_mock_client(fake_connection_params) + broker = faststream_stomp.StompBroker(client, receipt_timeout=broker_default) + + await broker.publish(faker.pystr(), faker.pystr(), receipt_timeout=call_override) + + assert client_mock.send.await_args.kwargs["receipt_timeout"] == expected + + @pytest.mark.parametrize( + ("broker_default", "publisher_default", "call_override", "expected"), + [ + (None, None, None, None), + (3.0, None, None, 3.0), + (3.0, 5.0, None, 5.0), + (None, 5.0, None, 5.0), + (3.0, 5.0, 1.0, 1.0), + ], + ) + async def test_publisher_publish( + self, + fake_connection_params: stompman.ConnectionParameters, + faker: faker.Faker, + *, + broker_default: float | None, + publisher_default: float | None, + call_override: float | None, + expected: float | None, + ) -> None: + client, client_mock = make_mock_client(fake_connection_params) + broker = faststream_stomp.StompBroker(client, receipt_timeout=broker_default) + publisher = broker.publisher(faker.pystr(), receipt_timeout=publisher_default) + + await publisher.publish(faker.pystr(), receipt_timeout=call_override) + + assert client_mock.send.await_args.kwargs["receipt_timeout"] == expected + + async def test_publisher_decorator_default(self, broker: faststream_stomp.StompBroker, faker: faker.Faker) -> None: + receipt_timeout = faker.pyfloat(min_value=1, max_value=10) + publisher = broker.publisher(faker.pystr(), receipt_timeout=receipt_timeout) + publish_mock = mock.AsyncMock() + command = PublishCommand(faker.pystr(), _publish_type=PublishType.PUBLISH) + + with mock.patch.object(broker.config.producer, "publish", publish_mock): + await publisher._publish(command, _extra_middlewares=()) + + assert publish_mock.await_args is not None + assert publish_mock.await_args.args[0].receipt_timeout == receipt_timeout + + async def test_publish_batch_does_not_claim_confirmation( + self, fake_connection_params: stompman.ConnectionParameters, faker: faker.Faker + ) -> None: + client, client_mock = make_mock_client(fake_connection_params) + broker = faststream_stomp.StompBroker(client, receipt_timeout=3.0) + transaction = mock.AsyncMock() + messages = (faker.pystr(), faker.pystr()) + client_mock.begin.return_value.__aenter__.return_value = transaction + + await broker.publish_batch(*messages, destination=faker.pystr()) + + assert transaction.send.await_count == len(messages) + assert all("receipt_timeout" not in call.kwargs for call in transaction.send.await_args_list) + + async def test_error_propagates_from_broker_publish( + self, fake_connection_params: stompman.ConnectionParameters, faker: faker.Faker + ) -> None: + client, client_mock = make_mock_client(fake_connection_params) + client_mock.send.side_effect = stompman.SendReceiptError(receipt_id=faker.pystr(), reason="timeout") + broker = faststream_stomp.StompBroker(client) + + with pytest.raises(stompman.SendReceiptError, match="timeout"): + await broker.publish(faker.pystr(), faker.pystr(), receipt_timeout=1.0) + + async def test_error_propagates_from_publisher( + self, fake_connection_params: stompman.ConnectionParameters, faker: faker.Faker + ) -> None: + client, client_mock = make_mock_client(fake_connection_params) + client_mock.send.side_effect = stompman.SendReceiptError(receipt_id=faker.pystr(), reason="rejected") + broker = faststream_stomp.StompBroker(client) + publisher = broker.publisher(faker.pystr(), receipt_timeout=1.0) + + with pytest.raises(stompman.SendReceiptError, match="rejected"): + await publisher.publish(faker.pystr()) + + async def test_testing_broker_does_not_wait_for_receipts( + self, broker: faststream_stomp.StompBroker, faker: faker.Faker + ) -> None: + @broker.subscriber(destination := faker.pystr()) + def handle(body: str) -> None: ... + + async with faststream_stomp.TestStompBroker(broker) as br: + async with asyncio.timeout(1): + await br.publish(faker.pystr(), destination, receipt_timeout=1.0) + + assert handle.mock + handle.mock.assert_called_once() + + def test_asyncapi_schema(faker: faker.Faker, broker: faststream_stomp.StompBroker) -> None: @broker.publisher(faker.pystr()) def _publisher() -> None: ... diff --git a/packages/faststream-stomp/test_faststream_stomp/test_typing.py b/packages/faststream-stomp/test_faststream_stomp/test_typing.py index 98d6d7e..49803fd 100644 --- a/packages/faststream-stomp/test_faststream_stomp/test_typing.py +++ b/packages/faststream-stomp/test_faststream_stomp/test_typing.py @@ -27,3 +27,9 @@ async def check_add_content_length_typing() -> None: await publisher.publish_batch("first", "second", add_content_length=True) broker.subscriber("destination", reply_add_content_length=False) + + async def check_receipt_timeout_typing() -> None: + await broker.publish("message", "destination", receipt_timeout=1.0) + + publisher = broker.publisher("destination", receipt_timeout=5.0) + await publisher.publish("message", receipt_timeout=1.0) diff --git a/packages/stompman/stompman/__init__.py b/packages/stompman/stompman/__init__.py index 481af2f..0266a62 100644 --- a/packages/stompman/stompman/__init__.py +++ b/packages/stompman/stompman/__init__.py @@ -6,6 +6,7 @@ Error, FailedAllConnectAttemptsError, FailedAllWriteAttemptsError, + SendReceiptError, StompProtocolConnectionIssue, SubscriptionError, UnsupportedProtocolVersion, @@ -66,6 +67,7 @@ "NackFrame", "ReceiptFrame", "SendFrame", + "SendReceiptError", "StompProtocolConnectionIssue", "SubscribeFrame", "SubscriptionError", diff --git a/packages/stompman/stompman/client.py b/packages/stompman/stompman/client.py index db6b409..6885176 100644 --- a/packages/stompman/stompman/client.py +++ b/packages/stompman/stompman/client.py @@ -24,6 +24,8 @@ SendFrame, ) from stompman.logger import LOGGER +from stompman.receipts import PendingReceipts +from stompman.send import send_with_receipt from stompman.subscription import ( AckableMessageFrame, ActiveSubscriptions, @@ -71,7 +73,8 @@ class Client: connection_class: type[AbstractConnection] = Connection _connection_manager: ConnectionManager = field(init=False) - _active_subscriptions: ActiveSubscriptions = field(default_factory=ActiveSubscriptions, init=False) + _receipts: PendingReceipts = field(default_factory=PendingReceipts, init=False) + _active_subscriptions: ActiveSubscriptions = field(init=False) _active_transactions: set[Transaction] = field(default_factory=set, init=False) _exit_stack: AsyncExitStack = field(default_factory=AsyncExitStack, init=False) _listen_task: asyncio.Task[None] = field(init=False, repr=False) @@ -79,6 +82,7 @@ class Client: _handler_semaphore: asyncio.Semaphore | None = field(init=False, default=None, repr=False) def __post_init__(self) -> None: + self._active_subscriptions = ActiveSubscriptions(receipts=self._receipts) self._connection_manager = ConnectionManager( servers=self.servers, lifespan_factory=partial( @@ -98,7 +102,7 @@ def __post_init__(self) -> None: check_server_alive_interval_factor=self.check_server_alive_interval_factor, no_message_restart_interval=self.no_message_restart_interval, keep_alive_on_connection_failure=self.keep_alive_on_connection_failure, - on_connection_lost=self._active_subscriptions.connection_lost, + on_connection_lost=self._receipts.connection_lost, restore_connection=self._restore_connection, ssl=self.ssl, ) @@ -153,7 +157,7 @@ def _release(_t: asyncio.Task[None], s: asyncio.Semaphore = semaphore) -> None: task.add_done_callback(_release) case ErrorFrame() | ReceiptFrame(): - self._handle_subscription_frame(frame, epoch=epoch) + self._handle_receipt_frame(frame, epoch=epoch) case HeartbeatFrame() | ConnectedFrame(): pass @@ -164,10 +168,10 @@ async def _reserve_handler_slot(self) -> bool: if not semaphore.locked(): await semaphore.acquire() return True - if self._active_subscriptions.pending_receipts: + if self._receipts.pending: return False capacity = asyncio.create_task(semaphore.acquire()) - confirmation = asyncio.create_task(self._active_subscriptions.confirmation_started.wait()) + confirmation = asyncio.create_task(self._receipts.started.wait()) reserved = False try: await asyncio.wait((capacity, confirmation), return_when=asyncio.FIRST_COMPLETED) @@ -214,11 +218,11 @@ async def _invoke_message_handler( ) await _run_handler_with_safety_net(handler) - def _handle_subscription_frame(self, frame: ErrorFrame | ReceiptFrame, *, epoch: int) -> None: + def _handle_receipt_frame(self, frame: ErrorFrame | ReceiptFrame, *, epoch: int) -> None: if isinstance(frame, ReceiptFrame): - self._active_subscriptions.handle_receipt(frame, epoch=epoch) + self._receipts.handle_receipt(frame, epoch=epoch) else: - self._active_subscriptions.handle_error(frame, epoch=epoch) + self._receipts.handle_error(frame, epoch=epoch) if self.on_error_frame: self.on_error_frame(frame) @@ -230,7 +234,20 @@ async def send( content_type: str | None = None, add_content_length: bool = True, headers: dict[str, str] | None = None, + receipt_timeout: float | None = None, ) -> None: + if receipt_timeout is not None: + await send_with_receipt( + connection_manager=self._connection_manager, + receipts=self._receipts, + receipt_timeout=receipt_timeout, + body=body, + destination=destination, + content_type=content_type, + add_content_length=add_content_length, + headers=headers, + ) + return await self._connection_manager.write_frame_reconnecting( SendFrame.build( body=body, diff --git a/packages/stompman/stompman/errors.py b/packages/stompman/stompman/errors.py index 3c1a961..56c3ddc 100644 --- a/packages/stompman/stompman/errors.py +++ b/packages/stompman/stompman/errors.py @@ -62,3 +62,12 @@ class SubscriptionError(Error): subscription_id: str reason: Literal["rejected", "timeout", "connection_lost", "unsubscribed"] frame: ErrorFrame | None = field(default=None, repr=False) + + +@dataclass(kw_only=True) +class SendReceiptError(Error): + """A receipt-confirmed send was not confirmed by the broker.""" + + receipt_id: str + reason: Literal["rejected", "timeout", "connection_lost"] + frame: ErrorFrame | None = field(default=None, repr=False) diff --git a/packages/stompman/stompman/frames.py b/packages/stompman/stompman/frames.py index ab1e3da..6022d9b 100644 --- a/packages/stompman/stompman/frames.py +++ b/packages/stompman/stompman/frames.py @@ -27,6 +27,7 @@ "content-length": NotRequired[str], "content-type": NotRequired[str], "destination": str, + "receipt": NotRequired[str], "transaction": NotRequired[str], }, ) diff --git a/packages/stompman/stompman/receipts.py b/packages/stompman/stompman/receipts.py new file mode 100644 index 0000000..88d3ebc --- /dev/null +++ b/packages/stompman/stompman/receipts.py @@ -0,0 +1,84 @@ +import asyncio +from dataclasses import dataclass, field +from typing import Literal, Protocol +from uuid import uuid4 + +from stompman.connection import AbstractConnection +from stompman.errors import Error +from stompman.frames import ErrorFrame, ReceiptFrame + +ReceiptFailureReason = Literal["rejected", "timeout", "connection_lost"] + + +class ReceiptWaiter(Protocol): + def confirm(self, pending: "PendingReceipt") -> None: ... + def fail_on_error_frame(self, pending: "PendingReceipt", frame: ErrorFrame) -> None: ... + def fail_on_connection_loss(self, pending: "PendingReceipt") -> None: ... + + +@dataclass(kw_only=True, slots=True, frozen=True) +class PendingReceipt: + waiter: ReceiptWaiter + connection: AbstractConnection + receipt_id: str + epoch: int + deadline: float + result: asyncio.Future[Error | None] + + +@dataclass(kw_only=True, slots=True, frozen=True) +class PendingReceipts: + """Receipts awaited on the current connection, keyed by receipt ID.""" + + pending: dict[str, PendingReceipt] = field(default_factory=dict, init=False) + started: asyncio.Event = field(default_factory=asyncio.Event, init=False) + + def add(self, pending: PendingReceipt) -> None: + self.pending[pending.receipt_id] = pending + self.started.set() + + def discard(self, receipt_id: str) -> bool: + if self.pending.pop(receipt_id, None) is None: + return False + if not self.pending: + self.started.clear() + return True + + def handle_receipt(self, frame: ReceiptFrame, *, epoch: int) -> None: + if (pending := self.pending.get(frame.headers["receipt-id"])) and pending.epoch == epoch: + pending.waiter.confirm(pending) + + def handle_error(self, frame: ErrorFrame, *, epoch: int) -> None: + receipt_id = frame.headers.get("receipt-id") + if receipt_id is not None: + pending = self.pending.get(receipt_id) + affected = [pending] if pending is not None else [] + else: + # An uncorrelated ERROR cannot safely confirm anything in flight. + affected = list(self.pending.values()) + for pending in affected: + if pending.epoch == epoch: + pending.waiter.fail_on_error_frame(pending, frame) + + def connection_lost(self, connection: AbstractConnection) -> None: + for pending in list(self.pending.values()): + if pending.connection is connection: + pending.waiter.fail_on_connection_loss(pending) + + +async def wait_for_result(pending: PendingReceipt) -> Error | None: + """Wait for the receipt, raising TimeoutError only if the deadline passed with nothing delivered.""" + try: + if pending.result.done(): + return pending.result.result() + async with asyncio.timeout_at(pending.deadline): + return await asyncio.shield(pending.result) + except TimeoutError: + # A receipt already processed by the reader beats a deadline cancellation delivered before we were rescheduled. + if pending.result.done(): + return pending.result.result() + raise + + +def make_receipt_id() -> str: + return str(uuid4()) diff --git a/packages/stompman/stompman/send.py b/packages/stompman/stompman/send.py new file mode 100644 index 0000000..7a2ebc4 --- /dev/null +++ b/packages/stompman/stompman/send.py @@ -0,0 +1,111 @@ +import asyncio +import math +from dataclasses import dataclass + +from stompman.connection_manager import ConnectionManager +from stompman.errors import ConnectionLostError, SendReceiptError +from stompman.frames import ErrorFrame, SendFrame +from stompman.receipts import PendingReceipt, PendingReceipts, ReceiptFailureReason, make_receipt_id, wait_for_result + + +@dataclass(kw_only=True, slots=True, frozen=True) +class _SendConfirmation: + receipts: PendingReceipts + + def confirm(self, pending: PendingReceipt) -> None: + self.receipts.discard(pending.receipt_id) + if not pending.result.done(): + pending.result.set_result(None) + + def fail_on_error_frame(self, pending: PendingReceipt, frame: ErrorFrame) -> None: + self.fail(pending, "rejected", frame=frame) + + def fail_on_connection_loss(self, pending: PendingReceipt) -> None: + self.fail(pending, "connection_lost") + + def fail( + self, pending: PendingReceipt, reason: ReceiptFailureReason, *, frame: ErrorFrame | None = None + ) -> SendReceiptError: + self.receipts.discard(pending.receipt_id) + if pending.result.done() and isinstance(existing := pending.result.result(), SendReceiptError): + return existing + error = SendReceiptError(receipt_id=pending.receipt_id, reason=reason, frame=frame) + if not pending.result.done(): + pending.result.set_result(error) + return error + + +def _is_confirmed(pending: PendingReceipt) -> bool: + return pending.result.done() and pending.result.result() is None + + +async def send_with_receipt( + *, + connection_manager: ConnectionManager, + receipts: PendingReceipts, + receipt_timeout: float, + body: bytes, + destination: str, + content_type: str | None, + add_content_length: bool, + headers: dict[str, str] | None, +) -> None: + if not math.isfinite(receipt_timeout) or receipt_timeout <= 0: + msg = "receipt_timeout must be a finite positive number" + raise ValueError(msg) + if headers and "receipt" in headers: + msg = "receipt header is generated by the client when receipt_timeout is set" + raise ValueError(msg) + + connection_state = await connection_manager._get_restored_connection_state() + loop = asyncio.get_running_loop() + confirmation = _SendConfirmation(receipts=receipts) + pending = PendingReceipt( + waiter=confirmation, + connection=connection_state.connection, + receipt_id=make_receipt_id(), + epoch=connection_manager._reconnection_count, + deadline=loop.time() + receipt_timeout, + result=loop.create_future(), + ) + receipts.add(pending) + + try: + async with asyncio.timeout_at(pending.deadline): + await connection_state.connection.write_frame( + SendFrame.build( + body=body, + destination=destination, + transaction=None, + content_type=content_type, + add_content_length=add_content_length, + headers={**(headers or {}), "receipt": pending.receipt_id}, + ) + ) + except TimeoutError as error: + # The frame is fully buffered before draining, so a receipt can win this race. + if _is_confirmed(pending): + return + raise confirmation.fail(pending, "timeout") from error + except ConnectionLostError as error: + await connection_manager._discard_failed_connection_state(connection_state, error) + if _is_confirmed(pending): + return + raise confirmation.fail(pending, "connection_lost") from error + except BaseException: + receipts.discard(pending.receipt_id) + raise + + await _await_receipt(pending, confirmation) + + +async def _await_receipt(pending: PendingReceipt, confirmation: _SendConfirmation) -> None: + try: + error = await wait_for_result(pending) + except TimeoutError as timeout_error: + raise confirmation.fail(pending, "timeout") from timeout_error + except asyncio.CancelledError: + confirmation.receipts.discard(pending.receipt_id) + raise + if error is not None: + raise error diff --git a/packages/stompman/stompman/subscription.py b/packages/stompman/stompman/subscription.py index fa2a556..8488b75 100644 --- a/packages/stompman/stompman/subscription.py +++ b/packages/stompman/stompman/subscription.py @@ -15,30 +15,19 @@ ErrorFrame, MessageFrame, NackFrame, - ReceiptFrame, SubscribeFrame, UnsubscribeFrame, ) from stompman.logger import LOGGER - - -@dataclass(kw_only=True, slots=True, frozen=True) -class _PendingConfirmation: - subscription: "BaseSubscription" - connection: AbstractConnection - receipt_id: str - epoch: int - deadline: float - result: asyncio.Future[SubscriptionError | None] +from stompman.receipts import PendingReceipt, PendingReceipts, make_receipt_id, wait_for_result @dataclass(kw_only=True, slots=True, frozen=True) class ActiveSubscriptions: + receipts: PendingReceipts subscriptions: dict[str, "AutoAckSubscription | ManualAckSubscription"] = field(default_factory=dict, init=False) event: asyncio.Event = field(default_factory=asyncio.Event, init=False) - pending_receipts: dict[str, _PendingConfirmation] = field(default_factory=dict, init=False) confirmation_tasks: set[asyncio.Task[None]] = field(default_factory=set, init=False) - confirmation_started: asyncio.Event = field(default_factory=asyncio.Event, init=False) def __post_init__(self) -> None: self.event.set() @@ -68,40 +57,8 @@ def contains_by_id(self, subscription_id: str) -> bool: async def wait_until_empty(self) -> bool: return await self.event.wait() - def handle_receipt(self, frame: ReceiptFrame, *, epoch: int) -> None: - if (pending := self.pending_receipts.get(frame.headers["receipt-id"])) and pending.epoch == epoch: - del self.pending_receipts[pending.receipt_id] - if not self.pending_receipts: - self.confirmation_started.clear() - pending.subscription._pending_confirmation = None - if not pending.result.done(): - pending.result.set_result(None) - - def handle_error(self, frame: ErrorFrame, *, epoch: int) -> None: - receipt_id = frame.headers.get("receipt-id") - if receipt_id is not None: - pending = self.pending_receipts.get(receipt_id) - affected = [pending] if pending is not None else [] - else: - # An uncorrelated ERROR cannot safely confirm any in-flight - # subscription. Confirmed subscriptions remain eligible for recovery. - affected = list(self.pending_receipts.values()) - for pending in affected: - if pending.epoch != epoch: - continue - pending.subscription._fail_confirmation( - pending, SubscriptionError(subscription_id=pending.subscription.id, reason="rejected", frame=frame) - ) - - def connection_lost(self, connection: AbstractConnection) -> None: - for pending in list(self.pending_receipts.values()): - if pending.connection is connection: - pending.subscription._fail_confirmation( - pending, SubscriptionError(subscription_id=pending.subscription.id, reason="connection_lost") - ) - - def watch_confirmation(self, pending: _PendingConfirmation) -> None: - task = asyncio.create_task(pending.subscription._watch_confirmation(pending)) + def watch_confirmation(self, subscription: "BaseSubscription", pending: PendingReceipt) -> None: + task = asyncio.create_task(subscription._watch_confirmation(pending)) self.confirmation_tasks.add(task) task.add_done_callback(self.confirmation_tasks.discard) @@ -124,7 +81,7 @@ class BaseSubscription: _active_subscriptions: ActiveSubscriptions receipt_timeout: float | None = None on_subscription_error: Callable[[SubscriptionError], Any] | None = None - _pending_confirmation: _PendingConfirmation | None = field(default=None, init=False, repr=False) + _pending_confirmation: PendingReceipt | None = field(default=None, init=False, repr=False) async def _subscribe(self) -> None: if self.receipt_timeout is not None: @@ -157,19 +114,18 @@ async def unsubscribe(self) -> None: async def _send_confirmed_subscription( self, connection: AbstractConnection, *, restoring: bool = False - ) -> _PendingConfirmation: + ) -> PendingReceipt: assert self.receipt_timeout is not None # ruff: ignore[assert] - internal invariant - pending = _PendingConfirmation( - subscription=self, + pending = PendingReceipt( + waiter=self, connection=connection, - receipt_id=_make_confirmation_id(), + receipt_id=make_receipt_id(), epoch=self._connection_manager._reconnection_count, deadline=asyncio.get_running_loop().time() + self.receipt_timeout, result=asyncio.get_running_loop().create_future(), ) self._pending_confirmation = pending - self._active_subscriptions.pending_receipts[pending.receipt_id] = pending - self._active_subscriptions.confirmation_started.set() + self._active_subscriptions.receipts.add(pending) try: async with asyncio.timeout_at(pending.deadline): await connection.write_frame( @@ -206,25 +162,14 @@ async def _send_confirmed_subscription( raise return pending - async def _await_confirmation(self, pending: _PendingConfirmation) -> None: + async def _await_confirmation(self, pending: PendingReceipt) -> None: try: - if pending.result.done(): - error = pending.result.result() - else: - async with asyncio.timeout_at(pending.deadline): - error = await asyncio.shield(pending.result) + error = await wait_for_result(pending) except TimeoutError as timeout_error: - # Prefer a receipt already processed by the reader over a deadline - # cancellation delivered before this waiter was rescheduled. - if pending.result.done(): - error = pending.result.result() - if error is not None: - raise error from timeout_error - return - error = SubscriptionError(subscription_id=self.id, reason="timeout") - self._fail_confirmation(pending, error) + failure = SubscriptionError(subscription_id=self.id, reason="timeout") + self._fail_confirmation(pending, failure) await self._cleanup_confirmation(pending) - raise error from timeout_error + raise failure from timeout_error except asyncio.CancelledError: was_active = self._active_subscriptions.contains_by_id(self.id) pending = self._pending_confirmation or pending @@ -239,15 +184,27 @@ async def _await_confirmation(self, pending: _PendingConfirmation) -> None: if error is not None: raise error + def confirm(self, pending: PendingReceipt) -> None: + self._active_subscriptions.receipts.discard(pending.receipt_id) + self._pending_confirmation = None + if not pending.result.done(): + pending.result.set_result(None) + + def fail_on_error_frame(self, pending: PendingReceipt, frame: ErrorFrame) -> None: + self._fail_confirmation(pending, SubscriptionError(subscription_id=self.id, reason="rejected", frame=frame)) + + def fail_on_connection_loss(self, pending: PendingReceipt) -> None: + self._fail_confirmation(pending, SubscriptionError(subscription_id=self.id, reason="connection_lost")) + def _fail_confirmation( self, - pending: _PendingConfirmation, + pending: PendingReceipt, error: SubscriptionError, *, notify: bool = True, include_confirmed: bool = False, ) -> None: - if self._active_subscriptions.pending_receipts.pop(pending.receipt_id, None) is None: + if not self._active_subscriptions.receipts.discard(pending.receipt_id): if not include_confirmed: return # A write may fail after its receipt. Notify once, without removing @@ -259,8 +216,6 @@ def _fail_confirmation( or not self._active_subscriptions.contains_by_id(self.id) ): return - if not self._active_subscriptions.pending_receipts: - self._active_subscriptions.confirmation_started.clear() self._pending_confirmation = None self._active_subscriptions.delete_by_id(self.id) if not pending.result.done(): @@ -274,7 +229,7 @@ def _fail_confirmation( except Exception: # ruff: ignore[blind-except] - user callbacks must not terminate the frame reader LOGGER.exception("unhandled exception in subscription error callback") - async def _cleanup_confirmation(self, pending: _PendingConfirmation) -> None: + async def _cleanup_confirmation(self, pending: PendingReceipt) -> None: try: async with asyncio.timeout(self.receipt_timeout): await pending.connection.write_frame(UnsubscribeFrame(headers={"id": self.id})) @@ -285,7 +240,7 @@ async def _cleanup_confirmation(self, pending: _PendingConfirmation) -> None: state, ConnectionLostError(reason=error) ) - async def _watch_confirmation(self, pending: _PendingConfirmation) -> None: + async def _watch_confirmation(self, pending: PendingReceipt) -> None: # Failure notification and registry cleanup happen before waking us. with suppress(SubscriptionError): await self._await_confirmation(pending) @@ -306,7 +261,7 @@ async def _resubscribe(self, connection: AbstractConnection) -> None: pending = await self._send_confirmed_subscription(connection, restoring=True) except SubscriptionError: return - self._active_subscriptions.watch_confirmation(pending) + self._active_subscriptions.watch_confirmation(self, pending) async def _nack(self, frame: MessageFrame, *, received_at_reconnection_count: int) -> None: if not self._active_subscriptions.contains_by_id(self.id): @@ -416,10 +371,6 @@ def _make_subscription_id() -> str: return str(uuid4()) -def _make_confirmation_id() -> str: - return str(uuid4()) - - async def resubscribe_to_active_subscriptions( *, connection: AbstractConnection, active_subscriptions: ActiveSubscriptions ) -> None: diff --git a/packages/stompman/test_stompman/conftest.py b/packages/stompman/test_stompman/conftest.py index d1da9c7..e66aec1 100644 --- a/packages/stompman/test_stompman/conftest.py +++ b/packages/stompman/test_stompman/conftest.py @@ -4,6 +4,7 @@ from datetime import timedelta from ssl import SSLContext from typing import Any, Literal, Self, TypeVar +from unittest import mock import pytest import stompman @@ -167,3 +168,61 @@ def enrich_expected_frames( stompman.DisconnectFrame(headers={"receipt": "receipt-id-1"}), stompman.ReceiptFrame(headers={"receipt-id": "receipt-id-1"}), ] + + +Incoming = dict[int, asyncio.Queue[stompman.AnyServerFrame | stompman.ConnectionLostError]] +Outgoing = asyncio.Queue[tuple[AbstractConnection, stompman.AnyClientFrame]] + + +@pytest.fixture +def incoming() -> Incoming: + return {} + + +@pytest.fixture +def outgoing() -> Outgoing: + return asyncio.Queue() + + +@pytest.fixture +async def client( + monkeypatch: pytest.MonkeyPatch, incoming: Incoming, outgoing: Outgoing +) -> AsyncGenerator[stompman.Client, None]: + async def write_frame( # ruff: ignore[unused-async] + connection: AbstractConnection, frame: stompman.AnyClientFrame + ) -> None: + queue = incoming.setdefault(id(connection), asyncio.Queue()) + outgoing.put_nowait((connection, frame)) + if isinstance(frame, stompman.ConnectFrame): + queue.put_nowait(stompman.ConnectedFrame(headers={"version": "1.2", "heart-beat": "1000,1000"})) + elif isinstance(frame, stompman.DisconnectFrame): + queue.put_nowait(stompman.ReceiptFrame(headers={"receipt-id": frame.headers["receipt"]})) + + async def read_frames(connection: AbstractConnection) -> AsyncGenerator[stompman.AnyServerFrame, None]: + queue = incoming.setdefault(id(connection), asyncio.Queue()) + while True: + frame = await queue.get() + if isinstance(frame, stompman.ConnectionLostError): + raise frame + yield frame + + monkeypatch.setattr(BaseMockConnection, "write_frame", write_frame) + monkeypatch.setattr(BaseMockConnection, "read_frames", read_frames) + async with EnrichedClient( + connection_class=BaseMockConnection, + connect_retry_interval=0, + max_concurrent_handlers=1, + on_error_frame=mock.Mock(), + ) as instance: + try: + yield instance + finally: + for subscription in instance._active_subscriptions.get_all(): + await subscription.unsubscribe() + + +def remaining_frames(outgoing: Outgoing) -> list[stompman.AnyClientFrame]: + frames: list[stompman.AnyClientFrame] = [] + while not outgoing.empty(): + frames.append(outgoing.get_nowait()[1]) + return frames diff --git a/packages/stompman/test_stompman/test_integration.py b/packages/stompman/test_stompman/test_integration.py index 54720f2..0c1a487 100644 --- a/packages/stompman/test_stompman/test_integration.py +++ b/packages/stompman/test_stompman/test_integration.py @@ -115,7 +115,7 @@ async def handle_message(frame: stompman.AckableMessageFrame) -> None: # ruff: receipt_timeout=3, ) assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending await force_reconnect(client) healthy = await client.subscribe_with_manual_ack(destination, handle_message, ack="auto", receipt_timeout=3) try: @@ -127,6 +127,28 @@ async def handle_message(frame: stompman.AckableMessageFrame) -> None: # ruff: await healthy.unsubscribe() +@pytest.mark.anyio +@pytest.mark.parametrize("port", [9000, 9001], ids=["artemis", "classic"]) +async def test_confirmed_send_is_accepted_by_broker(port: int) -> None: + parameters = stompman.ConnectionParameters("127.0.0.1", port, "admin", ":=123") + destination = f"confirmed-send-{uuid4()}" + payload = str(uuid4()).encode() + received: asyncio.Future[bytes] = asyncio.get_running_loop().create_future() + + async def handle_message(frame: stompman.AckableMessageFrame) -> None: # ruff: ignore[unused-async] + if not received.done(): + received.set_result(frame.body) + + async with stompman.Client(servers=[parameters], connection_confirmation_timeout=10) as client: + subscription = await client.subscribe_with_manual_ack(destination, handle_message, ack="auto") + try: + await client.send(payload, destination, receipt_timeout=3) + assert not client._receipts.pending + assert await asyncio.wait_for(received, timeout=3) == payload + finally: + await subscription.unsubscribe() + + @pytest.mark.anyio async def test_ok(connection_parameters: stompman.ConnectionParameters) -> None: async def produce(destination: str) -> None: diff --git a/packages/stompman/test_stompman/test_send_confirmation.py b/packages/stompman/test_stompman/test_send_confirmation.py new file mode 100644 index 0000000..f48dfe8 --- /dev/null +++ b/packages/stompman/test_stompman/test_send_confirmation.py @@ -0,0 +1,364 @@ +import asyncio +import math +import typing +from unittest import mock + +import pytest +import stompman +from stompman.connection import AbstractConnection +from stompman.receipts import PendingReceipt, PendingReceipts, wait_for_result + +from test_stompman.conftest import BaseMockConnection, Incoming, Outgoing, remaining_frames + +pytestmark = [pytest.mark.anyio, pytest.mark.timeout(10)] + + +async def next_send(outgoing: Outgoing) -> tuple[AbstractConnection, stompman.SendFrame]: + async with asyncio.timeout(1): + while True: + connection, frame = await outgoing.get() + if isinstance(frame, stompman.SendFrame): + return connection, frame + + +def receipt(incoming: Incoming, connection: AbstractConnection, frame: stompman.SendFrame) -> None: + incoming[id(connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": frame.headers["receipt"]})) + + +async def test_send_without_receipt_timeout_does_not_wait_for_receipt( + client: stompman.Client, outgoing: Outgoing +) -> None: + await client.send(b"hi", "destination") + + _, frame = await next_send(outgoing) + assert "receipt" not in frame.headers + assert not client._receipts.pending + + +async def test_confirmed_send_waits_for_matching_receipt( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + headers = {"persistent": "true"} + task = asyncio.create_task(client.send(b"hi", "destination", headers=headers, receipt_timeout=1)) + + connection, frame = await next_send(outgoing) + assert frame.headers["receipt"] + assert headers == {"persistent": "true"} + incoming[id(connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": "unrelated"})) + incoming[id(connection)].put_nowait( + stompman.ErrorFrame(headers={"message": "unrelated", "receipt-id": "unrelated"}) + ) + await asyncio.sleep(0) + assert not task.done() + + receipt(incoming, connection, frame) + await task + assert not client._receipts.pending + + +async def test_concurrent_confirmed_sends_complete_independently( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + first = asyncio.create_task(client.send(b"first", "destination", receipt_timeout=1)) + connection, first_frame = await next_send(outgoing) + second = asyncio.create_task(client.send(b"second", "destination", receipt_timeout=1)) + _, second_frame = await next_send(outgoing) + assert first_frame.headers["receipt"] != second_frame.headers["receipt"] + + receipt(incoming, connection, second_frame) + await second + assert not first.done() + + receipt(incoming, connection, first_frame) + await first + assert not client._receipts.pending + + +async def test_correlated_error_rejects_send(client: stompman.Client, incoming: Incoming, outgoing: Outgoing) -> None: + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + connection, frame = await next_send(outgoing) + error_frame = stompman.ErrorFrame(headers={"message": "rejected", "receipt-id": frame.headers["receipt"]}) + incoming[id(connection)].put_nowait(error_frame) + + with pytest.raises(stompman.SendReceiptError) as exc_info: + await task + + assert exc_info.value.reason == "rejected" + assert exc_info.value.receipt_id == frame.headers["receipt"] + assert exc_info.value.frame == error_frame + assert "rejected" in repr(exc_info.value) + assert "message" not in repr(exc_info.value) + assert not client._receipts.pending + assert typing.cast("mock.Mock", client.on_error_frame).mock_calls == [mock.call(error_frame)] + + +async def test_uncorrelated_error_rejects_pending_send( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + connection, _ = await next_send(outgoing) + incoming[id(connection)].put_nowait(stompman.ErrorFrame(headers={"message": "broker is unhappy"})) + + with pytest.raises(stompman.SendReceiptError, match="rejected"): + await task + + assert not client._receipts.pending + + +async def test_timeout_raises_and_removes_pending_state(client: stompman.Client, outgoing: Outgoing) -> None: + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=0.05)) + await next_send(outgoing) + + with pytest.raises(stompman.SendReceiptError, match="timeout"): + await task + + assert not client._receipts.pending + + +async def test_connection_loss_raises_and_removes_pending_state( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + connection, _ = await next_send(outgoing) + incoming[id(connection)].put_nowait(stompman.ConnectionLostError(reason="connection lost before receipt")) + + with pytest.raises(stompman.SendReceiptError, match="connection_lost"): + await task + + assert not client._receipts.pending + + +async def test_cancellation_removes_pending_state(client: stompman.Client, outgoing: Outgoing) -> None: + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + await next_send(outgoing) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert not client._receipts.pending + + +@pytest.mark.parametrize(("receipt_timeout", "write_duration"), [(1, 0.05), (0.05, 0.15)]) +async def test_receipt_arriving_before_write_finishes_confirms_send( + client: stompman.Client, + incoming: Incoming, + monkeypatch: pytest.MonkeyPatch, + receipt_timeout: float, + write_duration: float, +) -> None: + write_frame = BaseMockConnection.write_frame + + async def write_then_receive_receipt(connection: BaseMockConnection, frame: stompman.AnyClientFrame) -> None: + await write_frame(connection, frame) + if isinstance(frame, stompman.SendFrame) and (receipt_id := frame.headers.get("receipt")): + incoming[id(connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": receipt_id})) + await asyncio.sleep(write_duration) + + monkeypatch.setattr(BaseMockConnection, "write_frame", write_then_receive_receipt) + + await client.send(b"hi", "destination", receipt_timeout=receipt_timeout) + + assert not client._receipts.pending + + +async def test_timeout_while_writing_raises_and_removes_pending_state( + client: stompman.Client, monkeypatch: pytest.MonkeyPatch +) -> None: + write_frame = BaseMockConnection.write_frame + + async def slow_write(connection: BaseMockConnection, frame: stompman.AnyClientFrame) -> None: + await write_frame(connection, frame) + if isinstance(frame, stompman.SendFrame): + await asyncio.sleep(0.2) + + monkeypatch.setattr(BaseMockConnection, "write_frame", slow_write) + + with pytest.raises(stompman.SendReceiptError, match="timeout"): + await client.send(b"hi", "destination", receipt_timeout=0.05) + + assert not client._receipts.pending + + +async def test_connection_loss_while_writing_raises_and_removes_pending_state( + client: stompman.Client, monkeypatch: pytest.MonkeyPatch +) -> None: + write_frame = BaseMockConnection.write_frame + + async def failing_write(connection: BaseMockConnection, frame: stompman.AnyClientFrame) -> None: + if isinstance(frame, stompman.SendFrame): + raise stompman.ConnectionLostError(reason="broken pipe") + await write_frame(connection, frame) + + monkeypatch.setattr(BaseMockConnection, "write_frame", failing_write) + + with pytest.raises(stompman.SendReceiptError, match="connection_lost"): + await client.send(b"hi", "destination", receipt_timeout=1) + + assert not client._receipts.pending + + +async def test_receipt_received_before_connection_loss_confirms_send( + client: stompman.Client, incoming: Incoming, monkeypatch: pytest.MonkeyPatch +) -> None: + write_frame = BaseMockConnection.write_frame + + async def receive_receipt_then_fail(connection: BaseMockConnection, frame: stompman.AnyClientFrame) -> None: + await write_frame(connection, frame) + if isinstance(frame, stompman.SendFrame) and (receipt_id := frame.headers.get("receipt")): + incoming[id(connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": receipt_id})) + await asyncio.sleep(0) + raise stompman.ConnectionLostError(reason="peer closed after receipt") + + monkeypatch.setattr(BaseMockConnection, "write_frame", receive_receipt_then_fail) + + await client.send(b"hi", "destination", receipt_timeout=1) + + assert not client._receipts.pending + + +async def test_cancellation_while_writing_removes_pending_state( + client: stompman.Client, outgoing: Outgoing, monkeypatch: pytest.MonkeyPatch +) -> None: + write_frame = BaseMockConnection.write_frame + blocked = asyncio.Event() + + async def blocked_write(connection: BaseMockConnection, frame: stompman.AnyClientFrame) -> None: + await write_frame(connection, frame) + if isinstance(frame, stompman.SendFrame): + await blocked.wait() + + monkeypatch.setattr(BaseMockConnection, "write_frame", blocked_write) + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + await next_send(outgoing) + task.cancel() + + with pytest.raises(asyncio.CancelledError): + await task + + assert not client._receipts.pending + + +async def test_receipt_delivered_at_the_deadline_confirms_send() -> None: + loop = asyncio.get_running_loop() + result: asyncio.Future[stompman.Error | None] = loop.create_future() + pending = PendingReceipt( + waiter=mock.Mock(), + connection=mock.Mock(), + receipt_id="receipt-id", + epoch=0, + deadline=loop.time() + 0.05, + result=result, + ) + loop.call_at(pending.deadline, result.set_result, None) + + assert await wait_for_result(pending) is None + + +async def test_receipt_from_previous_connection_does_not_confirm_new_send( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + lost = asyncio.create_task(client.send(b"first", "destination", receipt_timeout=1)) + old_connection, old_frame = await next_send(outgoing) + incoming[id(old_connection)].put_nowait(stompman.ConnectionLostError(reason="peer closed")) + with pytest.raises(stompman.SendReceiptError, match="connection_lost"): + await lost + + task = asyncio.create_task(client.send(b"second", "destination", receipt_timeout=1)) + new_connection, new_frame = await next_send(outgoing) + assert new_connection is not old_connection + incoming[id(new_connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": old_frame.headers["receipt"]})) + await asyncio.sleep(0) + assert not task.done() + + receipt(incoming, new_connection, new_frame) + await task + + +@pytest.mark.parametrize("error", [True, False]) +def test_pending_receipts_ignore_other_epoch(*, error: bool) -> None: + loop = asyncio.new_event_loop() + waiter = mock.Mock() + receipts = PendingReceipts() + try: + receipts.add( + PendingReceipt( + waiter=waiter, + connection=mock.Mock(), + receipt_id="receipt-id", + epoch=0, + deadline=loop.time(), + result=loop.create_future(), + ) + ) + + if error: + receipts.handle_error(stompman.ErrorFrame(headers={"message": "nope", "receipt-id": "receipt-id"}), epoch=1) + else: + receipts.handle_receipt(stompman.ReceiptFrame(headers={"receipt-id": "receipt-id"}), epoch=1) + finally: + loop.close() + + assert not waiter.mock_calls + assert receipts.pending + + +@pytest.mark.parametrize("receipt_timeout", [0, -1, math.nan, math.inf]) +async def test_invalid_receipt_timeout_fails_before_write( + client: stompman.Client, outgoing: Outgoing, receipt_timeout: float +) -> None: + with pytest.raises(ValueError, match="receipt_timeout"): + await client.send(b"hi", "destination", receipt_timeout=receipt_timeout) + + assert not any(isinstance(frame, stompman.SendFrame) for frame in remaining_frames(outgoing)) + assert not client._receipts.pending + + +async def test_caller_provided_receipt_header_is_rejected(client: stompman.Client, outgoing: Outgoing) -> None: + with pytest.raises(ValueError, match="receipt"): + await client.send(b"hi", "destination", headers={"receipt": "mine"}, receipt_timeout=1) + + assert not any(isinstance(frame, stompman.SendFrame) for frame in remaining_frames(outgoing)) + assert not client._receipts.pending + + +async def test_caller_provided_receipt_header_is_kept_without_confirmation( + client: stompman.Client, outgoing: Outgoing +) -> None: + await client.send(b"hi", "destination", headers={"receipt": "mine"}) + + _, frame = await next_send(outgoing) + assert frame.headers["receipt"] == "mine" + + +async def test_confirmed_send_completes_while_message_handlers_are_saturated( + client: stompman.Client, incoming: Incoming, outgoing: Outgoing +) -> None: + release_handler = asyncio.Event() + handling = asyncio.Event() + + async def handle_message(frame: stompman.AckableMessageFrame) -> None: + handling.set() + await release_handler.wait() + + subscription = await client.subscribe_with_manual_ack("destination", handle_message) + connection = next(iter(incoming)) + for _ in range(2): + incoming[connection].put_nowait( + stompman.MessageFrame( + headers={"destination": "destination", "message-id": "1", "subscription": subscription.id}, + body=b"hi", + ) + ) + await asyncio.wait_for(handling.wait(), timeout=1) + await asyncio.sleep(0) + + task = asyncio.create_task(client.send(b"hi", "destination", receipt_timeout=1)) + sending_connection, frame = await next_send(outgoing) + receipt(incoming, sending_connection, frame) + try: + await task + finally: + release_handler.set() + + assert not client._receipts.pending diff --git a/packages/stompman/test_stompman/test_subscription_confirmation.py b/packages/stompman/test_stompman/test_subscription_confirmation.py index 2c5e610..464b72f 100644 --- a/packages/stompman/test_stompman/test_subscription_confirmation.py +++ b/packages/stompman/test_stompman/test_subscription_confirmation.py @@ -1,64 +1,23 @@ import asyncio -from collections.abc import AsyncGenerator, Coroutine -from typing import Any +from typing import TYPE_CHECKING, Any from unittest import mock import pytest import stompman from stompman.connection import AbstractConnection -from test_stompman.conftest import BaseMockConnection, EnrichedClient, noop_message_handler +from test_stompman.conftest import ( + BaseMockConnection, + Incoming, + Outgoing, + noop_message_handler, + remaining_frames, +) -pytestmark = [pytest.mark.anyio, pytest.mark.timeout(10)] -Incoming = dict[int, asyncio.Queue[stompman.AnyServerFrame | stompman.ConnectionLostError]] -Outgoing = asyncio.Queue[tuple[AbstractConnection, stompman.AnyClientFrame]] - - -@pytest.fixture -def incoming() -> Incoming: - return {} - - -@pytest.fixture -def outgoing() -> Outgoing: - return asyncio.Queue() - - -@pytest.fixture -async def client( - monkeypatch: pytest.MonkeyPatch, incoming: Incoming, outgoing: Outgoing -) -> AsyncGenerator[stompman.Client, None]: - async def write_frame( # ruff: ignore[unused-async] - connection: AbstractConnection, frame: stompman.AnyClientFrame - ) -> None: - queue = incoming.setdefault(id(connection), asyncio.Queue()) - outgoing.put_nowait((connection, frame)) - if isinstance(frame, stompman.ConnectFrame): - queue.put_nowait(stompman.ConnectedFrame(headers={"version": "1.2", "heart-beat": "1000,1000"})) - elif isinstance(frame, stompman.DisconnectFrame): - queue.put_nowait(stompman.ReceiptFrame(headers={"receipt-id": frame.headers["receipt"]})) +if TYPE_CHECKING: + from collections.abc import Coroutine - async def read_frames(connection: AbstractConnection) -> AsyncGenerator[stompman.AnyServerFrame, None]: - queue = incoming.setdefault(id(connection), asyncio.Queue()) - while True: - frame = await queue.get() - if isinstance(frame, stompman.ConnectionLostError): - raise frame - yield frame - - monkeypatch.setattr(BaseMockConnection, "write_frame", write_frame) - monkeypatch.setattr(BaseMockConnection, "read_frames", read_frames) - async with EnrichedClient( - connection_class=BaseMockConnection, - connect_retry_interval=0, - max_concurrent_handlers=1, - on_error_frame=mock.Mock(), - ) as instance: - try: - yield instance - finally: - for subscription in instance._active_subscriptions.get_all(): - await subscription.unsubscribe() +pytestmark = [pytest.mark.anyio, pytest.mark.timeout(10)] async def next_subscribe(outgoing: Outgoing) -> tuple[AbstractConnection, stompman.SubscribeFrame]: @@ -73,13 +32,6 @@ def receipt(incoming: Incoming, connection: AbstractConnection, frame: stompman. incoming[id(connection)].put_nowait(stompman.ReceiptFrame(headers={"receipt-id": frame.headers["receipt"]})) -def remaining_frames(outgoing: Outgoing) -> list[stompman.AnyClientFrame]: - frames: list[stompman.AnyClientFrame] = [] - while not outgoing.empty(): - frames.append(outgoing.get_nowait()[1]) - return frames - - @pytest.mark.parametrize("manual_ack", [True, False]) async def test_subscribe_waits_for_matching_receipt( client: stompman.Client, incoming: Incoming, outgoing: Outgoing, *, manual_ack: bool @@ -103,7 +55,7 @@ async def test_subscribe_waits_for_matching_receipt( receipt(incoming, connection, frame) subscription = task.result() assert subscription.id == frame.headers["id"] - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending await subscription.unsubscribe() @@ -132,7 +84,7 @@ def on_failure(error: stompman.SubscriptionError) -> None: with pytest.raises(stompman.SubscriptionError, match="rejected"): await task assert len(failures) == 1 - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending incoming[id(connection)].put_nowait(stompman.ConnectionLostError(reason="peer closed after ERROR")) await asyncio.sleep(0) await client.send(b"still usable", "test") @@ -153,7 +105,7 @@ async def test_confirmation_timeout_and_cancellation_cleanup( with pytest.raises(stompman.SubscriptionError, match="timeout"): await task assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending assert any(isinstance(frame, stompman.UnsubscribeFrame) for frame in remaining_frames(outgoing)) @@ -166,7 +118,7 @@ async def test_connection_loss_fails_unconfirmed_subscription( with pytest.raises(stompman.SubscriptionError, match="connection_lost"): await task assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending @pytest.mark.parametrize("correlated", [True, False]) @@ -217,7 +169,7 @@ async def test_resubscription_failure_notifies_owner_and_cleans_up( error = await asyncio.wait_for(failure, timeout=1) assert error.reason == ("rejected" if reject else "timeout") assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending @pytest.mark.parametrize("cancel_send", [True, False]) @@ -263,7 +215,7 @@ async def delayed_write(connection: BaseMockConnection, frame: stompman.AnyClien send = tasks.create_task(client.send(b"after replay", "test")) await asyncio.sleep(0.2) assert not failures - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending assert client._active_subscriptions.get_all() == subscriptions assert not send.done() if cancel_send: @@ -318,7 +270,7 @@ async def blocked_write(connection: BaseMockConnection, frame: stompman.AnyClien await asyncio.wait_for(failed.wait(), timeout=1) assert [error.reason for error in failures] == ["connection_lost" if connection_lost else "timeout"] assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending await client.send(b"still usable", "test") assert all(not isinstance(frame, stompman.SubscribeFrame) for frame in remaining_frames(outgoing)) @@ -456,4 +408,4 @@ async def blocked_write(connection: BaseMockConnection, frame: stompman.AnyClien with pytest.raises(stompman.SubscriptionError, match="timeout"): await task assert not client._active_subscriptions.get_all() - assert not client._active_subscriptions.pending_receipts + assert not client._receipts.pending diff --git a/pyproject.toml b/pyproject.toml index 58219c0..57c3e66 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,15 +5,15 @@ version = "0" [dependency-groups] dev = [ - "anyio==4.14.2", - "faker==40.37.0", - "hypothesis==6.165.10", + "anyio==4.15.1", + "faker==40.39.0", + "hypothesis==6.168.0", "mypy==2.3.1", "polyfactory==3.3.0", "pytest==9.1.1", "pytest-cov==7.1.0", "pytest-timeout==2.4.0", - "ruff==0.16.5", + "ruff==0.16.8", "uvloop==0.22.1", ] ws = [ @@ -56,6 +56,8 @@ ignore = [ "PLC2701", "PLC2801", "PLR0913", + "RUF105", + "RUF201", "SLF001" ] extend-per-file-ignores = { "*/test_*/*" = ["S101", "SLF001", "ARG", "PLR6301"] }