web3 7.0.0b2__py3-none-any.whl → 7.7.0__py3-none-any.whl

web3/providers/persistent/async_ipc.py

@@ -1,9 +1,6 @@
 import asyncio
 import errno
 import json
-from json import (
-    JSONDecodeError,
-)
 import logging
 from pathlib import (
     Path,
@@ -16,10 +13,6 @@ from typing import (
     Union,
 )
 
-from eth_utils import (
-    to_text,
-)
-
 from web3.types import (
     RPCEndpoint,
     RPCResponse,
@@ -28,11 +21,11 @@ from web3.types import (
 from . import (
     PersistentConnectionProvider,
 )
-from ..._utils.caching import (
-    async_handle_request_caching,
-)
 from ...exceptions import (
+    PersistentConnectionClosedOK,
     ProviderConnectionError,
+    ReadBufferLimitReached,
+    Web3TypeError,
 )
 from ..ipc import (
     get_default_ipc_path,
@@ -40,7 +33,7 @@ from ..ipc import (
 
 
 async def async_get_ipc_socket(
-    ipc_path: str,
+    ipc_path: str, read_buffer_limit: int
 ) -> Tuple[asyncio.StreamReader, asyncio.StreamWriter]:
     if sys.platform == "win32":
         # On Windows named pipe is used. Simulate socket with it.
@@ -50,7 +43,7 @@ async def async_get_ipc_socket(
 
         return NamedPipe(ipc_path)
     else:
-        return await asyncio.open_unix_connection(ipc_path)
+        return await asyncio.open_unix_connection(ipc_path, limit=read_buffer_limit)
 
 
 class AsyncIPCProvider(PersistentConnectionProvider):
@@ -58,23 +51,24 @@ class AsyncIPCProvider(PersistentConnectionProvider):
 
     _reader: Optional[asyncio.StreamReader] = None
     _writer: Optional[asyncio.StreamWriter] = None
+    _decoder: json.JSONDecoder = json.JSONDecoder()
 
     def __init__(
         self,
         ipc_path: Optional[Union[str, Path]] = None,
-        max_connection_retries: int = 5,
+        read_buffer_limit: int = 20 * 1024 * 1024,  # 20 MB
         # `PersistentConnectionProvider` kwargs can be passed through
         **kwargs: Any,
     ) -> None:
+        # initialize the ipc_path before calling the super constructor
         if ipc_path is None:
             self.ipc_path = get_default_ipc_path()
         elif isinstance(ipc_path, str) or isinstance(ipc_path, Path):
             self.ipc_path = str(Path(ipc_path).expanduser().resolve())
         else:
-            raise TypeError("ipc_path must be of type string or pathlib.Path")
-
-        self._max_connection_retries = max_connection_retries
+            raise Web3TypeError("ipc_path must be of type string or pathlib.Path")
         super().__init__(**kwargs)
+        self.read_buffer_limit = read_buffer_limit
 
     def __str__(self) -> str:
         return f"<{self.__class__.__name__} {self.ipc_path}>"
@@ -84,81 +78,46 @@ class AsyncIPCProvider(PersistentConnectionProvider):
             return False
 
         try:
-            request_data = self.encode_rpc_request(
-                RPCEndpoint("web3_clientVersions"), []
-            )
-            self._writer.write(request_data)
-            current_request_id = json.loads(request_data)["id"]
-            await self._get_response_for_request_id(current_request_id, timeout=2)
+            await self.make_request(RPCEndpoint("web3_clientVersion"), [])
             return True
-        except (OSError, BrokenPipeError, ProviderConnectionError) as e:
+        except (OSError, ProviderConnectionError) as e:
             if show_traceback:
                 raise ProviderConnectionError(
                     f"Problem connecting to provider with error: {type(e)}: {e}"
                 )
             return False
 
-    async def connect(self) -> None:
-        _connection_attempts = 0
-        _backoff_rate_change = 1.75
-        _backoff_time = 1.75
-
-        while _connection_attempts != self._max_connection_retries:
-            try:
-                _connection_attempts += 1
-                self._reader, self._writer = await async_get_ipc_socket(self.ipc_path)
-                self._message_listener_task = asyncio.create_task(
-                    self._message_listener()
-                )
-                break
-            except OSError as e:
-                if _connection_attempts == self._max_connection_retries:
-                    raise ProviderConnectionError(
-                        f"Could not connect to endpoint: {self.endpoint_uri}. "
-                        f"Retries exceeded max of {self._max_connection_retries}."
-                    ) from e
-                self.logger.info(
-                    f"Could not connect to endpoint: {self.endpoint_uri}. Retrying in "
-                    f"{round(_backoff_time, 1)} seconds.",
-                    exc_info=True,
-                )
-                await asyncio.sleep(_backoff_time)
-                _backoff_time *= _backoff_rate_change
-
-    async def disconnect(self) -> None:
-        if self._writer and not self._writer.is_closing():
-            self._writer.close()
-            await self._writer.wait_closed()
-            self._writer = None
-            self.logger.debug(
-                f'Successfully disconnected from endpoint: "{self.endpoint_uri}'
-            )
-
-        try:
-            self._message_listener_task.cancel()
-            await self._message_listener_task
-            self._reader = None
-        except (asyncio.CancelledError, StopAsyncIteration):
-            pass
-
-        self._request_processor.clear_caches()
-
-    async def _reset_socket(self) -> None:
-        self._writer.close()
-        await self._writer.wait_closed()
-        self._reader, self._writer = await async_get_ipc_socket(self.ipc_path)
-
-    @async_handle_request_caching
-    async def make_request(self, method: RPCEndpoint, params: Any) -> RPCResponse:
-        request_data = self.encode_rpc_request(method, params)
-
+    async def socket_send(self, request_data: bytes) -> None:
         if self._writer is None:
             raise ProviderConnectionError(
                 "Connection to ipc socket has not been initiated for the provider."
             )
 
+        return await asyncio.wait_for(
+            self._socket_send(request_data), timeout=self.request_timeout
+        )
+
+    async def socket_recv(self) -> RPCResponse:
         try:
-            self._writer.write(request_data)
+            data = await self._reader.readline()
+        except ValueError as e:
+            if all(kw in str(e) for kw in ("limit", "chunk")):
+                raise ReadBufferLimitReached(
+                    f"Read buffer limit of `{self.read_buffer_limit}` bytes was "
+                    "reached. Consider increasing the ``read_buffer_limit`` on the "
+                    "AsyncIPCProvider."
+                ) from e
+            raise
+
+        if not data:
+            raise PersistentConnectionClosedOK("Socket reader received end of stream.")
+        return self.decode_rpc_response(data)
+
+    # -- private methods -- #
+
+    async def _socket_send(self, request_data: bytes) -> None:
+        try:
+            self._writer.write(request_data + b"\n")
             await self._writer.drain()
         except OSError as e:
             # Broken pipe
@@ -168,48 +127,29 @@ class AsyncIPCProvider(PersistentConnectionProvider):
                 self._writer.write(request_data)
                 await self._writer.drain()
 
-        current_request_id = json.loads(request_data)["id"]
-        response = await self._get_response_for_request_id(current_request_id)
-
-        return response
+    async def _reset_socket(self) -> None:
+        self._writer.close()
+        await self._writer.wait_closed()
+        self._reader, self._writer = await async_get_ipc_socket(
+            self.ipc_path, self.read_buffer_limit
+        )
 
-    async def _message_listener(self) -> None:
-        self.logger.info(
-            "IPC socket listener background task started. Storing all messages in "
-            "appropriate request processor queues / caches to be processed."
+    async def _provider_specific_connect(self) -> None:
+        self._reader, self._writer = await async_get_ipc_socket(
+            self.ipc_path, self.read_buffer_limit
         )
-        raw_message = ""
-        decoder = json.JSONDecoder()
-
-        while True:
-            # the use of sleep(0) seems to be the most efficient way to yield control
-            # back to the event loop to share the loop with other tasks.
-            await asyncio.sleep(0)
-
-            try:
-                raw_message += to_text(await self._reader.read(4096)).lstrip()
-
-                while raw_message:
-                    try:
-                        response, pos = decoder.raw_decode(raw_message)
-                    except JSONDecodeError:
-                        break
-
-                    is_subscription = response.get("method") == "eth_subscription"
-                    await self._request_processor.cache_raw_response(
-                        response, subscription=is_subscription
-                    )
-                    raw_message = raw_message[pos:].lstrip()
-            except Exception as e:
-                if not self.silence_listener_task_exceptions:
-                    loop = asyncio.get_event_loop()
-                    for task in asyncio.all_tasks(loop=loop):
-                        task.cancel()
-                    raise e
-
-                self.logger.error(
-                    "Exception caught in listener, error logging and keeping listener "
-                    f"background task alive.\n    error={e}"
-                )
-                # if only error logging, reset the ``raw_message`` buffer and continue
-                raw_message = ""
+
+    async def _provider_specific_disconnect(self) -> None:
+        # this should remain idempotent
+        if self._writer and not self._writer.is_closing():
+            self._writer.close()
+            await self._writer.wait_closed()
+            self._writer = None
+        if self._reader:
+            self._reader = None
+
+    async def _provider_specific_socket_reader(self) -> RPCResponse:
+        return await self.socket_recv()
+
+    def _error_log_listener_task_exception(self, e: Exception) -> None:
+        super()._error_log_listener_task_exception(e)

web3/providers/persistent/persistent.py

@@ -1,17 +1,44 @@
 from abc import (
     ABC,
+    abstractmethod,
 )
 import asyncio
 import logging
+import signal
 from typing import (
+    TYPE_CHECKING,
+    Any,
+    Callable,
+    Coroutine,
+    List,
     Optional,
+    Tuple,
+    Union,
+    cast,
 )
 
+from websockets import (
+    ConnectionClosed,
+    WebSocketException,
+)
+
+from web3._utils.batching import (
+    BATCH_REQUEST_ID,
+    sort_batch_response_by_response_ids,
+)
 from web3._utils.caching import (
     generate_cache_key,
 )
+from web3._utils.caching.caching_utils import (
+    async_handle_recv_caching,
+    async_handle_send_caching,
+)
 from web3.exceptions import (
+    PersistentConnectionClosedOK,
+    ProviderConnectionError,
+    TaskNotRunning,
     TimeExhausted,
+    Web3AttributeError,
 )
 from web3.providers.async_base import (
     AsyncJSONBaseProvider,
@@ -20,47 +47,327 @@ from web3.providers.persistent.request_processor import (
     RequestProcessor,
 )
 from web3.types import (
+    RPCEndpoint,
     RPCId,
+    RPCRequest,
     RPCResponse,
 )
 
+if TYPE_CHECKING:
+    from web3 import AsyncWeb3  # noqa: F401
+    from web3.middleware.base import MiddlewareOnion  # noqa: F401
+
+
 DEFAULT_PERSISTENT_CONNECTION_TIMEOUT = 30.0
 
 
 class PersistentConnectionProvider(AsyncJSONBaseProvider, ABC):
     logger = logging.getLogger("web3.providers.PersistentConnectionProvider")
     has_persistent_connection = True
-    endpoint_uri: Optional[str] = None
 
-    _request_processor: RequestProcessor
-    _message_listener_task: Optional["asyncio.Task[None]"] = None
-    _listen_event: asyncio.Event = asyncio.Event()
+    _send_func_cache: Tuple[int, Callable[..., Coroutine[Any, Any, RPCRequest]]] = (
+        None,
+        None,
+    )
+    _recv_func_cache: Tuple[int, Callable[..., Coroutine[Any, Any, RPCResponse]]] = (
+        None,
+        None,
+    )
 
     def __init__(
         self,
         request_timeout: float = DEFAULT_PERSISTENT_CONNECTION_TIMEOUT,
         subscription_response_queue_size: int = 500,
         silence_listener_task_exceptions: bool = False,
+        max_connection_retries: int = 5,
+        **kwargs: Any,
     ) -> None:
-        super().__init__()
+        super().__init__(**kwargs)
         self._request_processor = RequestProcessor(
             self,
             subscription_response_queue_size=subscription_response_queue_size,
         )
+        self._message_listener_task: Optional["asyncio.Task[None]"] = None
+        self._batch_request_counter: Optional[int] = None
+        self._listen_event: asyncio.Event = asyncio.Event()
+        self._max_connection_retries = max_connection_retries
+
         self.request_timeout = request_timeout
         self.silence_listener_task_exceptions = silence_listener_task_exceptions
 
+    async def send_func(
+        self, async_w3: "AsyncWeb3", middleware_onion: "MiddlewareOnion"
+    ) -> Callable[..., Coroutine[Any, Any, RPCRequest]]:
+        """
+        Cache the middleware chain for `send`.
+        """
+        middleware = middleware_onion.as_tuple_of_middleware()
+        cache_key = hash(tuple(id(mw) for mw in middleware))
+
+        if cache_key != self._send_func_cache[0]:
+
+            async def send_function(method: RPCEndpoint, params: Any) -> RPCRequest:
+                for mw in middleware:
+                    initialized = mw(async_w3)
+                    method, params = await initialized.async_request_processor(
+                        method, params
+                    )
+
+                return await self.send_request(method, params)
+
+            self._send_func_cache = (cache_key, send_function)
+
+        return self._send_func_cache[1]
+
+    async def recv_func(
+        self, async_w3: "AsyncWeb3", middleware_onion: "MiddlewareOnion"
+    ) -> Any:
+        """
+        Cache and compose the middleware stack for `recv`.
+        """
+        middleware = middleware_onion.as_tuple_of_middleware()
+        cache_key = hash(tuple(id(mw) for mw in middleware))
+
+        if cache_key != self._recv_func_cache[0]:
+
+            async def recv_function(rpc_request: RPCRequest) -> RPCResponse:
+                # first, retrieve the response
+                response = await self.recv_for_request(rpc_request)
+                method = rpc_request["method"]
+                for mw in reversed(middleware):
+                    initialized = mw(async_w3)
+                    response = await initialized.async_response_processor(
+                        method, response
+                    )
+                return response
+
+            self._recv_func_cache = (cache_key, recv_function)
+
+        return self._recv_func_cache[1]
+
+    def get_endpoint_uri_or_ipc_path(self) -> str:
+        if hasattr(self, "endpoint_uri"):
+            return str(self.endpoint_uri)
+        elif hasattr(self, "ipc_path"):
+            return str(self.ipc_path)
+        else:
+            raise Web3AttributeError(
+                "`PersistentConnectionProvider` must have either `endpoint_uri` or "
+                "`ipc_path` attribute."
+            )
+
     async def connect(self) -> None:
-        raise NotImplementedError("Must be implemented by subclasses")
+        _connection_attempts = 0
+        _backoff_rate_change = 1.75
+        _backoff_time = 1.75
+
+        while _connection_attempts != self._max_connection_retries:
+            try:
+                _connection_attempts += 1
+                self.logger.info(
+                    f"Connecting to: {self.get_endpoint_uri_or_ipc_path()}"
+                )
+                await self._provider_specific_connect()
+                self._message_listener_task = asyncio.create_task(
+                    self._message_listener()
+                )
+                self._message_listener_task.add_done_callback(
+                    self._message_listener_callback
+                )
+                self.logger.info(
+                    f"Successfully connected to: {self.get_endpoint_uri_or_ipc_path()}"
+                )
+                break
+            except (WebSocketException, OSError) as e:
+                if _connection_attempts == self._max_connection_retries:
+                    raise ProviderConnectionError(
+                        f"Could not connect to: {self.get_endpoint_uri_or_ipc_path()}. "
+                        f"Retries exceeded max of {self._max_connection_retries}."
+                    ) from e
+                self.logger.info(
+                    f"Could not connect to: {self.get_endpoint_uri_or_ipc_path()}. "
+                    f"Retrying in {round(_backoff_time, 1)} seconds.",
+                    exc_info=True,
+                )
+                await asyncio.sleep(_backoff_time)
+                _backoff_time *= _backoff_rate_change
 
     async def disconnect(self) -> None:
+        # this should remain idempotent
+        try:
+            if self._message_listener_task:
+                self._message_listener_task.cancel()
+                await self._message_listener_task
+        except (asyncio.CancelledError, StopAsyncIteration, ConnectionClosed):
+            pass
+        finally:
+            self._message_listener_task = None
+            self.logger.info("Message listener background task successfully shut down.")
+
+        await self._provider_specific_disconnect()
+        self._request_processor.clear_caches()
+        self.logger.info(
+            f"Successfully disconnected from: {self.get_endpoint_uri_or_ipc_path()}"
+        )
+
+    @async_handle_send_caching
+    async def send_request(self, method: RPCEndpoint, params: Any) -> RPCRequest:
+        request_dict = self.form_request(method, params)
+        await self.socket_send(self.encode_rpc_dict(request_dict))
+        return request_dict
+
+    @async_handle_recv_caching
+    async def recv_for_request(self, rpc_request: RPCRequest) -> RPCResponse:
+        return await self._get_response_for_request_id(rpc_request["id"])
+
+    async def make_request(
+        self,
+        method: RPCEndpoint,
+        params: Any,
+    ) -> RPCResponse:
+        rpc_request = await self.send_request(method, params)
+        return await self.recv_for_request(rpc_request)
+
+    async def make_batch_request(
+        self, requests: List[Tuple[RPCEndpoint, Any]]
+    ) -> List[RPCResponse]:
+        request_data = self.encode_batch_rpc_request(requests)
+        await self.socket_send(request_data)
+
+        response = cast(
+            List[RPCResponse], await self._get_response_for_request_id(BATCH_REQUEST_ID)
+        )
+        return response
+
+    # -- abstract methods -- #
+
+    @abstractmethod
+    async def socket_send(self, request_data: bytes) -> None:
+        """
+        Send an encoded RPC request to the provider over the persistent connection.
+        """
         raise NotImplementedError("Must be implemented by subclasses")
 
-    async def _message_listener(self) -> None:
+    @abstractmethod
+    async def socket_recv(self) -> RPCResponse:
+        """
+        Receive, decode, and return an RPC response from the provider over the
+        persistent connection.
+        """
+        raise NotImplementedError("Must be implemented by subclasses")
+
+    # -- private methods -- #
+
+    async def _provider_specific_connect(self) -> None:
         raise NotImplementedError("Must be implemented by subclasses")
 
+    async def _provider_specific_disconnect(self) -> None:
+        # this method should be idempotent
+        raise NotImplementedError("Must be implemented by subclasses")
+
+    async def _provider_specific_socket_reader(self) -> RPCResponse:
+        raise NotImplementedError("Must be implemented by subclasses")
+
+    def _set_signal_handlers(self) -> None:
+        def extended_handler(sig: int, frame: Any, existing_handler: Any) -> None:
+            loop = asyncio.get_event_loop()
+
+            # invoke the existing handler, if callable
+            if callable(existing_handler):
+                existing_handler(sig, frame)
+            loop.create_task(self.disconnect())
+
+        existing_sigint_handler = signal.getsignal(signal.SIGINT)
+        existing_sigterm_handler = signal.getsignal(signal.SIGTERM)
+
+        # extend the existing signal handlers to include the disconnect method
+        signal.signal(
+            signal.SIGINT,
+            lambda sig, frame: extended_handler(sig, frame, existing_sigint_handler),
+        )
+        signal.signal(
+            signal.SIGTERM,
+            lambda sig, frame: extended_handler(sig, frame, existing_sigterm_handler),
+        )
+
+    def _message_listener_callback(
+        self, message_listener_task: "asyncio.Task[None]"
+    ) -> None:
+        # Puts a `TaskNotRunning` in appropriate queues to signal the end of the
+        # listener task to any listeners relying on the queues.
+        self._request_processor._subscription_response_queue.put_nowait(
+            TaskNotRunning(message_listener_task)
+        )
+        self._request_processor._handler_subscription_queue.put_nowait(
+            TaskNotRunning(message_listener_task)
+        )
+
+    async def _message_listener(self) -> None:
+        self.logger.info(
+            f"{self.__class__.__qualname__} listener background task started. Storing "
+            "all messages in appropriate request processor queues / caches to be "
+            "processed."
+        )
+        while True:
+            # the use of sleep(0) seems to be the most efficient way to yield control
+            # back to the event loop to share the loop with other tasks.
+            await asyncio.sleep(0)
+
+            try:
+                response = await self._provider_specific_socket_reader()
+
+                if isinstance(response, list):
+                    response = sort_batch_response_by_response_ids(response)
+
+                subscription = (
+                    response.get("method") == "eth_subscription"
+                    if not isinstance(response, list)
+                    else False
+                )
+                await self._request_processor.cache_raw_response(
+                    response, subscription=subscription
+                )
+            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
+                else:
+                    self._error_log_listener_task_exception(e)
+
+    def _error_log_listener_task_exception(self, e: Exception) -> None:
+        """
+        When silencing listener task exceptions, this method is used to log the
+        exception and keep the listener task alive. Override this method to fine-tune
+        error logging behavior for the implementation class.
+        """
+        self.logger.error(
+            "Exception caught in listener, error logging and keeping "
+            "listener background task alive."
+            f"\n    error={e.__class__.__name__}: {e}"
+        )
+
+    def _handle_listener_task_exceptions(self) -> None:
+        """
+        Should be called every time a `PersistentConnectionProvider` is polling for
+        messages in the main loop. If the message listener task has completed and an
+        exception was recorded, raise the exception in the main loop.
+        """
+        msg_listener_task = getattr(self, "_message_listener_task", None)
+        if (
+            msg_listener_task
+            and msg_listener_task.done()
+            and msg_listener_task.exception()
+        ):
+            raise msg_listener_task.exception()
+
     async def _get_response_for_request_id(
-        self, request_id: RPCId, timeout: Optional[float] = None
+        self, request_id: Union[RPCId, List[RPCId]], timeout: Optional[float] = None
     ) -> RPCResponse:
         if timeout is None:
             timeout = self.request_timeout
@@ -69,24 +376,24 @@ class PersistentConnectionProvider(AsyncJSONBaseProvider, ABC):
             request_cache_key = generate_cache_key(request_id)
 
             while True:
-                # sleep(0) here seems to be the most efficient way to yield control
-                # back to the event loop while waiting for the response to be in the
-                # queue.
-                await asyncio.sleep(0)
-
                 if request_cache_key in self._request_processor._request_response_cache:
                     self.logger.debug(
                         f"Popping response for id {request_id} from cache."
                     )
-                    popped_response = self._request_processor.pop_raw_response(
+                    popped_response = await self._request_processor.pop_raw_response(
                         cache_key=request_cache_key,
                     )
                     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:
             # Add the request timeout around the while loop that checks the request
-            # cache and tried to recv(). If the request is neither in the cache, nor
-            # received within the request_timeout, raise ``TimeExhausted``.
+            # cache. If the request is not in the cache within the request_timeout,
+            # raise ``TimeExhausted``.
             return await asyncio.wait_for(_match_response_id_to_request_id(), timeout)
         except asyncio.TimeoutError:
             raise TimeExhausted(