|
3 | 3 | from dataclasses import (
|
4 | 4 | dataclass,
|
5 | 5 | )
|
| 6 | +import time |
6 | 7 | from typing import (
|
7 | 8 | TYPE_CHECKING,
|
8 | 9 | Any,
|
|
38 | 39 | from web3.middleware import (
|
39 | 40 | ExtraDataToPOAMiddleware,
|
40 | 41 | )
|
| 42 | +from web3.providers.persistent.request_processor import ( |
| 43 | + TaskReliantQueue, |
| 44 | +) |
41 | 45 | from web3.types import (
|
42 | 46 | BlockData,
|
43 | 47 | FormattedEthSubscriptionResponse,
|
@@ -914,3 +918,59 @@ async def test_run_forever_starts_with_0_subs_and_runs_until_task_cancelled(
|
914 | 918 |
|
915 | 919 | # cleanup
|
916 | 920 | await clean_up_task(run_forever_task)
|
| 921 | + |
| 922 | + @pytest.mark.asyncio |
| 923 | + async def test_high_throughput_subscription_task_based( |
| 924 | + self, async_w3: AsyncWeb3 |
| 925 | + ) -> None: |
| 926 | + async_w3.provider._request_processor._handler_subscription_queue = ( |
| 927 | + TaskReliantQueue(maxsize=5_000) |
| 928 | + ) |
| 929 | + sub_manager = async_w3.subscription_manager |
| 930 | + sub_manager.task_based = True # turn on task-based processing |
| 931 | + |
| 932 | + class Counter: |
| 933 | + val: int = 0 |
| 934 | + |
| 935 | + counter = Counter() |
| 936 | + |
| 937 | + async def high_throughput_handler( |
| 938 | + handler_context: Any, |
| 939 | + ) -> None: |
| 940 | + handler_context.counter.val += 1 |
| 941 | + if handler_context.counter.val == 5_000: |
| 942 | + await handler_context.subscription.unsubscribe() |
| 943 | + # if we awaited all 5_000 messages, we would sleep at least 5 seconds |
| 944 | + await asyncio.sleep(5 // 5_000) |
| 945 | + |
| 946 | + # build a meaningless subscription since we are fabricating the messages |
| 947 | + sub_id = await async_w3.eth.subscribe( |
| 948 | + "syncing", |
| 949 | + handler=high_throughput_handler, |
| 950 | + handler_context={"counter": counter}, |
| 951 | + ) |
| 952 | + async_w3.provider._request_processor.cache_request_information( |
| 953 | + request_id=sub_id, |
| 954 | + method=RPCEndpoint("eth_subscribe"), |
| 955 | + params=[], |
| 956 | + response_formatters=((), (), ()), # type: ignore |
| 957 | + ) |
| 958 | + |
| 959 | + # put 5_000 messages in the queue |
| 960 | + for _ in range(5_000): |
| 961 | + async_w3.provider._request_processor._handler_subscription_queue.put_nowait( |
| 962 | + { |
| 963 | + "jsonrpc": "2.0", |
| 964 | + "method": "eth_subscription", |
| 965 | + "params": {"subscription": HexBytes(sub_id), "result": False}, |
| 966 | + } |
| 967 | + ) |
| 968 | + |
| 969 | + start = time.time() |
| 970 | + await sub_manager.handle_subscriptions() |
| 971 | + stop = time.time() |
| 972 | + |
| 973 | + assert counter.val == 5_000 |
| 974 | + |
| 975 | + assert sub_manager.total_handler_calls == 5_000 |
| 976 | + assert stop - start < 3, "subscription handling took too long!" |
0 commit comments