feat(gateway): Client Latency support and introduce ready state (#548)

* feat(gateway): Client Latency support, introduce ready state, fix headers/docs

* feat: Extend helper method to the client class.

* fix: Add attempt to fix gateway recursion, minor refactor gateway header usage.

* docs: Document new gateway latency attributes and ready state.

Co-authored-by: EdVraz <88881326+EdVraz@users.noreply.github.com>

Author: FayeDel

interactions/api/gateway.py

@@ -13,8 +13,8 @@
     new_event_loop,
     sleep,
 )
-from logging import Logger
 from sys import platform, version_info
+from time import perf_counter
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 from aiohttp import WSMessage
@@ -30,7 +30,7 @@
 from .models.misc import MISSING
 from .models.presence import ClientPresence
 
-log: Logger = get_logger("gateway")
+log = get_logger("gateway")
 
 
 __all__ = ("_Heartbeat", "WebSocketClient")
@@ -69,9 +69,13 @@ class WebSocketClient:
     :ivar _Heartbeat __heartbeater: The context state of a "heartbeat" made to the Gateway.
     :ivar Optional[List[Tuple[int]]] __shard: The shards used during connection.
     :ivar Optional[ClientPresence] __presence: The presence used in connection.
+    :ivar Event ready: The ready state of the client as an ``asyncio.Event``.
     :ivar Task __task: The closing task for ending connections.
-    :ivar int session_id: The ID of the ongoing session.
-    :ivar str sequence: The sequence identifier of the ongoing session.
+    :ivar Optional[str] session_id: The ID of the ongoing session.
+    :ivar Optional[int] sequence: The sequence identifier of the ongoing session.
+    :ivar float _last_send: The latest time of the last send_packet function call since connection creation, in seconds.
+    :ivar float _last_ack: The latest time of the last ``HEARTBEAT_ACK`` event since connection creation, in seconds.
+    :ivar float latency: The latency of the connection, in seconds.
     """
 
     __slots__ = (
@@ -90,13 +94,16 @@ class WebSocketClient:
         "session_id",
         "sequence",
         "ready",
+        "_last_send",
+        "_last_ack",
+        "latency",
     )
 
     def __init__(
         self,
         token: str,
         intents: Intents,
-        session_id: Optional[int] = MISSING,
+        session_id: Optional[str] = MISSING,
         sequence: Optional[int] = MISSING,
     ) -> None:
         """
@@ -105,7 +112,7 @@ def __init__(
         :param intents: The Gateway intents of the application for event dispatch.
         :type intents: Intents
         :param session_id?: The ID of the session if trying to reconnect. Defaults to ``None``.
-        :type session_id: Optional[int]
+        :type session_id: Optional[str]
         :param sequence?: The identifier sequence if trying to reconnect. Defaults to ``None``.
         :type sequence: Optional[int]
         """
@@ -132,7 +139,12 @@ def __init__(
         self.__task = None
         self.session_id = None if session_id is MISSING else session_id
         self.sequence = None if sequence is MISSING else sequence
-        self.ready = None
+        self.ready = Event(loop=self._loop) if version_info < (3, 10) else Event()
+
+        self._last_send = perf_counter()
+        self._last_ack = perf_counter()
+        self.latency: float("nan")  # noqa: F821
+        # self.latency has to be noqa, this is valid in python but not in Flake8.
 
     async def _manage_heartbeat(self) -> None:
         """Manages the heartbeat loop."""
@@ -186,7 +198,7 @@ async def _establish_connection(
                     continue
                 if self._client.close_code in range(4010, 4014) or self._client.close_code == 4004:
                     raise GatewayException(self._client.close_code)
-                elif self._client.close_code is not None:
+                elif self._closed:  # Redundant conditional.
                     await self._establish_connection()
 
                 await self._handle_connection(stream, shard, presence)
@@ -226,8 +238,10 @@ async def _handle_connection(
             if op == OpCodeType.HEARTBEAT:
                 await self.__heartbeat()
             if op == OpCodeType.HEARTBEAT_ACK:
+                self._last_ack = perf_counter()
                 log.debug("HEARTBEAT_ACK")
                 self.__heartbeater.event.set()
+                self.latency = self._last_ack - self._last_send
             if op in (OpCodeType.INVALIDATE_SESSION, OpCodeType.RECONNECT):
                 log.debug("INVALID_SESSION/RECONNECT")
 
@@ -243,10 +257,15 @@ async def _handle_connection(
             self.sequence = stream["s"]
             self._dispatch.dispatch("on_ready")
             log.debug(f"READY (session_id: {self.session_id}, seq: {self.sequence})")
+            self.ready.set()
         else:
             log.debug(f"{event}: {data}")
             self._dispatch_event(event, data)
 
+    async def wait_until_ready(self):
+        """Waits for the client to become ready according to the Gateway."""
+        await self.ready.wait()
+
     def _dispatch_event(self, event: str, data: dict) -> None:
         """
         Dispatches an event from the Gateway.
@@ -282,7 +301,7 @@ def _dispatch_event(self, event: str, data: dict) -> None:
                     )
                 else:
                     _context = self.__contextualize(data)
-                    _name: str
+                    _name: str = ""
                     __args: list = [_context]
                     __kwargs: dict = {}
 
@@ -339,9 +358,9 @@ def _dispatch_event(self, event: str, data: dict) -> None:
 
                         self._dispatch.dispatch("on_modal", _context)
 
-                self._dispatch.dispatch(_name, *__args, **__kwargs)
-                self._dispatch.dispatch("on_interaction", _context)
-                self._dispatch.dispatch("on_interaction_create", _context)
+                    self._dispatch.dispatch(_name, *__args, **__kwargs)
+                    self._dispatch.dispatch("on_interaction", _context)
+                    self._dispatch.dispatch("on_interaction_create", _context)
 
         self._dispatch.dispatch("raw_socket_create", data)
 
@@ -459,6 +478,7 @@ async def _send_packet(self, data: Dict[str, Any]) -> None:
         :param data: The data to send to the Gateway.
         :type data: Dict[str, Any]
         """
+        self._last_send = perf_counter()
         packet: str = dumps(data).decode("utf-8") if isinstance(data, dict) else data
         await self._client.send_str(packet)
         log.debug(packet)

interactions/api/gateway.pyi

@@ -4,20 +4,19 @@ from asyncio import (
     Task,
 )
 from logging import Logger
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union, Iterable
 
 from aiohttp import ClientWebSocketResponse
 
-from ..base import get_logger
+from ..models import Option
 from ..api.models.misc import MISSING
 from ..api.models.presence import ClientPresence
 from .dispatch import Listener
 from .http import HTTPClient
 from .models.flags import Intents
 
-log: Logger = get_logger("gateway")
-
-__all__ = ("_Heartbeat", "WebSocketClient")
+log: Logger
+__all__: Iterable[str]
 
 class _Heartbeat:
     event: Event
@@ -37,8 +36,13 @@ class WebSocketClient:
     __shard: Optional[List[Tuple[int]]]
     __presence: Optional[ClientPresence]
     __task: Optional[Task]
-    session_id: int
-    sequence: str
+    session_id: Optional[str]
+    sequence: Optional[int]
+    _last_send: float
+    _last_ack: float
+    latency: float
+    ready: Event
+
     def __init__(
         self,
         token: str,
@@ -57,6 +61,11 @@ class WebSocketClient:
         shard: Optional[List[Tuple[int]]] = MISSING,
         presence: Optional[ClientPresence] = MISSING,
     ) -> None: ...
+    async def wait_until_ready(self) -> None: ...
+    def _dispatch_event(self, event: str, data: dict) -> None: ...
+    def __contextualize(self, data: dict) -> object: ...
+    def __sub_command_context(self, data: Union[dict, Option]) -> Union[Tuple[str], dict]: ...
+    def __option_type_context(self, context: object, type: int) -> dict: ...
     @property
     async def __receive_packet_stream(self) -> Optional[Dict[str, Any]]: ...
     async def _send_packet(self, data: Dict[str, Any]) -> None: ...

interactions/client.py

@@ -98,6 +98,12 @@ def __init__(
         data = self._loop.run_until_complete(self._http.get_current_bot_information())
         self.me = Application(**data)
 
+    @property
+    def latency(self) -> float:
+        """Returns the connection latency in milliseconds."""
+
+        return self._websocket.latency * 1000
+
     def start(self) -> None:
         """Starts the client session."""
         self._loop.run_until_complete(self._ready())
@@ -305,6 +311,10 @@ async def _login(self) -> None:
         while not self._websocket._closed:
             await self._websocket._establish_connection(self._shard, self._presence)
 
+    async def wait_until_ready(self) -> None:
+        """Helper method that waits until the websocket is ready."""
+        await self._websocket.wait_until_ready()
+
     def event(self, coro: Coroutine, name: Optional[str] = MISSING) -> Callable[..., Any]:
         """
         A decorator for listening to events dispatched from the

interactions/client.pyi

@@ -34,6 +34,8 @@ class Client:
         token: str,
         **kwargs,
     ) -> None: ...
+    @property
+    def latency(self) -> float: ...
     def start(self) -> None: ...
     def __register_events(self) -> None: ...
     async def __compare_sync(self, data: dict) -> None: ...
@@ -44,6 +46,7 @@ class Client:
     async def _synchronize(self, payload: Optional[dict] = None) -> None: ...
     async def _ready(self) -> None: ...
     async def _login(self) -> None: ...
+    async def wait_until_ready(self) -> None: ...
     def event(self, coro: Coroutine, name: Optional[str] = None) -> Callable[..., Any]: ...
     def __check_command(
         self,