1 Commits

Author SHA1 Message Date
e649f356be Turn KeepAliveHandler into an asyncio Task 2021-10-07 17:22:32 +01:00
3 changed files with 28 additions and 71 deletions

View File

@ -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 = []
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
attachments = components = reactions = stickers = 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.

View File

@ -153,38 +153,38 @@ class GatewayRatelimiter:
await asyncio.sleep(delta)
class KeepAliveHandler(threading.Thread):
def __init__(self, *args: Any, **kwargs: Any) -> None:
ws = kwargs.pop("ws")
interval = kwargs.pop("interval", None)
shard_id = kwargs.pop("shard_id", None)
threading.Thread.__init__(self, *args, **kwargs)
class KeepAliveHandler:
def __init__(self, *, ws: DiscordWebSocket, shard_id: int = None, interval: float = None) -> None:
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._last_ack: float = time.perf_counter()
self._stop_ev: asyncio.Event = asyncio.Event()
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:
while not self._stop_ev.wait(self.interval):
async def run(self) -> None:
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]
except KeyError:
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}"
stack = "".join(traceback.format_stack())
msg = f"{self.block_msg}\nLoop 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()

View File

@ -94,18 +94,3 @@ class Object(Hashable):
def created_at(self) -> datetime.datetime:
""":class:`datetime.datetime`: Returns the snowflake's creation time in UTC."""
return utils.snowflake_time(self.id)
@property
def worker_id(self) -> int:
""":class:`int`: Returns the worker id that made the snowflake."""
return (self.id & 0x3E0000) >> 17
@property
def process_id(self) -> int:
""":class:`int`: Returns the process id that made the snowflake."""
return (self.id & 0x1F000) >> 12
@property
def increment_id(self) -> int:
""":class:`int`: Returns the increment id that made the snowflake."""
return (self.id & 0xFFF)