Repository: hikari-py/hikari
Pull request #1648 (2023-06-25 20:17:17)

Retry REST requests on connection errors
Continuation of #1132




Patch:
diff --git a/changes/1648.bugfix.md b/changes/1648.bugfix.md
new file mode 100644
index 00000000..64c2b659
--- /dev/null
+++ b/changes/1648.bugfix.md
@@ -0,0 +1,1 @@
+Retry REST requests on connection errors
diff --git a/dev-requirements/mypy.txt b/dev-requirements/mypy.txt
index 8ad26ea2..49dc1503 100644
--- a/dev-requirements/mypy.txt
+++ b/dev-requirements/mypy.txt
@@ -1,1 +1,1 @@
-mypy==1.4.0
+mypy==1.4.1
diff --git a/dev-requirements/pyright.txt b/dev-requirements/pyright.txt
index 675d73c4..d023f962 100644
--- a/dev-requirements/pyright.txt
+++ b/dev-requirements/pyright.txt
@@ -1,1 +1,1 @@
-pyright==1.1.315
+pyright==1.1.316
diff --git a/hikari/impl/rest.py b/hikari/impl/rest.py
index 77030bdc..024704ff 100644
--- a/hikari/impl/rest.py
+++ b/hikari/impl/rest.py
@@ -736,8 +736,10 @@ async def _request(
 
         raise errors.ComponentStateConflictError("The REST client was closed mid-request")
 
+    # Ignore too long and too complex, respectively
+    # We rather keep everything we can here inline.
     @typing.final
-    async def _perform_request(
+    async def _perform_request(  # noqa: CFQ001, C901
         self,
         compiled_route: routes.CompiledRoute,
         *,
@@ -779,13 +781,14 @@ async def _perform_request(
         url = compiled_route.create_url(self._rest_url)
 
         stack = contextlib.AsyncExitStack()
-        # This is initiated the first time we hit a 5xx error to save a little memory when nothing goes wrong
+        # This is initiated the first time we time out or hit a 5xx error to
+        # save a little memory when nothing goes wrong
         backoff: typing.Optional[rate_limits.ExponentialBackOff] = None
         retry_count = 0
         trace_logging_enabled = _LOGGER.isEnabledFor(ux.TRACE)
 
         while True:
-            async with stack:
+            try:
                 if form_builder:
                     data = await form_builder.build(stack, executor=self._executor)
 
@@ -832,6 +835,28 @@ async def _perform_request(
                 # Ensure we are not rate limited, and update rate limiting headers where appropriate.
                 time_before_retry = await self._parse_ratelimits(compiled_route, auth, response)
 
+            except (asyncio.TimeoutError, aiohttp.ClientConnectionError) as ex:
+                if retry_count >= self._max_retries:
+                    raise errors.HTTPError(message=str(ex)) from ex
+
+                if backoff is None:
+                    backoff = rate_limits.ExponentialBackOff(maximum=_MAX_BACKOFF_DURATION)
+
+                sleep_time = next(backoff)
+                _LOGGER.warning(
+                    "Connection error (%s), backing off for %.2fs and retrying. Retries remaining: %s",
+                    type(ex).__name__,
+                    sleep_time,
+                    self._max_retries - retry_count,
+                )
+                retry_count += 1
+
+                await asyncio.sleep(sleep_time)
+                continue
+
+            finally:
+                await stack.aclose()
+
             if time_before_retry is not None:
                 await asyncio.sleep(time_before_retry)
                 continue


Test patch:
diff --git a/tests/hikari/impl/test_rest.py b/tests/hikari/impl/test_rest.py
index 972251b6..15132aa7 100644
--- a/tests/hikari/impl/test_rest.py
+++ b/tests/hikari/impl/test_rest.py
@@ -26,6 +26,7 @@
 import re
 import typing
 
+import aiohttp
 import mock
 import pytest
 
@@ -1913,6 +1914,33 @@ class StubResponse:
         asyncio_sleep.assert_has_awaits([mock.call(1), mock.call(2), mock.call(3)])
         generate_error_response.assert_called_once_with(rest_client._client_session.request.return_value)
 
+    @hikari_test_helpers.timeout()
+    @pytest.mark.parametrize("exception", [asyncio.TimeoutError, aiohttp.ClientConnectionError])
+    async def test_perform_request_when_connection_error_will_retry_until_exhausted(self, rest_client, exception):
+        route = routes.Route("GET", "/something/{channel}/somewhere").compile(channel=123)
+        mock_session = mock.AsyncMock(request=mock.AsyncMock(side_effect=exception))
+        rest_client._max_retries = 3
+        rest_client._parse_ratelimits = mock.AsyncMock()
+        rest_client._client_session = mock_session
+
+        stack = contextlib.ExitStack()
+        stack.enter_context(pytest.raises(errors.HTTPError))
+        exponential_backoff = stack.enter_context(
+            mock.patch.object(
+                rate_limits,
+                "ExponentialBackOff",
+                return_value=mock.Mock(__next__=mock.Mock(side_effect=[1, 2, 3, 4, 5])),
+            )
+        )
+        asyncio_sleep = stack.enter_context(mock.patch.object(asyncio, "sleep"))
+
+        with stack:
+            await rest_client._perform_request(route)
+
+        assert exponential_backoff.return_value.__next__.call_count == 3
+        exponential_backoff.assert_called_once_with(maximum=16)
+        asyncio_sleep.assert_has_awaits([mock.call(1), mock.call(2), mock.call(3)])
+
     @pytest.mark.parametrize("enabled", [True, False])
     @hikari_test_helpers.timeout()
     async def test_perform_request_logger(self, rest_client, enabled):