[Persistent Connection Providers] Improve request info management:

- Instead of guessing what the request `id` will be and having to bump
  this cache key for internal requests, split persistent Connection
  provider calls into send and recv. Use the build request after sending
  to get the `id` and store the request information for when the response
  is received.

- Add request caching support and request mocker support for these new
  methods. Use the `send_request` and `recv_for_request` methods on the
  `Module` class so we properly handle the request information for
  processing the response.

Author: fselmo

tests/core/providers/test_async_ipc_provider.py

@@ -364,5 +364,5 @@ async def test_async_ipc_provider_write_messages_end_with_new_line_delimiter(
 
         await w3.provider.make_request("method", [])
 
-        request_data = b'{"jsonrpc": "2.0", "method": "method", "params": [], "id": 0}'
+        request_data = b'{"id": 0, "jsonrpc": "2.0", "method": "method", "params": []}'
         w3.provider._writer.write.assert_called_with(request_data + b"\n")

web3/_utils/caching/caching_utils.py

@@ -64,10 +64,12 @@
     from web3.providers import (  # noqa: F401
         AsyncBaseProvider,
         BaseProvider,
+        PersistentConnectionProvider,
     )
     from web3.types import (  # noqa: F401
         AsyncMakeRequestFn,
         MakeRequestFn,
+        RPCRequest,
         RPCResponse,
     )
 
@@ -367,3 +369,65 @@ async def wrapper(
     # save a reference to the decorator on the wrapped function
     wrapper._decorator = async_handle_request_caching  # type: ignore
     return wrapper
+
+
+def async_handle_send_caching(
+    func: Callable[
+        [ASYNC_PROVIDER_TYPE, RPCEndpoint, Any],
+        Coroutine[Any, Any, "RPCRequest"],
+    ],
+) -> Callable[..., Coroutine[Any, Any, "RPCRequest"]]:
+    async def wrapper(
+        provider: ASYNC_PROVIDER_TYPE, method: RPCEndpoint, params: Any
+    ) -> "RPCRequest":
+        if is_cacheable_request(provider, method, params):
+            request_cache = provider._request_cache
+            cache_key = generate_cache_key(
+                f"{threading.get_ident()}:{(method, params)}"
+            )
+            cached_response = request_cache.get_cache_entry(cache_key)
+            if cached_response is not None:
+                # The request data isn't used, this just prevents a cached request from
+                # being sent - return an empty request object
+                return {"id": -1, "method": RPCEndpoint(""), "params": []}
+        return await func(provider, method, params)
+
+    # save a reference to the decorator on the wrapped function
+    wrapper._decorator = async_handle_send_caching  # type: ignore
+    return wrapper
+
+
+def async_handle_recv_caching(
+    func: Callable[
+        ["PersistentConnectionProvider", "RPCRequest"],
+        Coroutine[Any, Any, "RPCResponse"],
+    ]
+) -> Callable[..., Coroutine[Any, Any, "RPCResponse"]]:
+    async def wrapper(
+        provider: "PersistentConnectionProvider",
+        rpc_request: "RPCRequest",
+    ) -> "RPCResponse":
+        method = rpc_request["method"]
+        params = rpc_request["params"]
+        if is_cacheable_request(provider, method, params):
+            request_cache = provider._request_cache
+            cache_key = generate_cache_key(
+                f"{threading.get_ident()}:{(method, params)}"
+            )
+            cache_result = request_cache.get_cache_entry(cache_key)
+            if cache_result is not None:
+                return cache_result
+            else:
+                response = await func(provider, rpc_request)
+                if await _async_should_cache_response(
+                    provider, method, params, response
+                ):
+                    async with provider._request_cache_lock:
+                        request_cache.cache(cache_key, response)
+                return response
+        else:
+            return await func(provider, rpc_request)
+
+    # save a reference to the decorator on the wrapped function
+    wrapper._decorator = async_handle_recv_caching  # type: ignore
+    return wrapper

web3/_utils/caching/request_caching_validation.py

@@ -19,6 +19,7 @@
     from web3.providers import (  # noqa: F401
         AsyncBaseProvider,
         BaseProvider,
+        PersistentConnectionProvider,
     )
 
 UNCACHEABLE_BLOCK_IDS = {"finalized", "safe", "latest", "pending"}

web3/_utils/module_testing/utils.py

@@ -10,10 +10,14 @@
     cast,
 )
 
-from toolz import (
+from eth_utils.toolz import (
     merge,
 )
 
+from web3.providers.persistent import (
+    PersistentConnectionProvider,
+)
+
 if TYPE_CHECKING:
     from web3 import (  # noqa: F401
         AsyncWeb3,
@@ -26,6 +30,7 @@
         AsyncMakeRequestFn,
         MakeRequestFn,
         RPCEndpoint,
+        RPCRequest,
         RPCResponse,
     )
 
@@ -101,9 +106,21 @@ def __init__(
         self.mock_results = mock_results or {}
         self.mock_errors = mock_errors or {}
         self.mock_responses = mock_responses or {}
-        self._make_request: Union[
-            "AsyncMakeRequestFn", "MakeRequestFn"
-        ] = w3.provider.make_request
+        if isinstance(w3.provider, PersistentConnectionProvider):
+            self._send_request = w3.provider.send_request
+            self._recv_for_request = w3.provider.recv_for_request
+        else:
+            self._make_request: Union[
+                "AsyncMakeRequestFn", "MakeRequestFn"
+            ] = w3.provider.make_request
+
+    def _build_request_id(self) -> int:
+        request_id = (
+            next(copy.deepcopy(self.w3.provider.request_counter))
+            if hasattr(self.w3.provider, "request_counter")
+            else 1
+        )
+        return request_id
 
     def __enter__(self) -> "Self":
         # mypy error: Cannot assign to a method
@@ -131,11 +148,7 @@ def _mock_request_handler(
         ):
             return self._make_request(method, params)
 
-        request_id = (
-            next(copy.deepcopy(self.w3.provider.request_counter))
-            if hasattr(self.w3.provider, "request_counter")
-            else 1
-        )
+        request_id = self._build_request_id()
         response_dict = {"jsonrpc": "2.0", "id": request_id}
 
         if method in self.mock_responses:
@@ -176,35 +189,34 @@ def _mock_request_handler(
     # -- async -- #
 
     async def __aenter__(self) -> "Self":
-        # mypy error: Cannot assign to a method
-        self.w3.provider.make_request = self._async_mock_request_handler  # type: ignore[method-assign]  # noqa: E501
-        # reset request func cache to re-build request_func with mocked make_request
-        self.w3.provider._request_func_cache = (None, None)
+        if not isinstance(self.w3.provider, PersistentConnectionProvider):
+            # mypy error: Cannot assign to a method
+            self.w3.provider.make_request = self._async_mock_request_handler  # type: ignore[method-assign]  # noqa: E501
+            # reset request func cache to re-build request_func w/ mocked make_request
+            self.w3.provider._request_func_cache = (None, None)
+        else:
+            self.w3.provider.send_request = self._async_mock_send_handler  # type: ignore[method-assign]  # noqa: E501
+            self.w3.provider.recv_for_request = self._async_mock_recv_handler  # type: ignore[method-assign]  # noqa: E501
+            self.w3.provider._send_func_cache = (None, None)
+            self.w3.provider._recv_func_cache = (None, None)
         return self
 
     async def __aexit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
-        # mypy error: Cannot assign to a method
-        self.w3.provider.make_request = self._make_request  # type: ignore[assignment]
-        # reset request func cache to re-build request_func with original make_request
-        self.w3.provider._request_func_cache = (None, None)
+        if not isinstance(self.w3.provider, PersistentConnectionProvider):
+            # mypy error: Cannot assign to a method
+            self.w3.provider.make_request = self._make_request  # type: ignore[assignment]  # noqa: E501
+            # reset request func cache to re-build request_func w/ original make_request
+            self.w3.provider._request_func_cache = (None, None)
+        else:
+            self.w3.provider.send_request = self._send_request  # type: ignore[method-assign]  # noqa: E501
+            self.w3.provider.recv_for_request = self._recv_for_request  # type: ignore[method-assign]  # noqa: E501
+            self.w3.provider._send_func_cache = (None, None)
+            self.w3.provider._recv_func_cache = (None, None)
 
-    async def _async_mock_request_handler(
-        self, method: "RPCEndpoint", params: Any
+    async def _async_build_mock_result(
+        self, method: "RPCEndpoint", params: Any, request_id: int = None
     ) -> "RPCResponse":
-        self.w3 = cast("AsyncWeb3", self.w3)
-        self._make_request = cast("AsyncMakeRequestFn", self._make_request)
-
-        if all(
-            method not in mock_dict
-            for mock_dict in (self.mock_errors, self.mock_results, self.mock_responses)
-        ):
-            return await self._make_request(method, params)
-
-        request_id = (
-            next(copy.deepcopy(self.w3.provider.request_counter))
-            if hasattr(self.w3.provider, "request_counter")
-            else 1
-        )
+        request_id = request_id if request_id else self._build_request_id()
         response_dict = {"jsonrpc": "2.0", "id": request_id}
 
         if method in self.mock_responses:
@@ -244,6 +256,19 @@ async def _async_mock_request_handler(
         else:
             raise Exception("Invariant: unreachable code path")
 
+        return mocked_result
+
+    async def _async_mock_request_handler(
+        self, method: "RPCEndpoint", params: Any
+    ) -> "RPCResponse":
+        self.w3 = cast("AsyncWeb3", self.w3)
+        self._make_request = cast("AsyncMakeRequestFn", self._make_request)
+        if all(
+            method not in mock_dict
+            for mock_dict in (self.mock_errors, self.mock_results, self.mock_responses)
+        ):
+            return await self._make_request(method, params)
+        mocked_result = await self._async_build_mock_result(method, params)
         decorator = getattr(self._make_request, "_decorator", None)
         if decorator is not None:
             # If the original make_request was decorated, we need to re-apply
@@ -259,6 +284,47 @@ async def _coro(
         else:
             return mocked_result
 
+    async def _async_mock_send_handler(
+        self, method: "RPCEndpoint", params: Any
+    ) -> "RPCRequest":
+        if all(
+            method not in mock_dict
+            for mock_dict in (self.mock_errors, self.mock_results, self.mock_responses)
+        ):
+            return await self._send_request(method, params)
+        else:
+            request_id = self._build_request_id()
+            return {"id": request_id, "method": method, "params": params}
+
+    async def _async_mock_recv_handler(
+        self, rpc_request: "RPCRequest"
+    ) -> "RPCResponse":
+        self.w3 = cast("AsyncWeb3", self.w3)
+        method = rpc_request["method"]
+        request_id = rpc_request["id"]
+        if all(
+            method not in mock_dict
+            for mock_dict in (self.mock_errors, self.mock_results, self.mock_responses)
+        ):
+            return await self._recv_for_request(request_id)
+        mocked_result = await self._async_build_mock_result(
+            method, rpc_request["params"], request_id=int(request_id)
+        )
+        decorator = getattr(self._recv_for_request, "_decorator", None)
+        if decorator is not None:
+            # If the original recv_for_request was decorated, we need to re-apply
+            # the decorator to the mocked recv_for_request. This is necessary for
+            # the request caching decorator to work properly.
+
+            async def _coro(
+                _provider: Any, _rpc_request: "RPCRequest"
+            ) -> "RPCResponse":
+                return mocked_result
+
+            return await decorator(_coro)(self.w3.provider, rpc_request)
+        else:
+            return mocked_result
+
     @staticmethod
     def _create_error_object(error: Dict[str, Any]) -> Dict[str, Any]:
         code = error.get("code", -32000)

web3/manager.py

@@ -6,6 +6,7 @@
     AsyncGenerator,
     Callable,
     Coroutine,
+    Dict,
     List,
     Optional,
     Sequence,
@@ -71,6 +72,7 @@
 from web3.types import (
     FormattedEthSubscriptionResponse,
     RPCEndpoint,
+    RPCRequest,
     RPCResponse,
 )
 
@@ -502,31 +504,62 @@ def _format_batched_response(
     # -- persistent connection -- #
 
     async def socket_request(
-        self, method: RPCEndpoint, params: Any
-    ) -> Union[RPCResponse, FormattedEthSubscriptionResponse]:
+        self,
+        method: RPCEndpoint,
+        params: Any,
+        response_formatters: Optional[
+            Tuple[Dict[str, Callable[..., Any]], Callable[..., Any], Callable[..., Any]]
+        ] = None,
+    ) -> RPCResponse:
         provider = cast(PersistentConnectionProvider, self._provider)
-        request_func = await provider.request_func(
-            cast("AsyncWeb3", self.w3), cast("MiddlewareOnion", self.middleware_onion)
-        )
         self.logger.debug(
             "Making request to open socket connection and waiting for response: "
-            f"{provider.get_endpoint_uri_or_ipc_path()}, method: {method}"
+            f"{provider.get_endpoint_uri_or_ipc_path()},\n    method: {method},\n"
+            f"    params: {params}"
         )
-        response = await request_func(method, params)
-        return await self._process_response(response)
+        rpc_request = await self.send(method, params)
+        provider._request_processor.cache_request_information(
+            rpc_request["id"],
+            rpc_request["method"],
+            rpc_request["params"],
+            response_formatters=response_formatters or ((), (), ()),
+        )
+        return await self.recv_for_request(rpc_request)
 
-    async def send(self, method: RPCEndpoint, params: Any) -> None:
+    async def send(self, method: RPCEndpoint, params: Any) -> RPCRequest:
         provider = cast(PersistentConnectionProvider, self._provider)
-        # run through the request processors of the middleware
-        for mw_class in self.middleware_onion.as_tuple_of_middleware():
-            mw = mw_class(self.w3)
-            method, params = mw.request_processor(method, params)
-
+        async_w3 = cast("AsyncWeb3", self.w3)
+        middleware_onion = cast("MiddlewareOnion", self.middleware_onion)
+        send_func = await provider.send_func(
+            async_w3,
+            middleware_onion,
+        )
         self.logger.debug(
             "Sending request to open socket connection: "
-            f"{provider.get_endpoint_uri_or_ipc_path()}, method: {method}"
+            f"{provider.get_endpoint_uri_or_ipc_path()},\n    method: {method},\n"
+            f"    params: {params}"
+        )
+        return await send_func(method, params)
+
+    async def recv_for_request(self, rpc_request: RPCRequest) -> RPCResponse:
+        provider = cast(PersistentConnectionProvider, self._provider)
+        async_w3 = cast("AsyncWeb3", self.w3)
+        middleware_onion = cast("MiddlewareOnion", self.middleware_onion)
+        recv_func = await provider.recv_func(
+            async_w3,
+            middleware_onion,
+        )
+        self.logger.debug(
+            "Getting response for request from open socket connection:\n"
+            f"    request: {rpc_request}"
         )
-        await provider.socket_send(provider.encode_rpc_request(method, params))
+        response = await recv_func(rpc_request)
+        try:
+            return cast(RPCResponse, await self._process_response(response))
+        except Exception:
+            response_id_key = generate_cache_key(response["id"])
+            provider._request_processor._request_information_cache.pop(response_id_key)
+            raise
 
     async def recv(self) -> Union[RPCResponse, FormattedEthSubscriptionResponse]:
         provider = cast(PersistentConnectionProvider, self._provider)

web3/method.py

@@ -203,7 +203,7 @@ def input_munger(self, module: "Module", args: Any, kwargs: Any) -> List[Any]:
     def process_params(
         self, module: "Module", *args: Any, **kwargs: Any
     ) -> Tuple[
-        Tuple[Union[RPCEndpoint, Callable[..., RPCEndpoint]], Tuple[Any, ...]],
+        Tuple[Union[RPCEndpoint, Callable[..., RPCEndpoint]], Tuple[RPCEndpoint, ...]],
         Tuple[
             Union[TReturn, Dict[str, Callable[..., Any]]],
             Callable[..., Any],