Repository: hikari-py/hikari
Pull request #1645 (2023-06-22 10:33:56)

Handle connection reset error


Patch:
diff --git a/changes/1645.bugfix.md b/changes/1645.bugfix.md
new file mode 100644
index 00000000..9133cea8
--- /dev/null
+++ b/changes/1645.bugfix.md
@@ -0,0 +1,1 @@
+Handle connection reset error on shards.
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 cd5698b8..d023f962 100644
--- a/dev-requirements/pyright.txt
+++ b/dev-requirements/pyright.txt
@@ -1,1 +1,1 @@
-pyright==1.1.314
+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
diff --git a/hikari/impl/shard.py b/hikari/impl/shard.py
index 97d6ff42..607eb831 100644
--- a/hikari/impl/shard.py
+++ b/hikari/impl/shard.py
@@ -886,6 +886,9 @@ async def _keep_alive(self) -> None:
                 # Since nothing went wrong, we can reset the backoff and try again
                 backoff.reset()
 
+            except ConnectionResetError:
+                self._logger.warning("connection got reset by server. Will retry shortly")
+
             except errors.GatewayConnectionError as ex:
                 self._logger.warning("failed to communicate with server, reason was: %r. Will retry shortly", ex.reason)
 


Test patch:
diff --git a/dev-requirements/pytest.txt b/dev-requirements/pytest.txt
index 60318ad3..e6753936 100644
--- a/dev-requirements/pytest.txt
+++ b/dev-requirements/pytest.txt
@@ -1,7 +1,7 @@
 # mock in the stdlib changes between versions of Python, so are not consistent in behaviour like the backport is
 mock==5.0.2
 
-pytest==7.3.2
+pytest==7.4.0
 pytest-asyncio==0.21.0
 pytest-cov==4.1.0
 pytest-randomly==3.12.0
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):