Fixed more edge case scenarios for gateway termination and fixed tests.

hikari/impl/shard.py

@@ -666,14 +666,14 @@ async def _identify(self) -> None:
         await self._ws.send_json(payload)  # type: ignore[union-attr]
 
     async def _heartbeat(self, heartbeat_interval: float) -> bool:
-        # Return True if zombied.
+        # Return True if zombied or should reconnect, false if time to die forever.
         # Prevent immediately zombie-ing.
         self._last_heartbeat_ack_received = date.monotonic()
         self._logger.debug("starting heartbeat with interval %ss", heartbeat_interval)
 
-        while True:
+        while not self._closing.is_set() and not self._closed.is_set():
             if self._last_heartbeat_ack_received <= self._last_heartbeat_sent:
-                # Gateway is zombie
+                # Gateway is zombie, close and request reconnect.
                 self._logger.warning(
                     "connection has not received a HEARTBEAT_ACK for approx %.1fs and is being disconnected, "
                     "expect a reconnect shortly",
@@ -690,11 +690,14 @@ async def _heartbeat(self, heartbeat_interval: float) -> bool:
             try:
                 await asyncio.wait_for(self._closing.wait(), timeout=heartbeat_interval)
                 # We are closing
-                return False
+                break
             except asyncio.TimeoutError:
                 # We should continue
                 continue
 
+        self._logger.debug("heartbeat task is finishing now")
+        return False
+
     async def _poll_events(self) -> typing.Optional[bool]:
         payload = await self._ws.receive_json(timeout=5)  # type: ignore[union-attr]
 
@@ -742,6 +745,7 @@ async def _resume(self) -> None:
 
     async def _run(self) -> None:
         self._closed.clear()
+        self._closing.clear()
         last_started_at = -float("inf")
 
         backoff = rate_limits.ExponentialBackOff(
@@ -751,7 +755,7 @@ async def _run(self) -> None:
         )
 
         try:
-            while not self._closing.is_set() and not self._closed:
+            while not self._closing.is_set() and not self._closed.is_set():
                 if date.monotonic() - last_started_at < _BACKOFF_WINDOW:
                     time = next(backoff)
                     self._logger.info("backing off reconnecting for %.2fs", time)
@@ -801,10 +805,10 @@ async def _run(self) -> None:
                     self._logger.error("encountered some unhandled error", exc_info=ex)
                     raise
         finally:
+            self._closing.set()
             self._closed.set()
 
     async def _run_once(self) -> bool:
-        self._closing.clear()
         self._handshake_completed.clear()
         dispatch_disconnect = False
 

tests/hikari/impl/test_shard.py

@@ -970,7 +970,8 @@ async def test__identify_when_intents(self, client):
     async def test__heartbeat(self, client):
         client._last_heartbeat_sent = 5
         client._logger = mock.Mock()
-        client._closing = mock.Mock()
+        client._closing = mock.Mock(is_set=mock.Mock(return_value=False))
+        client._closed = mock.Mock(is_set=mock.Mock(return_value=False))
         client._send_heartbeat = mock.AsyncMock()
 
         with mock.patch.object(date, "monotonic", return_value=10):