diff --git a/ably/realtime/connectionmanager.py b/ably/realtime/connectionmanager.py index 8b51fb0f..8a479936 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 @@ -125,7 +125,7 @@ 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, asyncio.Future[None]] = {} self.__timeout_in_secs: float = self.options.realtime_request_timeout / 1000 self.transport: WebSocketTransport | None = None self.__connection_details: ConnectionDetails | None = None @@ -332,29 +332,21 @@ 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) + 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 + ping_id = get_random_id() + echo = self.__pending_pings[ping_id] = asyncio.get_running_loop().create_future() + start_time = time.monotonic() try: - await asyncio.wait_for(self.__ping_future, self.__timeout_in_secs) + await self.send_protocol_message({"action": ProtocolMessageAction.HEARTBEAT, "id": ping_id}) + await asyncio.wait_for(echo, self.__timeout_in_secs) except asyncio.TimeoutError: raise AblyException("Timeout waiting for ping response", 504, 50003) from None - - ping_end_time = datetime.now().timestamp() - response_time_ms = (ping_end_time - ping_start_time) * 1000 - return round(response_time_ms, 2) + finally: + self.__pending_pings.pop(ping_id, None) + return round((time.monotonic() - start_time) * 1000, 2) def on_connected(self, connection_details: ConnectionDetails, connection_id: str, reason: AblyException | None = None) -> None: @@ -458,12 +450,10 @@ 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 + echo = self.__pending_pings.pop(id, None) + # the echo can arrive while wait_for is still cancelling a timed-out ping + if echo is not None and not echo.done(): + echo.set_result(None) def on_ack( self, serial: int, count: int, res: list[PublishResult] | None diff --git a/test/ably/realtime/realtimeconnection_test.py b/test/ably/realtime/realtimeconnection_test.py index 2593eb3e..133212cc 100644 --- a/test/ably/realtime/realtimeconnection_test.py +++ b/test/ably/realtime/realtimeconnection_test.py @@ -127,8 +127,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 +136,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 +147,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 +212,33 @@ 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_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):