Separate handler subscriptions into their own queue:

- Make use of a ``TaskReliant`` queue that is specific for handler
  subscriptions. Put a ``SubscriptionProcessingFinished`` exceptions
  in the queue to signal when all handler subscriptions have been
  unsubscribed from.

- Add testing around running the ``handler_subscriptions`` method while
  listening for un-handled subscriptions via the socket.

- Create a container to manage all subscriptions, since we store subs
  in more than one way: by ``id``, by ``label``, etc...

- The ``SubscriptionContainer`` class has an object reference, so we changes
  share this with the provider's ``RequestProcessor`` class. The
  request_processor can use the container to make decisions about which queue
  to store the subscription in, based on the presence or absence of a
  ``handler``.

Author: fselmo

tests/core/providers/test_async_ipc_provider.py

@@ -180,9 +180,11 @@ async def test_disconnect_cleanup(
     provider._request_processor._request_response_cache.cache("0", "0x1337")
     provider._request_processor._request_information_cache.cache("0", "0x1337")
     provider._request_processor._subscription_response_queue.put_nowait({"id": "0"})
+    provider._request_processor._handler_subscription_queue.put_nowait({"id": "0"})
     assert len(provider._request_processor._request_response_cache) == 1
     assert len(provider._request_processor._request_information_cache) == 1
     assert provider._request_processor._subscription_response_queue.qsize() == 1
+    assert provider._request_processor._handler_subscription_queue.qsize() == 1
 
     await w3.provider.disconnect()
 
@@ -192,6 +194,7 @@ async def test_disconnect_cleanup(
     assert len(provider._request_processor._request_response_cache) == 0
     assert len(provider._request_processor._request_information_cache) == 0
     assert provider._request_processor._subscription_response_queue.empty()
+    assert provider._request_processor._handler_subscription_queue.empty()
 
 
 async def _raise_connection_closed(*_args, **_kwargs):

tests/core/providers/test_websocket_provider.py

@@ -34,6 +34,9 @@
 from web3.types import (
     RPCEndpoint,
 )
+from web3.utils import (
+    EthSubscription,
+)
 
 
 def _mock_ws(provider):
@@ -80,16 +83,19 @@ async def test_disconnect_cleanup():
     provider._request_processor._request_response_cache.cache("0", "0x1337")
     provider._request_processor._request_information_cache.cache("0", "0x1337")
     provider._request_processor._subscription_response_queue.put_nowait({"id": "0"})
+    provider._request_processor._handler_subscription_queue.put_nowait({"id": "0"})
     assert len(provider._request_processor._request_response_cache) == 1
     assert len(provider._request_processor._request_information_cache) == 1
     assert provider._request_processor._subscription_response_queue.qsize() == 1
+    assert provider._request_processor._handler_subscription_queue.qsize() == 1
 
     await provider.disconnect()
 
     assert provider._ws is None
     assert len(provider._request_processor._request_response_cache) == 0
     assert len(provider._request_processor._request_information_cache) == 0
     assert provider._request_processor._subscription_response_queue.empty()
+    assert provider._request_processor._handler_subscription_queue.empty()
 
 
 @pytest.mark.asyncio
@@ -278,6 +284,9 @@ async def test_listen_event_awaits_msg_processing_when_subscription_queue_is_ful
         generate_cache_key(sub_id),
         sub_request_information,
     )
+    sub = EthSubscription()
+    sub._id = sub_id
+    async_w3.subscription_manager._add_subscription(sub)
 
     mocked_sub = {
         "jsonrpc": "2.0",
@@ -381,3 +390,45 @@ async def test_connection_closed_ok_breaks_message_iteration():
         w3 = await AsyncWeb3(WebSocketProvider("ws://mocked"))
         async for _ in w3.socket.process_subscriptions():
             pytest.fail("Should not reach this point.")
+
+
+@pytest.mark.asyncio
+async def test_listener_task_breaks_out_of_stream_when_cancelled():
+    with patch(
+        "web3.providers.persistent.websocket.connect",
+        new=lambda *_1, **_2: _mocked_ws_conn(),
+    ):
+        async_w3 = await AsyncWeb3(WebSocketProvider("ws://mocked"))
+
+    async_w3.provider._message_listener_task.cancel()
+    sub = EthSubscription()
+    sub._id = "0x1"
+    async_w3.subscription_manager._add_subscription(sub)
+    # this should hang indefinitely if the listener task does not put a
+    # ``TaskNotRunning`` in the ``_subscription_response_queue`` to break out of
+    # listening. The call to ``provider._handle_listener_task_exceptions`` bubbles up
+    # the exception.
+    with pytest.raises(asyncio.CancelledError):
+        async for _ in async_w3.socket.process_subscriptions():
+            ...
+
+
+@pytest.mark.asyncio
+async def test_listener_task_breaks_out_of_handle_subscriptions_when_cancelled():
+    with patch(
+        "web3.providers.persistent.websocket.connect",
+        new=lambda *_1, **_2: _mocked_ws_conn(),
+    ):
+        async_w3 = await AsyncWeb3(WebSocketProvider("ws://mocked"))
+
+    async_w3.provider._message_listener_task.cancel()
+    sub = EthSubscription(handler=AsyncMock())
+    sub._id = "0x1"
+
+    async_w3.subscription_manager._add_subscription(sub)
+    # this should hang indefinitely if the listener task does not put a
+    # ``TaskNotRunning`` in the ``_handler_subscription_queue`` to break out of
+    # listening. The call to ``provider._handle_listener_task_exceptions`` bubbles
+    # up the exception.
+    with pytest.raises(asyncio.CancelledError):
+        await async_w3.subscription_manager.handle_subscriptions(run_forever=True)

tests/core/subscriptions/test_subscription_manager.py

@@ -56,8 +56,11 @@ async def test_subscription_manager_raises_for_sub_with_the_same_label(
 
     # make sure the subscription was subscribed to and not added to the manager
     assert subscription_manager.subscriptions == [sub1]
-    assert subscription_manager._subscriptions_by_label == {"foo": sub1}
-    assert subscription_manager._subscriptions_by_id == {"0x0": sub1}
+    sub_container = subscription_manager._subscription_container
+    assert len(sub_container) == 1
+    assert sub_container.subscriptions == [sub1]
+    assert sub_container.subscriptions_by_id == {"0x0": sub1}
+    assert sub_container.subscriptions_by_label == {"foo": sub1}
 
 
 @pytest.mark.asyncio
@@ -97,8 +100,13 @@ async def test_unsubscribe_all_clears_all_subscriptions(subscription_manager):
     sub1 = NewHeadsSubscription(label="foo")
     sub2 = PendingTxSubscription(label="bar")
     await subscription_manager.subscribe([sub1, sub2])
+    assert subscription_manager.subscriptions == [sub1, sub2]
 
     await subscription_manager.unsubscribe_all()
     assert subscription_manager.subscriptions == []
-    assert subscription_manager._subscriptions_by_id == {}
-    assert subscription_manager._subscriptions_by_label == {}
+
+    sub_container = subscription_manager._subscription_container
+    assert len(sub_container) == 0
+    assert sub_container.subscriptions == []
+    assert sub_container.subscriptions_by_id == {}
+    assert sub_container.subscriptions_by_label == {}

web3/_utils/module_testing/persistent_connection_provider.py

@@ -7,7 +7,9 @@
     TYPE_CHECKING,
     Any,
     Dict,
+    List,
     Tuple,
+    Union,
     cast,
 )
 
@@ -61,6 +63,9 @@
         AsyncContract,
         AsyncContractFunction,
     )
+    from web3.providers.persistent.subscription_manager import (
+        SubscriptionContainer,
+    )
 
 
 # LogIndexedAndNotIndexed event args
@@ -158,6 +163,12 @@ async def logs_handler(
     assert await sub.unsubscribe()
 
 
+async def idle_handler(
+    _handler_context: Any,
+) -> None:
+    pass
+
+
 async def emit_contract_event(
     async_w3: AsyncWeb3,
     acct: ChecksumAddress,
@@ -194,6 +205,14 @@ async def log_indexed_and_non_indexed_args_task(
     )
 
 
+def assert_no_subscriptions_left(sub_container: "SubscriptionContainer") -> None:
+    assert len(sub_container) == 0
+    assert len(sub_container.subscriptions) == 0
+    assert len(sub_container.subscriptions_by_id) == 0
+    assert len(sub_container.subscriptions_by_label) == 0
+    assert len(sub_container.handler_subscriptions) == 0
+
+
 class PersistentConnectionProviderTest:
     @staticmethod
     async def seed_transactions_to_geth(
@@ -295,6 +314,13 @@ async def test_async_eth_subscribe_syncing_mocked(
             await async_w3.eth.unsubscribe(sub_id)
             break
 
+        assert_no_subscriptions_left(
+            async_w3.subscription_manager._subscription_container
+        )
+
+        # cleanup
+        async_w3.provider._request_processor.clear_caches()
+
     @pytest.mark.asyncio
     async def test_async_eth_subscribe_new_heads(self, async_w3: AsyncWeb3) -> None:
         sub_id = await async_w3.eth.subscribe("newHeads")
@@ -308,14 +334,16 @@ async def test_async_eth_subscribe_new_heads(self, async_w3: AsyncWeb3) -> None:
             break
 
         assert await async_w3.eth.unsubscribe(sub_id)
+        assert_no_subscriptions_left(
+            async_w3.subscription_manager._subscription_container
+        )
 
     @pytest.mark.asyncio
     async def test_async_eth_subscribe_creates_and_handles_new_heads_subscription_type(
         self,
         async_w3: AsyncWeb3,
     ) -> None:
         sub_manager = async_w3.subscription_manager
-
         new_heads_handler_test = SubscriptionHandlerTest()
 
         sub_id = await async_w3.eth.subscribe(
@@ -376,9 +404,9 @@ async def test_async_eth_subscribe_new_and_process_pending_tx_true(
 
         # cleanup
         assert await async_w3.eth.unsubscribe(sub_id)
-        assert len(async_w3.subscription_manager.subscriptions) == 0
-        assert len(async_w3.subscription_manager._subscriptions_by_id) == 0
-        assert len(async_w3.subscription_manager._subscriptions_by_label) == 0
+        assert_no_subscriptions_left(
+            async_w3.subscription_manager._subscription_container
+        )
         async_w3.provider._request_processor.clear_caches()
         await async_w3.eth.wait_for_transaction_receipt(tx_hash)
         tx_seeder_task.cancel()
@@ -414,9 +442,9 @@ async def test_async_eth_subscribe_and_process_pending_tx_false(
 
         # cleanup
         await async_w3.eth.unsubscribe(sub_id)
-        assert len(async_w3.subscription_manager.subscriptions) == 0
-        assert len(async_w3.subscription_manager._subscriptions_by_id) == 0
-        assert len(async_w3.subscription_manager._subscriptions_by_label) == 0
+        assert_no_subscriptions_left(
+            async_w3.subscription_manager._subscription_container
+        )
         await async_w3.eth.wait_for_transaction_receipt(tx_hash)
         tx_seeder_task.cancel()
 
@@ -762,3 +790,76 @@ async def test_sub_handler(
 
         # cleanup
         sub_manager.total_handler_calls = 0
+
+    @pytest.mark.asyncio
+    async def test_subscriptions_with_handler_and_without(
+        self, async_w3: AsyncWeb3
+    ) -> None:
+        handler_test = SubscriptionHandlerTest()
+        stream_passed = False
+
+        async def test_sub_handler(
+            handler_context: NewHeadsSubscriptionContext,
+        ) -> None:
+            handler_context.handler_test.passed = True
+            await handler_context.subscription.unsubscribe()
+
+        async def handle_subscription_stream() -> None:
+            nonlocal stream_passed
+            async for msg in async_w3.socket.process_subscriptions():
+                response = cast(FormattedEthSubscriptionResponse, msg)
+                assert sub_manager.get_by_id(response["subscription"]) is not None
+                assert response["result"] is not None
+                # wait for the handler to unsubscribe:
+                stream_passed = True
+                await async_w3.eth.unsubscribe(response["subscription"])
+                break
+
+        await async_w3.eth.subscribe(
+            "newHeads",
+            handler=test_sub_handler,
+            label="managed",
+            handler_context={"handler_test": handler_test},
+        )
+        await async_w3.eth.subscribe("newHeads", label="streamed")
+
+        sub_manager = async_w3.subscription_manager
+        assert len(sub_manager.subscriptions) == 2
+
+        await asyncio.gather(
+            sub_manager.handle_subscriptions(),
+            handle_subscription_stream(),
+        )
+
+        assert len(sub_manager.subscriptions) == 0
+        assert sub_manager.total_handler_calls == 1
+        assert handler_test.passed
+        assert stream_passed
+
+        # cleanup
+        sub_manager.total_handler_calls = 0
+
+    @pytest.mark.asyncio
+    async def test_handle_subscriptions_breaks_on_unsubscribe(
+        self,
+        async_w3: AsyncWeb3,
+    ) -> None:
+        async def unsubscribe_subs(
+            subs: List[Union[NewHeadsSubscription, LogsSubscription]]
+        ) -> None:
+            for sub in subs:
+                await sub.unsubscribe()
+
+        sub_manager = async_w3.subscription_manager
+        sub1 = NewHeadsSubscription(label="foo", handler=idle_handler)
+        sub2 = LogsSubscription(label="bar", handler=idle_handler)
+        await sub_manager.subscribe([sub1, sub2])
+        assert sub_manager.subscriptions == [sub1, sub2]
+
+        asyncio.create_task(unsubscribe_subs([sub1, sub2]))
+        # With no subscriptions in the queue, ``handle_subscriptions`` should hang
+        # indefinitely. Test that when the last subscription is unsubscribed from,
+        # the method breaks out of the loop. This is done via a raised
+        # ``SubscriptionProcessingFinished`` within the ``TaskReliantQueue``.
+        await sub_manager.handle_subscriptions()
+        assert_no_subscriptions_left(sub_manager._subscription_container)

web3/exceptions.py

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