From e7d918f5140deb87df410344d3e38059598d5561 Mon Sep 17 00:00:00 2001
From: Luke Van Seters <lukevanseters@gmail.com>
Date: Sat, 11 Sep 2021 14:35:09 -0600
Subject: [PATCH] Add backoff to http retry middleware

---
 .pre-commit-config.yaml  |  1 +
 mev_inspect/retry.py     | 59 ++++++++++++++++++++++++++++++++++++++++
 pyproject.toml           |  1 -
 scripts/inspect_block.py | 16 +++++++++--
 4 files changed, 74 insertions(+), 3 deletions(-)
 create mode 100644 mev_inspect/retry.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5e9c7df..397ce75 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -18,3 +18,4 @@ repos:
       - id: 'mypy'
         additional_dependencies:
           - 'pydantic'
+          - 'types-requests'
diff --git a/mev_inspect/retry.py b/mev_inspect/retry.py
new file mode 100644
index 0000000..c85acdf
--- /dev/null
+++ b/mev_inspect/retry.py
@@ -0,0 +1,59 @@
+import time
+from typing import (
+    Any,
+    Callable,
+    Collection,
+    Type,
+)
+
+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.types import (
+    RPCEndpoint,
+    RPCResponse,
+)
+
+
+def exception_retry_with_backoff_middleware(
+    make_request: Callable[[RPCEndpoint, Any], RPCResponse],
+    web3: Web3,  # pylint: disable=unused-argument
+    errors: Collection[Type[BaseException]],
+    retries: int = 5,
+    backoff_time_seconds: float = 0.1,
+) -> Callable[[RPCEndpoint, Any], RPCResponse]:
+    """
+    Creates middleware that retries failed HTTP requests. Is a default
+    middleware for HTTPProvider.
+    """
+
+    def middleware(method: RPCEndpoint, params: Any) -> RPCResponse:
+        if check_if_retry_on_failure(method):
+            for i in range(retries):
+                try:
+                    return make_request(method, params)
+                # https://github.com/python/mypy/issues/5349
+                except errors:  # type: ignore
+                    if i < retries - 1:
+                        time.sleep(backoff_time_seconds)
+                        continue
+                    else:
+                        raise
+            return None
+        else:
+            return make_request(method, params)
+
+    return middleware
+
+
+def http_retry_with_backoff_request_middleware(
+    make_request: Callable[[RPCEndpoint, Any], Any], web3: Web3
+) -> Callable[[RPCEndpoint, Any], Any]:
+    return exception_retry_with_backoff_middleware(
+        make_request, web3, (ConnectionError, HTTPError, Timeout, TooManyRedirects)
+    )
diff --git a/pyproject.toml b/pyproject.toml
index 04937a5..3d7b66d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -32,7 +32,6 @@ build-backend = "poetry.core.masonry.api"
 lint = 'scripts.poetry.dev_tools:lint'
 test = 'scripts.poetry.dev_tools:test'
 isort = 'scripts.poetry.dev_tools:isort'
-mypy = 'scripts.poetry.dev_tools:mypy'
 black = 'scripts.poetry.dev_tools:black'
 pre_commit = 'scripts.poetry.dev_tools:pre_commit'
 start = 'scripts.poetry.docker:start'
diff --git a/scripts/inspect_block.py b/scripts/inspect_block.py
index caa174e..7a72bfa 100644
--- a/scripts/inspect_block.py
+++ b/scripts/inspect_block.py
@@ -22,6 +22,7 @@ from mev_inspect.crud.swaps import delete_swaps_for_block, write_swaps
 from mev_inspect.db import get_session
 from mev_inspect.miner_payments import get_miner_payments
 from mev_inspect.swaps import get_swaps
+from mev_inspect.retry import http_retry_with_backoff_request_middleware
 
 
 @click.group()
@@ -34,7 +35,7 @@ def cli():
 @click.argument("rpc")
 @click.option("--cache/--no-cache", default=True)
 def inspect_block(block_number: int, rpc: str, cache: bool):
-    base_provider = Web3.HTTPProvider(rpc)
+    base_provider = _get_base_provider(rpc)
     w3 = Web3(base_provider)
 
     if not cache:
@@ -49,7 +50,7 @@ def inspect_block(block_number: int, rpc: str, cache: bool):
 @click.argument("rpc")
 @click.option("--cache/--no-cache", default=True)
 def inspect_many_blocks(after_block: int, before_block: int, rpc: str, cache: bool):
-    base_provider = Web3.HTTPProvider(rpc)
+    base_provider = _get_base_provider(rpc)
     w3 = Web3(base_provider)
 
     if not cache:
@@ -156,5 +157,16 @@ def get_stats(classified_traces) -> dict:
     return stats
 
 
+def _get_base_provider(rpc: str) -> Web3.HTTPProvider:
+    base_provider = Web3.HTTPProvider(rpc)
+    base_provider.middlewares.remove("http_retry_request")
+    base_provider.middlewares.add(
+        http_retry_with_backoff_request_middleware,
+        "http_retry_with_backoff",
+    )
+
+    return base_provider
+
+
 if __name__ == "__main__":
     cli()