Gracefully exit message iterator if ConnectionClosedOK is raised

- Similar changes as #3424 that were implemented for ``v6``.
- Raise ``PersistentConnectionClosedOK`` if a ``ConnectionClosedOK`` is
  raised by the server for a websocket provider. Handle this appropriately
  in a pattern than can be applicable to other persistent connections
  by creating a new internal exception that is then handled, rather than
  handling the specific websocket exception.

Author: fselmo

newsfragments/3432.bugfix.rst

@@ -0,0 +1 @@
+Handle ``ConnectionClosedOK`` case for ``WebSocketProvider``. If a persistent connection is closed gracefully, log and raise a silent ``PersistentConnectionClosedOK`` exception, triggering an end to the message listener task and breaking out of the ``process_subscriptions()`` iterator.

tests/core/providers/test_websocket_provider.py

@@ -12,6 +12,7 @@
 )
 from websockets import (
     ConnectionClosed,
+    ConnectionClosedOK,
 )
 
 from web3 import (
@@ -365,3 +366,16 @@ async def test_async_iterator_pattern_exception_handling_for_subscriptions():
             pytest.fail("Expected `ConnectionClosed` exception.")
 
         assert iterations == 3
+
+
+@pytest.mark.asyncio
+async def test_connection_closed_ok_breaks_message_iteration():
+    with patch(
+        "web3.providers.persistent.websocket.connect",
+        new=lambda *_1, **_2: WebSocketMessageStreamMock(
+            raise_exception=ConnectionClosedOK(None, None)
+        ),
+    ):
+        w3 = await AsyncWeb3(WebSocketProvider("ws://mocked"))
+        async for _ in w3.socket.process_subscriptions():
+            pytest.fail("Should not reach this point.")

web3/_utils/module_testing/module_testing_utils.py

@@ -1,6 +1,4 @@
-from collections import (
-    deque,
-)
+import asyncio
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -179,7 +177,9 @@ class WebSocketMessageStreamMock:
     def __init__(
         self, messages: Collection[bytes] = None, raise_exception: Exception = None
     ) -> None:
-        self.messages = deque(messages) if messages else deque()
+        self.queue = asyncio.Queue()  # type: ignore  # py38 issue
+        for msg in messages or []:
+            self.queue.put_nowait(msg)
         self.raise_exception = raise_exception
 
     def __await__(self) -> Generator[Any, Any, "Self"]:
@@ -192,13 +192,12 @@ def __aiter__(self) -> "Self":
         return self
 
     async def __anext__(self) -> bytes:
+        return await self.queue.get()
+
+    async def recv(self) -> bytes:
         if self.raise_exception:
             raise self.raise_exception
-
-        elif len(self.messages) == 0:
-            raise StopAsyncIteration
-
-        return self.messages.popleft()
+        return await self.queue.get()
 
     @staticmethod
     async def pong() -> Literal[False]:

web3/exceptions.py

@@ -327,6 +327,12 @@ def __init__(
         super().__init__(message)
 
 
+class PersistentConnectionClosedOK(Web3Exception):
+    """
+    Raised when a persistent connection is closed gracefully by the server.
+    """
+
+
 class Web3RPCError(Web3Exception):
     """
     Raised when a JSON-RPC response contains an error field.

web3/manager.py

@@ -20,9 +20,6 @@
 from hexbytes import (
     HexBytes,
 )
-from websockets.exceptions import (
-    ConnectionClosedOK,
-)
 
 from web3._utils.batching import (
     RequestBatcher,
@@ -487,7 +484,7 @@ async def send(self, method: RPCEndpoint, params: Any) -> RPCResponse:
     def _persistent_message_stream(self) -> "_AsyncPersistentMessageStream":
         return _AsyncPersistentMessageStream(self)
 
-    async def _get_next_message(self) -> Any:
+    async def _get_next_message(self) -> Optional[RPCResponse]:
         return await self._message_stream().__anext__()
 
     async def _message_stream(self) -> AsyncGenerator[RPCResponse, None]:
@@ -515,12 +512,13 @@ async def _message_stream(self) -> AsyncGenerator[RPCResponse, None]:
                     # if response is an active subscription response, process it
                     yield await self._process_response(response)
             except TaskNotRunning:
+                await asyncio.sleep(0)
                 self._provider._handle_listener_task_exceptions()
                 self.logger.error(
                     "Message listener background task has stopped unexpectedly. "
                     "Stopping message stream."
                 )
-                raise StopAsyncIteration
+                return
 
     async def _process_response(self, response: RPCResponse) -> RPCResponse:
         provider = cast(PersistentConnectionProvider, self._provider)
@@ -586,7 +584,4 @@ def __aiter__(self) -> Self:
         return self
 
     async def __anext__(self) -> RPCResponse:
-        try:
-            return await self.manager._get_next_message()
-        except ConnectionClosedOK:
-            raise StopAsyncIteration
+        return await self.manager._get_next_message()

web3/providers/persistent/persistent.py

@@ -19,6 +19,7 @@
     generate_cache_key,
 )
 from web3.exceptions import (
+    PersistentConnectionClosedOK,
     ProviderConnectionError,
     TaskNotRunning,
     TimeExhausted,
@@ -160,6 +161,13 @@ async def _message_listener(self) -> None:
             await asyncio.sleep(0)
             try:
                 await self._provider_specific_message_listener()
+            except PersistentConnectionClosedOK as e:
+                self.logger.info(
+                    "Message listener background task has ended gracefully: "
+                    f"{e.user_message}"
+                )
+                # trigger a return to end the listener task and initiate the callback fn
+                return
             except Exception as e:
                 if not self.silence_listener_task_exceptions:
                     raise e
@@ -202,10 +210,6 @@ async def _match_response_id_to_request_id() -> RPCResponse:
             request_cache_key = generate_cache_key(request_id)
 
             while True:
-                # check if an exception was recorded in the listener task and raise it
-                # in the main loop if so
-                self._handle_listener_task_exceptions()
-
                 if request_cache_key in self._request_processor._request_response_cache:
                     self.logger.debug(
                         f"Popping response for id {request_id} from cache."
@@ -215,6 +219,9 @@ async def _match_response_id_to_request_id() -> RPCResponse:
                     )
                     return popped_response
                 else:
+                    # check if an exception was recorded in the listener task and raise
+                    # it in the main loop if so
+                    self._handle_listener_task_exceptions()
                     await asyncio.sleep(0)
 
         try:

web3/providers/persistent/persistent_connection.py

@@ -2,6 +2,7 @@
     TYPE_CHECKING,
     Any,
     Dict,
+    Optional,
 )
 
 from web3.types import (
@@ -35,7 +36,7 @@ def subscriptions(self) -> Dict[str, Any]:
     async def send(self, method: RPCEndpoint, params: Any) -> RPCResponse:
         return await self._manager.send(method, params)
 
-    async def recv(self) -> Any:
+    async def recv(self) -> Optional[RPCResponse]:
         return await self._manager._get_next_message()
 
     def process_subscriptions(self) -> "_AsyncPersistentMessageStream":

web3/providers/persistent/websocket.py

@@ -25,6 +25,7 @@
     connect,
 )
 from websockets.exceptions import (
+    ConnectionClosedOK,
     WebSocketException,
 )
 
@@ -36,6 +37,7 @@
     async_handle_request_caching,
 )
 from web3.exceptions import (
+    PersistentConnectionClosedOK,
     ProviderConnectionError,
     Web3ValidationError,
 )
@@ -169,8 +171,13 @@ async def make_batch_request(
         return response
 
     async def _provider_specific_message_listener(self) -> None:
-        async for raw_message in self._ws:
-            await asyncio.sleep(0)
+        while True:
+            try:
+                raw_message = await self._ws.recv()
+            except ConnectionClosedOK:
+                raise PersistentConnectionClosedOK(
+                    user_message="WebSocket connection received `ConnectionClosedOK`."
+                )
 
             response = json.loads(raw_message)
             if isinstance(response, list):