feat: Add a subscription manager to manage subscriptions

- Create a SubscriptionManager class to manage subscriptions.
- Implement subscribing to one or many subscriptions.
- Implement unsubscribing from one or many subscriptions.

Author: fselmo

web3/eth/async_eth.py

@@ -99,8 +99,12 @@
     _Hash32,
 )
 from web3.utils import (
+    EthSubscription,
     async_handle_offchain_lookup,
 )
+from web3.utils.subscriptions import (
+    EthSubscriptionHandler,
+)
 
 if TYPE_CHECKING:
     from web3 import AsyncWeb3  # noqa: F401
@@ -714,17 +718,22 @@ async def subscribe(
                 bool,  # newPendingTransactions, full_transactions
             ]
         ] = None,
+        handler: Optional[EthSubscriptionHandler] = None,
+        label: Optional[str] = None,
     ) -> HexStr:
         if not isinstance(self.w3.provider, PersistentConnectionProvider):
             raise MethodNotSupported(
                 "eth_subscribe is only supported with providers that support "
                 "persistent connections."
             )
 
-        if subscription_arg is None:
-            return await self._subscribe(subscription_type)
-
-        return await self._subscribe_with_args(subscription_type, subscription_arg)
+        params = (
+            (subscription_type, subscription_arg)
+            if subscription_arg is not None
+            else (subscription_type,)
+        )
+        sx = EthSubscription(subscription_params=params, handler=handler, label=label)
+        return await self.w3.subscription_manager.subscribe(sx)
 
     _unsubscribe: Method[Callable[[HexStr], Awaitable[bool]]] = Method(
         RPC.eth_unsubscribe,
@@ -738,7 +747,14 @@ async def unsubscribe(self, subscription_id: HexStr) -> bool:
                 "persistent connections."
             )
 
-        return await self._unsubscribe(subscription_id)
+        for sx in self.w3.subscription_manager.subscriptions:
+            if sx._id == subscription_id:
+                return await sx.unsubscribe()
+
+        raise Web3ValueError(
+            f"Cannot unsubscribe subscription with id `{subscription_id}`. "
+            "Subscription not found."
+        )
 
     # -- contract methods -- #
 

web3/main.py

@@ -148,6 +148,9 @@
 from web3.types import (
     Wei,
 )
+from web3.providers.persistent.subscription_manager import (
+    SubscriptionManager,
+)
 
 if TYPE_CHECKING:
     from web3._utils.batching import RequestBatcher  # noqa: F401
@@ -508,7 +511,20 @@ def ens(self, new_ens: Union[AsyncENS, "Empty"]) -> None:
             new_ens.w3 = self  # set self object reference for ``AsyncENS.w3``
         self._ens = new_ens
 
-    # -- persistent connection methods -- #
+    # -- persistent connection settings -- #
+
+    _subscription_manager: SubscriptionManager = None
+
+    @property
+    @persistent_connection_provider_method()
+    def subscription_manager(self) -> SubscriptionManager:
+        """
+        Access the subscription manager for the current PersistentConnectionProvider.
+        """
+        if not self._subscription_manager:
+            self._subscription_manager = SubscriptionManager(self)
+
+        return self._subscription_manager
 
     @property
     @persistent_connection_provider_method()

web3/manager.py

@@ -543,15 +543,18 @@ async def recv(self) -> RPCResponse:
     def _persistent_message_stream(self) -> "_AsyncPersistentMessageStream":
         return _AsyncPersistentMessageStream(self)
 
-    async def _get_next_message(self) -> RPCResponse:
+    async def _get_next_message(self) -> Optional[RPCResponse]:
         return await self._message_stream().__anext__()
 
-    async def _message_stream(self) -> AsyncGenerator[RPCResponse, None]:
+    async def _message_stream(
+        self,
+    ) -> AsyncGenerator[Optional[RPCResponse], None]:
         if not isinstance(self._provider, PersistentConnectionProvider):
             raise Web3TypeError(
                 "Only providers that maintain an open, persistent connection "
                 "can listen to streams."
             )
+        async_w3 = cast("AsyncWeb3", self.w3)
 
         if self._provider._message_listener_task is None:
             raise ProviderConnectionError(
@@ -563,13 +566,25 @@ async def _message_stream(self) -> AsyncGenerator[RPCResponse, None]:
                 response = await self._request_processor.pop_raw_response(
                     subscription=True
                 )
-                if (
-                    response is not None
-                    and response.get("params", {}).get("subscription")
-                    in self._request_processor.active_subscriptions
-                ):
-                    # if response is an active subscription response, process it
-                    yield await self._process_response(response)
+                if response is not None:
+                    sx_response = await self._process_response(response)
+
+                    # call the handler is there is one, else yield the response to any
+                    # listeners
+                    for sx in async_w3.subscription_manager.subscriptions:
+                        if (
+                            sx_response.get("subscription") == sx.id
+                            and sx._handler is not None
+                        ):
+                            await sx._handler(async_w3, sx, sx_response["result"])
+                    yield sx_response
+                else:
+                    # if response is not an active subscription response, log it
+                    self.logger.debug(
+                        "Received inactive subscription from socket:\n"
+                        f"    {self._provider.get_endpoint_uri_or_ipc_path()}, "
+                        f"    response: {response}"
+                    )
             except TaskNotRunning:
                 await asyncio.sleep(0)
                 self._provider._handle_listener_task_exceptions()
@@ -642,5 +657,5 @@ def __init__(self, manager: RequestManager, *args: Any, **kwargs: Any) -> None:
     def __aiter__(self) -> Self:
         return self
 
-    async def __anext__(self) -> RPCResponse:
+    async def __anext__(self) -> Optional[RPCResponse]:
         return await self.manager._get_next_message()

web3/providers/persistent/persistent.py

@@ -65,6 +65,7 @@ def __init__(
         subscription_response_queue_size: int = 500,
         silence_listener_task_exceptions: bool = False,
         max_connection_retries: int = 5,
+        label: Optional[str] = None,
         **kwargs: Any,
     ) -> None:
         super().__init__(**kwargs)
@@ -75,6 +76,7 @@ def __init__(
         self.request_timeout = request_timeout
         self.silence_listener_task_exceptions = silence_listener_task_exceptions
         self._max_connection_retries = max_connection_retries
+        self.label = label or self.get_endpoint_uri_or_ipc_path()
 
     def get_endpoint_uri_or_ipc_path(self) -> str:
         if hasattr(self, "endpoint_uri"):

web3/providers/persistent/subscription_manager.py

@@ -0,0 +1,103 @@
+from typing import (
+    TYPE_CHECKING,
+    List,
+    Sequence,
+    Union,
+    cast,
+    overload,
+)
+
+from eth_typing import (
+    HexStr,
+)
+from typing_extensions import (
+    TypeVar,
+)
+
+from web3.exceptions import (
+    Web3TypeError,
+    Web3ValueError,
+)
+from web3.utils import (
+    EthSubscription,
+)
+
+if TYPE_CHECKING:
+    from web3 import AsyncWeb3  # noqa: F401
+    from web3.providers.persistent import PersistentConnectionProvider  # noqa: F401
+
+
+T = TypeVar("T", bound="EthSubscription")
+
+
+class SubscriptionManager:
+    def __init__(self, w3: "AsyncWeb3") -> None:
+        self._w3 = w3
+        self._provider = cast("PersistentConnectionProvider", w3.provider)
+        self.subscriptions: List[EthSubscription] = []
+
+    @overload
+    async def subscribe(self, subscriptions: T) -> HexStr:
+        ...
+
+    @overload
+    async def subscribe(self, subscriptions: Sequence[T]) -> List[HexStr]:
+        ...
+
+    async def subscribe(
+        self, subscriptions: Union[T, Sequence[T]]
+    ) -> Union[HexStr, List[HexStr]]:
+        if isinstance(subscriptions, EthSubscription):
+            subscriptions._manager = self
+            sx_id = await self._w3.eth._subscribe(*subscriptions.subscription_params)
+            subscriptions._id = sx_id
+            self.subscriptions.append(subscriptions)
+            self._provider.logger.info(
+                f"Successfully subscribed to subscription:\n    "
+                f"label: {subscriptions.label}\n    id: {sx_id}"
+            )
+            return sx_id
+        elif isinstance(subscriptions, Sequence):
+            if len(subscriptions) == 0:
+                raise Web3ValueError("No subscriptions provided.")
+
+            sx_ids = []
+            for sx in subscriptions:
+                await self.subscribe(sx)
+                sx_ids.append(sx._id)
+
+            return sx_ids
+        else:
+            raise Web3TypeError(
+                "Expected a Subscription or a sequence of Subscriptions."
+            )
+
+    async def unsubscribe(self, subscription: EthSubscription) -> bool:
+        if subscription not in self.subscriptions:
+            raise Web3ValueError(
+                f"Subscription not found or is not being managed by the subscription "
+                f"manager.\n    label: {subscription.label}\n    id: {subscription._id}"
+            )
+
+        if await self._w3.eth._unsubscribe(subscription.id):
+            self.subscriptions.remove(subscription)
+            return True
+        return False
+
+    async def unsubscribe_all(self) -> None:
+        for sx in self.subscriptions:
+            await self.unsubscribe(sx)
+
+        self._provider.logger.info("Successfully unsubscribed from all subscriptions.")
+
+    async def _handle_subscriptions(self, run_forever: bool = False) -> None:
+        self._provider.logger.info("Subscription manager processing started.")
+        while True:
+            if not run_forever and len(self.subscriptions) == 0:
+                break
+            await self._w3.manager._get_next_message()
+
+        self._provider.logger.info("Subscription manager processing ended.")
+
+    async def handle_subscriptions(self, run_forever: bool = False) -> None:
+        await self._handle_subscriptions(run_forever=run_forever)

web3/utils/__init__.py

@@ -42,6 +42,9 @@
 from .exception_handling import (
     handle_offchain_lookup,
 )
+from .subscriptions import (
+    EthSubscription,
+)
 
 __all__ = [
     "abi_to_signature",
@@ -70,5 +73,6 @@
     "async_handle_offchain_lookup",
     "RequestCacheValidationThreshold",
     "SimpleCache",
+    "EthSubscription",
     "handle_offchain_lookup",
 ]