Turn KeepAliveHandler into an asyncio Task

discord/ext/commands/bot.py

@@ -28,7 +28,6 @@ from __future__ import annotations
 import asyncio
 import collections
 import collections.abc
-from functools import cached_property
 
 import inspect
 import importlib.util
@@ -73,9 +72,7 @@ from .cog import Cog
 if TYPE_CHECKING:
     import importlib.machinery
 
-    from discord.role import Role
     from discord.message import Message
-    from discord.abc import PartialMessageableChannel
     from ._types import (
         Check,
         CoroFunc,
@@ -97,16 +94,9 @@ CXT = TypeVar("CXT", bound="Context")
 
 class _FakeSlashMessage(discord.PartialMessage):
     activity = application = edited_at = reference = webhook_id = None
-    attachments = components = reactions = stickers = []
+    attachments = components = reactions = stickers = mentions = []
-    tts = False
-
-    raw_mentions = discord.Message.raw_mentions
-    clean_content = discord.Message.clean_content
-    channel_mentions = discord.Message.channel_mentions
-    raw_role_mentions = discord.Message.raw_role_mentions
-    raw_channel_mentions = discord.Message.raw_channel_mentions
-
     author: Union[discord.User, discord.Member]
+    tts = False
 
     @classmethod
     def from_interaction(
@@ -118,22 +108,6 @@ class _FakeSlashMessage(discord.PartialMessage):
 
         return self
 
-    @cached_property
-    def mentions(self) -> List[Union[discord.Member, discord.User]]:
-        client = self._state._get_client()
-        if self.guild:
-            ensure_user = lambda id: self.guild.get_member(id) or client.get_user(id) # type: ignore
-        else:
-            ensure_user = client.get_user
-
-        return discord.utils._unique(filter(None, map(ensure_user, self.raw_mentions)))
-
-    @cached_property
-    def role_mentions(self) -> List[Role]:
-        if self.guild is None:
-            return []
-        return discord.utils._unique(filter(None, map(self.guild.get_role, self.raw_role_mentions)))
-
 
 def when_mentioned(bot: Union[Bot, AutoShardedBot], msg: Message) -> List[str]:
     """A callable that implements a command prefix equivalent to being mentioned.

discord/gateway.py

@@ -153,38 +153,38 @@ class GatewayRatelimiter:
                 await asyncio.sleep(delta)
 
 
-class KeepAliveHandler(threading.Thread):
+class KeepAliveHandler:
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
+    def __init__(self, *, ws: DiscordWebSocket, shard_id: int = None, interval: float = None) -> None:
-        ws = kwargs.pop("ws")
-        interval = kwargs.pop("interval", None)
-        shard_id = kwargs.pop("shard_id", None)
-        threading.Thread.__init__(self, *args, **kwargs)
         self.ws: DiscordWebSocket = ws
-        self._main_thread_id: int = ws.thread_id
-        self.interval: Optional[float] = interval
-        self.daemon: bool = True
         self.shard_id: Optional[int] = shard_id
+        self.interval: Optional[float] = interval
+        self.heartbeat_timeout: float = self.ws._max_heartbeat_timeout
+
         self.msg: str = "Keeping shard ID %s websocket alive with sequence %s."
         self.block_msg: str = "Shard ID %s heartbeat blocked for more than %s seconds."
         self.behind_msg: str = "Can't keep up, shard ID %s websocket is %.1fs behind."
-        self._stop_ev: threading.Event = threading.Event()
+        self._stop_ev: asyncio.Event = asyncio.Event()
-        self._last_ack: float = time.perf_counter()
         self._last_send: float = time.perf_counter()
         self._last_recv: float = time.perf_counter()
+        self._last_ack: float = time.perf_counter()
         self.latency: float = float("inf")
-        self.heartbeat_timeout: float = ws._max_heartbeat_timeout
 
-    def run(self) -> None:
+    async def run(self) -> None:
-        while not self._stop_ev.wait(self.interval):
+        while True:
+            try:
+                await asyncio.wait_for(self._stop_ev.wait(), timeout=self.interval)
+            except asyncio.TimeoutError:
+                pass
+            else:
+                return
+
             if self._last_recv + self.heartbeat_timeout < time.perf_counter():
                 _log.warning(
                     "Shard ID %s has stopped responding to the gateway. Closing and restarting.", self.shard_id
                 )
-                coro = self.ws.close(4000)
-                f = asyncio.run_coroutine_threadsafe(coro, loop=self.ws.loop)
 
                 try:
-                    f.result()
+                    await self.ws.close(4000)
                 except Exception:
                     _log.exception("An error occurred while stopping the gateway. Ignoring.")
                 finally:
@@ -193,24 +193,18 @@ class KeepAliveHandler(threading.Thread):
 
             data = self.get_payload()
             _log.debug(self.msg, self.shard_id, data["d"])
-            coro = self.ws.send_heartbeat(data)
-            f = asyncio.run_coroutine_threadsafe(coro, loop=self.ws.loop)
             try:
                 # block until sending is complete
                 total = 0
                 while True:
                     try:
-                        f.result(10)
+                        await asyncio.wait_for(self.ws.send_heartbeat(data), timeout=10)
                         break
-                    except concurrent.futures.TimeoutError:
+                    except asyncio.TimeoutError:
                         total += 10
-                        try:
+
-                            frame = sys._current_frames()[self._main_thread_id]
+                        stack = "".join(traceback.format_stack())
-                        except KeyError:
+                        msg = f"{self.block_msg}\nLoop traceback (most recent call last):\n{stack}"
-                            msg = self.block_msg
-                        else:
-                            stack = "".join(traceback.format_stack(frame))
-                            msg = f"{self.block_msg}\nLoop thread traceback (most recent call last):\n{stack}"
                         _log.warning(msg, self.shard_id, total)
 
             except Exception:
@@ -225,6 +219,10 @@ class KeepAliveHandler(threading.Thread):
             "d": self.ws.sequence,  # type: ignore
         }
 
+
+    def start(self) -> None:
+        self.ws.loop.create_task(self.run())
+
     def stop(self) -> None:
         self._stop_ev.set()
 

discord/utils.py

@@ -84,7 +84,6 @@ __all__ = (
     "escape_mentions",
     "as_chunks",
     "format_dt",
-    "generate_snowflake",
 )
 
 DISCORD_EPOCH = 1420070400000
@@ -1021,23 +1020,3 @@ def format_dt(dt: datetime.datetime, /, style: Optional[TimestampStyle] = None)
     if style is None:
         return f"<t:{int(dt.timestamp())}>"
     return f"<t:{int(dt.timestamp())}:{style}>"
-
-
-def generate_snowflake(dt: Optional[datetime.datetime] = None) -> int:
-    """Returns a numeric snowflake pretending to be created at the given date but more accurate and random than time_snowflake.
-    If No dt is not passed, it makes one from the current time using utcnow.
-
-    Parameters
-    -----------
-    dt: :class:`datetime.datetime`
-        A datetime object to convert to a snowflake.
-        If naive, the timezone is assumed to be local time.
-
-    Returns
-    --------
-    :class:`int`
-        The snowflake representing the time given.
-    """
-
-    dt = dt or utcnow()
-    return int(dt.timestamp() * 1000 - DISCORD_EPOCH) << 22 | 0x3fffff

docs/api.rst

@@ -1136,8 +1136,6 @@ Utility Functions
 
 .. autofunction:: discord.utils.as_chunks
 
-.. autofunction:: discord.utils.generate_snowflake
-
 .. _discord-api-enums:
 
 Enumerations