Retry websocket connections closed with code 1001

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)

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))

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