diff --git a/ably/realtime/connectionmanager.py b/ably/realtime/connectionmanager.py index 8b51fb0f..f910dfe6 100644 --- a/ably/realtime/connectionmanager.py +++ b/ably/realtime/connectionmanager.py @@ -2,8 +2,8 @@ import asyncio import logging +import time from collections import deque -from datetime import datetime from itertools import zip_longest from typing import TYPE_CHECKING @@ -120,12 +120,29 @@ def clear(self) -> None: self.messages.clear() +class PendingPing: + """Represents a ping awaiting its heartbeat echo from the server""" + + def __init__(self, id: str): + self.id = id + self.future: asyncio.Future[None] = asyncio.get_running_loop().create_future() + self.start_time: float = time.monotonic() + + @property + def message(self) -> dict: + return {"action": ProtocolMessageAction.HEARTBEAT, "id": self.id} + + @property + def response_time_ms(self) -> float: + return round((time.monotonic() - self.start_time) * 1000, 2) + + class ConnectionManager(EventEmitter): def __init__(self, realtime: AblyRealtime, initial_state): self.options = realtime.options self.__ably = realtime self.__state: ConnectionState = initial_state - self.__ping_future: asyncio.Future | None = None + self.__pending_pings: dict[str, PendingPing] = {} self.__timeout_in_secs: float = self.options.realtime_request_timeout / 1000 self.transport: WebSocketTransport | None = None self.__connection_details: ConnectionDetails | None = None @@ -150,6 +167,16 @@ def enact_state_change(self, state: ConnectionState, reason: AblyException | Non if reason: self.__error_reason = reason + # RTN13b: a ping cannot complete from these states, and the echo of a heartbeat already + # sent will never arrive, so fail pending pings rather than leaving them to time out + if state in ( + ConnectionState.SUSPENDED, + ConnectionState.CLOSING, + ConnectionState.CLOSED, + ConnectionState.FAILED, + ): + self.__fail_pending_pings(reason or ConnectionErrors[state]) + # RTN16d: Clear connection state when entering SUSPENDED or terminal states if state == ConnectionState.SUSPENDED or state in ( ConnectionState.CLOSED, @@ -332,29 +359,43 @@ def fail_queued_messages(self, err) -> None: self.pending_message_queue.complete_all_messages(error) async def ping(self) -> float: - if self.__ping_future: - try: - response = await self.__ping_future - except asyncio.CancelledError: - raise AblyException("Ping request cancelled due to request timeout", 504, 50003) from None - return response - - self.__ping_future = asyncio.Future() - if self.__state in [ConnectionState.CONNECTED, ConnectionState.CONNECTING]: - self.__ping_id = get_random_id() - ping_start_time = datetime.now().timestamp() - await self.send_protocol_message({"action": ProtocolMessageAction.HEARTBEAT, - "id": self.__ping_id}) - else: - raise AblyException("Cannot send ping request. Calling ping in invalid state", 40000, 400) + # RTN13b + if self.__state not in (ConnectionState.CONNECTED, ConnectionState.CONNECTING): + raise AblyException("Cannot send ping request. Calling ping in invalid state", 400, 40000) + + # RTN13e: the id tells this ping's echo apart from server heartbeats and other pings + pending_ping = PendingPing(get_random_id()) + self.__pending_pings[pending_ping.id] = pending_ping try: - await asyncio.wait_for(self.__ping_future, self.__timeout_in_secs) + # RTN13d: while connecting, the heartbeat goes out once the connection is established + if self.__state == ConnectionState.CONNECTED: + await self.__send_heartbeat(pending_ping) + # RTN13c + await asyncio.wait_for(pending_ping.future, self.__timeout_in_secs) except asyncio.TimeoutError: raise AblyException("Timeout waiting for ping response", 504, 50003) from None + finally: + self.__pending_pings.pop(pending_ping.id, None) + return pending_ping.response_time_ms - ping_end_time = datetime.now().timestamp() - response_time_ms = (ping_end_time - ping_start_time) * 1000 - return round(response_time_ms, 2) + async def __send_heartbeat(self, pending_ping: PendingPing) -> None: + pending_ping.start_time = time.monotonic() + try: + await self.send_protocol_message(pending_ping.message) + except Exception as error: + # the caller is waiting on the future, so the send failure is surfaced there + if not pending_ping.future.done(): + pending_ping.future.set_exception(error) + + def __send_pending_pings(self) -> None: + """RTN13d: send the heartbeat for each ping waiting on the connection + + A ping registered while the connection was establishing has not sent its heartbeat yet, + and one whose heartbeat went out before the transport went away will never see that echo, + so both are (re)sent here and their round trip is measured from this point. + """ + for pending_ping in list(self.__pending_pings.values()): + asyncio.create_task(self.__send_heartbeat(pending_ping)) def on_connected(self, connection_details: ConnectionDetails, connection_id: str, reason: AblyException | None = None) -> None: @@ -458,12 +499,18 @@ def on_channel_message(self, msg: dict) -> None: self.__ably.channels._on_channel_message(msg) def on_heartbeat(self, id: str | None) -> None: - if self.__ping_future: - # Resolve on heartbeat from ping request. - if self.__ping_id == id: - if not self.__ping_future.cancelled(): - self.__ping_future.set_result(None) - self.__ping_future = None + if id is None: + return + pending_ping = self.__pending_pings.pop(id, None) + # the echo can arrive while wait_for is still cancelling a timed-out ping + if pending_ping is not None and not pending_ping.future.done(): + pending_ping.future.set_result(None) + + def __fail_pending_pings(self, error: AblyException) -> None: + pending, self.__pending_pings = self.__pending_pings, {} + for pending_ping in pending.values(): + if not pending_ping.future.done(): + pending_ping.future.set_exception(error) def on_ack( self, serial: int, count: int, res: list[PublishResult] | None @@ -630,6 +677,7 @@ def notify_state(self, state: ConnectionState, reason: AblyException | None = No if state == ConnectionState.CONNECTED: self.send_queued_messages() + self.__send_pending_pings() elif state in ( ConnectionState.CLOSING, ConnectionState.CLOSED, diff --git a/test/ably/realtime/realtimeconnection_test.py b/test/ably/realtime/realtimeconnection_test.py index 2593eb3e..dbd20698 100644 --- a/test/ably/realtime/realtimeconnection_test.py +++ b/test/ably/realtime/realtimeconnection_test.py @@ -1,4 +1,5 @@ import asyncio +import time import pytest from websockets import connect as _ws_connect @@ -127,8 +128,8 @@ async def test_connection_ping_initialized(self): assert ably.connection.state == ConnectionState.INITIALIZED with pytest.raises(AblyException) as exception: await ably.connection.ping() - assert exception.value.code == 400 - assert exception.value.status_code == 40000 + assert exception.value.code == 40000 + assert exception.value.status_code == 400 async def test_connection_ping_failed(self): ably = await TestApp.get_ably_realtime(key=self.valid_key_format) @@ -136,8 +137,8 @@ async def test_connection_ping_failed(self): assert ably.connection.state == ConnectionState.FAILED with pytest.raises(AblyException) as exception: await ably.connection.ping() - assert exception.value.code == 400 - assert exception.value.status_code == 40000 + assert exception.value.code == 40000 + assert exception.value.status_code == 400 await ably.close() async def test_connection_ping_closed(self): @@ -147,8 +148,8 @@ async def test_connection_ping_closed(self): await ably.close() with pytest.raises(AblyException) as exception: await ably.connection.ping() - assert exception.value.code == 400 - assert exception.value.status_code == 40000 + assert exception.value.code == 40000 + assert exception.value.status_code == 400 async def test_auto_connect(self): ably = await TestApp.get_ably_realtime() @@ -212,6 +213,107 @@ async def new_send_protocol_message(protocol_message): assert exception.value.code == 50003 assert exception.value.status_code == 504 + + # one timed-out ping must not break later pings + ably.connection.connection_manager.send_protocol_message = original_send_protocol_message + response_time_ms = await asyncio.wait_for(ably.connection.ping(), timeout=5) + assert type(response_time_ms) is float + await ably.close() + + async def test_ping_after_invalid_state_ping(self): + # a ping rejected for bad state must not break later pings + ably = await TestApp.get_ably_realtime(auto_connect=False) + with pytest.raises(AblyException) as exception: + await ably.connection.ping() + assert exception.value.code == 40000 + + ably.connect() + await asyncio.wait_for(ably.connection.once_async(ConnectionState.CONNECTED), timeout=5) + response_time_ms = await asyncio.wait_for(ably.connection.ping(), timeout=5) + assert type(response_time_ms) is float + await ably.close() + + async def test_ping_survives_connection_drop(self): + # RTN13d: a heartbeat lost with a dropped connection is sent again once the connection is + # back, rather than the ping failing or waiting out realtime_request_timeout + outage_secs = 1 + async with WsProxy(self.test_vars["host"]) as proxy: + ably = await TestApp.get_ably_realtime( + realtime_request_timeout=20000, + tls=False, + endpoint=proxy.endpoint, + ) + try: + await asyncio.wait_for(ably.connection.once_async(ConnectionState.CONNECTED), timeout=10) + + connection_manager = ably.connection.connection_manager + original_send_protocol_message = connection_manager.send_protocol_message + + async def drop_heartbeats(protocol_message): + if protocol_message.get('action') == ProtocolMessageAction.HEARTBEAT: + return + await original_send_protocol_message(protocol_message) + + connection_manager.send_protocol_message = drop_heartbeats + ping_start_time = time.monotonic() + ping = asyncio.ensure_future(ably.connection.ping()) + await asyncio.sleep(outage_secs) + assert not ping.done() + + # Simulate server sending a normal WS close frame + connection_manager.send_protocol_message = original_send_protocol_message + await proxy.close_active_connection() + + response_time_ms = await asyncio.wait_for(ping, timeout=15) + elapsed_ms = (time.monotonic() - ping_start_time) * 1000 + assert type(response_time_ms) is float + # the round trip is measured from the heartbeat that was sent on reconnection, + # so it excludes the time spent waiting for the connection + assert response_time_ms < elapsed_ms - outage_secs * 1000 + finally: + await ably.close() + + async def test_ping_fails_when_connection_closes(self): + # RTN13b: a ping cannot complete once the connection has transitioned to CLOSING + ably = await TestApp.get_ably_realtime(realtime_request_timeout=20000) + await asyncio.wait_for(ably.connection.once_async(ConnectionState.CONNECTED), timeout=5) + + connection_manager = ably.connection.connection_manager + original_send_protocol_message = connection_manager.send_protocol_message + + async def drop_heartbeats(protocol_message): + if protocol_message.get('action') == ProtocolMessageAction.HEARTBEAT: + return + await original_send_protocol_message(protocol_message) + + connection_manager.send_protocol_message = drop_heartbeats + ping = asyncio.ensure_future(ably.connection.ping()) + await asyncio.sleep(0.5) + assert not ping.done() + + await ably.close() + + with pytest.raises(AblyException) as exception: + await asyncio.wait_for(ping, timeout=5) + assert exception.value.code == 80017 + assert exception.value.status_code == 400 + + async def test_ping_while_connecting(self): + # RTN13d: a ping requested while connecting is done once the connection is established + ably = await TestApp.get_ably_realtime(auto_connect=False) + ably.connect() + assert ably.connection.state == ConnectionState.CONNECTING + response_time_ms = await asyncio.wait_for(ably.connection.ping(), timeout=15) + assert type(response_time_ms) is float + await ably.close() + + async def test_concurrent_pings(self): + ably = await TestApp.get_ably_realtime() + await asyncio.wait_for(ably.connection.once_async(ConnectionState.CONNECTED), timeout=5) + results = await asyncio.wait_for( + asyncio.gather(ably.connection.ping(), ably.connection.ping(), ably.connection.ping()), timeout=5 + ) + assert all(type(response_time_ms) is float for response_time_ms in results) await ably.close() async def test_disconnected_retry_timeout(self):