From 8a5edfec40c0a5bfffc2b866d1281265f98bc7fa Mon Sep 17 00:00:00 2001 From: Patrick Ribbsaeter Date: Tue, 11 Aug 2026 15:37:58 +0200 Subject: [PATCH] Retry websocket connections closed with code 1001 Treat remote Going Away closures as recoverable while preserving deliberate client shutdown, with local-server regression coverage for reconnect and resubscribe. --- coinbase/websocket/websocket_base.py | 26 ++++++++++++++++++++++++++ tests/websocket/mock_ws_server.py | 16 +++++++++++++--- tests/websocket/test_websocket_base.py | 26 ++++++++++++++++++++++++++ 3 files changed, 65 insertions(+), 3 deletions(-) diff --git a/coinbase/websocket/websocket_base.py b/coinbase/websocket/websocket_base.py index a65b8f4..955053b 100644 --- a/coinbase/websocket/websocket_base.py +++ b/coinbase/websocket/websocket_base.py @@ -108,6 +108,7 @@ def __init__( self.subscriptions = {} self._background_exception = None self._retrying = False + self._closing = False def open(self) -> None: """ @@ -136,6 +137,7 @@ async def open_async(self) -> None: Open the websocket client connection asynchronously. """ self._ensure_websocket_not_open() + self._closing = False headers = self._set_headers() @@ -200,6 +202,7 @@ async def close_async(self) -> None: self._ensure_websocket_open() logger.debug("Closing connection to %s", self.base_url) + self._closing = True try: await self.websocket.close() self.websocket = None @@ -210,6 +213,7 @@ async def close_async(self) -> None: if self.on_close: self.on_close() except (websockets.exceptions.WebSocketException, OSError) as wse: + self._closing = False logger.error("Failed to close WebSocket connection: %s", wse) raise WSClientException("Failed to close WebSocket connection.") from wse @@ -508,6 +512,28 @@ async def _message_handler(self): self.on_message(message) except websockets.exceptions.ConnectionClosedOK as cco: logger.debug("Connection closed (OK): %s", cco) + close_code = getattr(cco.rcvd, "code", cco.rcvd) + if close_code == 1001 and not self._closing: + if self.on_close: + self.on_close() + + if self.retry: + self._retrying = True + try: + logger.debug("Retrying connection after Going Away closure") + await self._retry_connection() + self._retrying = False + continue + except WSClientException: + logger.error( + "Connection closed with Going Away status. Retry attempts failed." + ) + self._background_exception = WSClientConnectionClosedException( + "Connection closed with Going Away status. Retry attempts failed." + ) + self.subscriptions = {} + self._retrying = False + self._retry_count = 0 break except websockets.exceptions.ConnectionClosedError as cce: logger.error("Connection closed (ERROR): %s", cce) diff --git a/tests/websocket/mock_ws_server.py b/tests/websocket/mock_ws_server.py index beea62d..45df966 100644 --- a/tests/websocket/mock_ws_server.py +++ b/tests/websocket/mock_ws_server.py @@ -50,13 +50,23 @@ async def restart_with_error(self): await asyncio.sleep(1) # Short delay to ensure the port is freed up await self.start() + async def restart_with_going_away(self): + await self.trigger_connection_going_away() + await self.stop() + await asyncio.sleep(1) # Short delay to ensure the port is freed up + await self.start() + async def trigger_connection_closed_error(self): + await self.trigger_connection_closed(4000, "Abnormal closure") + + async def trigger_connection_going_away(self): + await self.trigger_connection_closed(1001, "Going away") + + async def trigger_connection_closed(self, code, reason): WebSocketTask = namedtuple("WebSocketTask", ["ws", "task"]) tasks = [ - WebSocketTask( - ws, asyncio.create_task(ws.close(code=4000, reason="Abnormal closure")) - ) + WebSocketTask(ws, asyncio.create_task(ws.close(code=code, reason=reason))) for ws in self.active_websockets ] await asyncio.gather(*(task.task for task in tasks)) diff --git a/tests/websocket/test_websocket_base.py b/tests/websocket/test_websocket_base.py index 2bbecec..4441482 100644 --- a/tests/websocket/test_websocket_base.py +++ b/tests/websocket/test_websocket_base.py @@ -458,6 +458,10 @@ async def mock_send(self, message): async def asyncSetUp(self): self.messages_queue = asyncio.Queue() + self.connection_closed_event = asyncio.Event() + self.on_close_mock = unittest.mock.Mock( + side_effect=self.connection_closed_event.set + ) self.server = await mock_ws_server.start_mock_server() def on_message(msg): @@ -468,6 +472,7 @@ def on_message(msg): TEST_API_SECRET, base_url="ws://localhost:8765", on_message=on_message, + on_close=self.on_close_mock, retry=False, ) @@ -524,6 +529,27 @@ async def test_reconnect(self): self.assertEqual(resubscribe_2_json["product_ids"], ["BTC-USD"]) self.assertEqual(resubscribe_2_json["channel"], "heartbeats") + async def test_reconnect_after_going_away(self): + self.ws.retry = True + + await self.ws.open_async() + await self.ws.subscribe_async( + product_ids=["BTC-USD", "ETH-USD"], channels=["ticker"] + ) + await self.messages_queue.get() + + await self.server.restart_with_going_away() + + await asyncio.wait_for(self.connection_closed_event.wait(), timeout=5) + self.on_close_mock.assert_called_once_with() + resubscribe = await asyncio.wait_for(self.messages_queue.get(), timeout=5) + resubscribe_json = json.loads(resubscribe) + self.assertEqual(resubscribe_json["type"], SUBSCRIBE_MESSAGE_TYPE) + self.assertEqual( + sorted(resubscribe_json["product_ids"]), ["BTC-USD", "ETH-USD"] + ) + self.assertEqual(resubscribe_json["channel"], "ticker") + async def test_reconnect_fail(self): # tests that client can catch WSClientConnectionClosedException after failed reconnection self.ws.retry = True