Add some configurations for task based subscription processing

- Propagate subscription manager tasks exceptions to the main loop.
- Allow configuring exceptions to be logged silently, rather than raised.
- Add task cleanup.

Author: fselmo

tests/core/subscriptions/test_subscription_manager.py

@@ -1,5 +1,8 @@
 import pytest
+import asyncio
 import itertools
+import logging
+import time
 from unittest.mock import (
     AsyncMock,
 )
@@ -14,8 +17,12 @@
     PersistentConnectionProvider,
 )
 from web3.exceptions import (
+    SubscriptionHandlerTaskException,
     Web3ValueError,
 )
+from web3.providers.persistent.request_processor import (
+    TaskReliantQueue,
+)
 from web3.providers.persistent.subscription_manager import (
     SubscriptionManager,
 )
@@ -213,3 +220,161 @@ async def test_unsubscribe_with_subscriptions_reference_does_not_mutate_the_list
 
     await subscription_manager.unsubscribe_all()
     assert subscription_manager.subscriptions == []
+
+
+@pytest.mark.asyncio
+async def test_high_throughput_subscription_task_based(
+    subscription_manager,
+) -> None:
+    provider = subscription_manager._w3.provider
+    num_msgs = 5_000
+
+    provider._request_processor._handler_subscription_queue = TaskReliantQueue(
+        maxsize=num_msgs
+    )
+
+    # Turn on task-based processing. This test should fail the time constraint if this
+    # is not set to ``True`` (not task-based processing).
+    subscription_manager.task_based = True
+
+    class Counter:
+        val: int = 0
+
+    counter = Counter()
+
+    async def high_throughput_handler(handler_context) -> None:
+        # if we awaited all `num_msgs`, we would sleep at least 5 seconds total
+        await asyncio.sleep(5 / num_msgs)
+
+        handler_context.counter.val += 1
+        if handler_context.counter.val == num_msgs:
+            await handler_context.subscription.unsubscribe()
+
+    # build a meaningless subscription since we are fabricating the messages
+    sub_id = await subscription_manager.subscribe(
+        NewHeadsSubscription(
+            handler=high_throughput_handler, handler_context={"counter": counter}
+        ),
+    )
+    provider._request_processor.cache_request_information(
+        request_id=sub_id,
+        method="eth_subscribe",
+        params=[],
+        response_formatters=((), (), ()),
+    )
+
+    # put `num_msgs` messages in the queue
+    for _ in range(num_msgs):
+        provider._request_processor._handler_subscription_queue.put_nowait(
+            {
+                "jsonrpc": "2.0",
+                "method": "eth_subscription",
+                "params": {"subscription": sub_id, "result": "0x0"},
+            }
+        )
+
+    start = time.time()
+    await subscription_manager.handle_subscriptions()
+    stop = time.time()
+
+    assert counter.val == num_msgs
+
+    assert subscription_manager.total_handler_calls == num_msgs
+    assert stop - start < 3, "subscription handling took too long!"
+
+
+@pytest.mark.asyncio
+async def test_task_based_subscription_handling_error_propagation(
+    subscription_manager,
+) -> None:
+    provider = subscription_manager._w3.provider
+    subscription_manager.task_based = True
+
+    async def high_throughput_handler(_handler_context) -> None:
+        raise ValueError("Test error msg.")
+
+    # build a meaningless subscription since we are fabricating the messages
+    sub_id = await subscription_manager.subscribe(
+        NewHeadsSubscription(handler=high_throughput_handler)
+    )
+    provider._request_processor.cache_request_information(
+        request_id=sub_id,
+        method="eth_subscribe",
+        params=[],
+        response_formatters=((), (), ()),
+    )
+    provider._request_processor._handler_subscription_queue.put_nowait(
+        {
+            "jsonrpc": "2.0",
+            "method": "eth_subscription",
+            "params": {"subscription": sub_id, "result": "0x0"},
+        }
+    )
+
+    with pytest.raises(
+        SubscriptionHandlerTaskException,
+        match="Test error msg.",
+    ):
+        await subscription_manager.handle_subscriptions()
+
+
+@pytest.mark.asyncio
+async def test_task_based_subscription_handling_ignore_errors(
+    subscription_manager, caplog
+) -> None:
+    provider = subscription_manager._w3.provider
+    subscription_manager.task_based = True
+    subscription_manager.ignore_task_exceptions = True
+
+    class TestObject:
+        exception = None
+
+    async def sub_handler(handler_context) -> None:
+        if handler_context.obj.exception:
+            # on the second call, yield to loop so we log, unsubscribe, and return
+            await asyncio.sleep(0.01)
+            await handler_context.subscription.unsubscribe()
+            return
+
+        e = ValueError("Test error msg.")
+        handler_context.obj.exception = e
+        raise e
+
+    # build a meaningless subscription since we are fabricating the messages
+    test_obj = TestObject()
+    sub_id = await subscription_manager.subscribe(
+        NewHeadsSubscription(handler=sub_handler, handler_context={"obj": test_obj})
+    )
+    provider._request_processor.cache_request_information(
+        request_id=sub_id,
+        method="eth_subscribe",
+        params=[],
+        response_formatters=((), (), ()),
+    )
+    for _ in range(2):
+        provider._request_processor._handler_subscription_queue.put_nowait(
+            {
+                "jsonrpc": "2.0",
+                "method": "eth_subscription",
+                "params": {"subscription": sub_id, "result": "0x0"},
+            }
+        )
+
+    with caplog.at_level(
+        logging.WARNING, logger="web3.providers.persistent.subscription_manager"
+    ):
+        await subscription_manager.handle_subscriptions()
+
+    # find the warning so and assert it was logged
+    warning_records = [r for r in caplog.records if r.levelname == "WARNING"]
+    assert len(warning_records) == 1
+    record = warning_records[0]
+    assert (
+        "An exception occurred in a subscription handler task but was ignored, `"
+        "`ignore_task_exceptions==True``." in record.message
+    )
+    await subscription_manager.handle_subscriptions()
+
+    assert subscription_manager.total_handler_calls == 2
+    assert subscription_manager.subscriptions == []
+    assert subscription_manager._tasks == set()

web3/_utils/module_testing/persistent_connection_provider.py

@@ -3,7 +3,6 @@
 from dataclasses import (
     dataclass,
 )
-import time
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -39,9 +38,6 @@
 from web3.middleware import (
     ExtraDataToPOAMiddleware,
 )
-from web3.providers.persistent.request_processor import (
-    TaskReliantQueue,
-)
 from web3.types import (
     BlockData,
     FormattedEthSubscriptionResponse,
@@ -918,59 +914,3 @@ async def test_run_forever_starts_with_0_subs_and_runs_until_task_cancelled(
 
         # cleanup
         await clean_up_task(run_forever_task)
-
-    @pytest.mark.asyncio
-    async def test_high_throughput_subscription_task_based(
-        self, async_w3: AsyncWeb3
-    ) -> None:
-        async_w3.provider._request_processor._handler_subscription_queue = (
-            TaskReliantQueue(maxsize=5_000)
-        )
-        sub_manager = async_w3.subscription_manager
-        sub_manager.task_based = True  # turn on task-based processing
-
-        class Counter:
-            val: int = 0
-
-        counter = Counter()
-
-        async def high_throughput_handler(
-            handler_context: Any,
-        ) -> None:
-            handler_context.counter.val += 1
-            if handler_context.counter.val == 5_000:
-                await handler_context.subscription.unsubscribe()
-            # if we awaited all 5_000 messages, we would sleep at least 5 seconds
-            await asyncio.sleep(5 // 5_000)
-
-        # build a meaningless subscription since we are fabricating the messages
-        sub_id = await async_w3.eth.subscribe(
-            "syncing",
-            handler=high_throughput_handler,
-            handler_context={"counter": counter},
-        )
-        async_w3.provider._request_processor.cache_request_information(
-            request_id=sub_id,
-            method=RPCEndpoint("eth_subscribe"),
-            params=[],
-            response_formatters=((), (), ()),  # type: ignore
-        )
-
-        # put 5_000 messages in the queue
-        for _ in range(5_000):
-            async_w3.provider._request_processor._handler_subscription_queue.put_nowait(
-                {
-                    "jsonrpc": "2.0",
-                    "method": "eth_subscription",
-                    "params": {"subscription": HexBytes(sub_id), "result": False},
-                }
-            )
-
-        start = time.time()
-        await sub_manager.handle_subscriptions()
-        stop = time.time()
-
-        assert counter.val == 5_000
-
-        assert sub_manager.total_handler_calls == 5_000
-        assert stop - start < 3, "subscription handling took too long!"

web3/exceptions.py

@@ -353,13 +353,20 @@ class PersistentConnectionClosedOK(PersistentConnectionError):
     """
 
 
-class SubscriptionProcessingFinished(Web3Exception):
+class SubscriptionProcessingFinished(PersistentConnectionError):
     """
     Raised to alert the subscription manager that the processing of subscriptions
     has finished.
     """
 
 
+class SubscriptionHandlerTaskException(TaskNotRunning):
+    """
+    Raised to alert the subscription manager that an exception occurred in the
+    subscription processing task.
+    """
+
+
 class Web3RPCError(Web3Exception):
     """
     Raised when a JSON-RPC response contains an error field.

web3/manager.py

@@ -550,13 +550,10 @@ async def _message_stream(
                 else:
                     # if not an active sub, skip processing and continue
                     continue
-            except TaskNotRunning:
+            except TaskNotRunning as e:
                 await asyncio.sleep(0)
                 self._provider._handle_listener_task_exceptions()
-                self.logger.error(
-                    "Message listener background task has stopped unexpectedly. "
-                    "Stopping message stream."
-                )
+                self.logger.error("Stopping message stream: %s", e.message)
                 return
 
     async def _process_response(

web3/providers/persistent/persistent.py

@@ -164,7 +164,7 @@ async def send_batch_func(
         if cache_key != self._send_batch_func_cache[0]:
 
             async def send_func(
-                requests: List[Tuple[RPCEndpoint, Any]]
+                requests: List[Tuple[RPCEndpoint, Any]],
             ) -> List[RPCRequest]:
                 for mw in middleware:
                     initialized = mw(async_w3)
@@ -376,11 +376,12 @@ def _message_listener_callback(
     ) -> None:
         # Puts a `TaskNotRunning` in appropriate queues to signal the end of the
         # listener task to any listeners relying on the queues.
+        message = "Message listener task has ended."
         self._request_processor._subscription_response_queue.put_nowait(
-            TaskNotRunning(message_listener_task)
+            TaskNotRunning(message_listener_task, message=message)
         )
         self._request_processor._handler_subscription_queue.put_nowait(
-            TaskNotRunning(message_listener_task)
+            TaskNotRunning(message_listener_task, message=message)
         )
 
     def _raise_stray_errors_from_cache(self) -> None: