From 7707b818f0283668be40632f8581c193328a9f55 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Sat, 25 Dec 2021 17:23:21 -0500
Subject: [PATCH] Include new methods in retry-able methods

---
 mev_inspect/retry.py | 26 ++++++++++++++++++++++----
 1 file changed, 22 insertions(+), 4 deletions(-)

diff --git a/mev_inspect/retry.py b/mev_inspect/retry.py
index ce2663c..25273be 100644
--- a/mev_inspect/retry.py
+++ b/mev_inspect/retry.py
@@ -5,6 +5,7 @@ from time import sleep
 from typing import Any, Callable, Collection, Coroutine, Type
 
 from aiohttp.client_exceptions import (
+    ClientConnectorError,
     ClientOSError,
     ClientResponseError,
     ServerDisconnectedError,
@@ -12,20 +13,33 @@ from aiohttp.client_exceptions import (
 )
 from requests.exceptions import ConnectionError, HTTPError, Timeout, TooManyRedirects
 from web3 import Web3
-from web3.middleware.exception_retry_request import check_if_retry_on_failure
+from web3.middleware.exception_retry_request import whitelist
 from web3.types import RPCEndpoint, RPCResponse
 
 request_exceptions = (ConnectionError, HTTPError, Timeout, TooManyRedirects)
 aiohttp_exceptions = (
     ClientOSError,
+    ClientResponseError,
+    ClientConnectorError,
     ServerDisconnectedError,
     ServerTimeoutError,
-    ClientResponseError,
 )
 
+whitelist_additions = ["eth_getBlockReceipts", "trace_block", "eth_feeHistory"]
+
 logger = logging.getLogger(__name__)
 
 
+def check_if_retry_on_failure(method: RPCEndpoint) -> bool:
+    root = method.split("_")[0]
+    if root in (whitelist + whitelist_additions):
+        return True
+    elif method in (whitelist + whitelist_additions):
+        return True
+    else:
+        return False
+
+
 async def exception_retry_with_backoff_middleware(
     make_request: Callable[[RPCEndpoint, Any], Any],
     web3: Web3,  # pylint: disable=unused-argument
@@ -47,7 +61,7 @@ async def exception_retry_with_backoff_middleware(
                 # https://github.com/python/mypy/issues/5349
                 except errors:  # type: ignore
                     logger.error(
-                        f"Request for method {method}, block: {int(params[0], 16)}, retrying: {i}/{retries}"
+                        f"Request for method {method}, params: {params}, retrying: {i}/{retries}"
                     )
                     if i < (retries - 1):
                         backoff_time = backoff_time_seconds * (
@@ -72,5 +86,9 @@ async def http_retry_with_backoff_request_middleware(
     return await exception_retry_with_backoff_middleware(
         make_request,
         web3,
-        (request_exceptions + aiohttp_exceptions + (TimeoutError,)),
+        (
+            request_exceptions
+            + aiohttp_exceptions
+            + (TimeoutError, ConnectionRefusedError)
+        ),
     )