pincer.core.gateway.Gateway.start_loop() - Code Metrics - Inspection of ":bug: :sparkles: Added VERY basic sharding support..." - Pincer-org/Pincer - Measure and Improve Code Quality continuously with Scrutinizer

Passed

Pull Request — main (#329)

unknown

created 2021-12-30 02:45 UTC

pincer.core.gateway.Gateway.start_loop() A

↳ Parent: pincer.core.gateway

Complexity

Conditions

Size

Total Lines	24
Code Lines	15

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
eloc	15
dl	0
loc	24
rs	9.65
c	0
b	0
f	0
cc	4
nop	1

# Copyright Pincer 2021-Present

# Full MIT License can be found in `LICENSE` at the project root.


from __future__ import annotations

from asyncio import create_task, Task, ensure_future, sleep
from dataclasses import dataclass
from datetime import datetime
from itertools import repeat, count, chain
import logging
from platform import system
from random import random
from typing import TYPE_CHECKING, Any, Dict, Callable, Optional
from zlib import decompressobj

from aiohttp import (
    ClientSession, WSMsgType, ClientConnectorError, ClientWebSocketResponse
)

from . import __package__
from ..utils.api_object import APIObject
from .._config import GatewayConfig
from ..core.dispatch import GatewayDispatch
from ..exceptions import (
    InvalidTokenError, GatewayConnectionError, GatewayError, UnhandledException
)

if TYPE_CHECKING:
    from ..objects.app.intents import Intents
    Handler = Callable[[GatewayDispatch], None]

_log = logging.getLogger(__package__)

ZLIB_SUFFIX = b'\x00\x00\xff\xff'
inflator = decompressobj()


@dataclass
class SessionStartLimit(APIObject):
    """Session start limit info returned from the `gateway/bot` endpoint"""
    total: int
    remaining: int
    reset_after: int
    max_concurrency: int


@dataclass
class GatewayInfo(APIObject):
    """Gateway info returned from the `gateway/bot` endpoint"""
    url: str
    shards: int
    session_start_limit: SessionStartLimit


class Gateway:

    """The Gateway handles all interactions with the Discord Websocket API.
    This also contains the main event loop, and handles the heartbeat.

    Running the Gateway will create a connection with the
    Discord Websocket API on behalf of the provided token.

    This token must be a bot token.
    (Which can be found on
    `<https://discord.com/developers/applications/>`_)

    Parameters
    ----------
    token : str.
        The token for this bot
    intents : :class:`~pincer.objects.app.intents.Intents`
        The itents to use. More information can be found at
        `<https://discord.com/developers/docs/topics/gateway#gateway-intents>`_.
    url : str
        The gateway url.
    shard : int
        The ID of the shard to run.
    num_shards : int
        Number used to route traffic to the current. This should usually be the total
        number of shards that will be run. More information at
        `<https://discord.com/developers/docs/topics/gateway#sharding>`_.
    """
    def __init__(
        self,

        token: str, *,

        intents: Intents,

        url: str,

        shard: int,

        num_shards: int

    ) -> None:
        if len(token) != 59:
            raise InvalidTokenError(
                "Discord Token must have exactly 59 characters."
            )

        self.token = token
        self.intents = intents
        self.url = url
        self.shard = shard
        self.num_shards = num_shards
        self.shard_key = [shard, num_shards]

        self.__dispatch_handlers: Dict[int, Handler] = {
            1: self.handle_heartbeat_req,
            7: self.handle_reconnect,
            9: self.handle_invalid_session,
            10: self.identify_and_handle_hello,
            11: self.handle_heartbeat
        }

        # 4000 and 4009 are not included. The client will reconnect when receiving
        # either.
        self.__close_codes: Dict[int, GatewayError] = {
            4001: GatewayError("Invalid opcode was sent"),
            4002: GatewayError("Invalid payload was sent."),
            4003: GatewayError("Payload was sent prior to identifying"),
            4004: GatewayError("Token is not valid"),
            4005: GatewayError(
                "Authentication was sent after client already authenticated"
            ),
            4007: GatewayError("Invalid sequence sent when starting new session"),
            4008: GatewayError("Client was rate limited"),
            4010: GatewayError("Invalid shard"),
            4011: GatewayError("Sharding required"),
            4012: GatewayError("Invalid API version"),
            4013: GatewayError("Invalid intents"),
            4014: GatewayError("Disallowed intents")
        }

        # ClientSession to be used for this Dispatcher
        self.__session: Optional[ClientSession] = None

        # This type `_WSRequestContextManager` isn't exposed by aiohttp.
        # `ClientWebSocketResponse` is a parent class.
        self.__socket: Optional[ClientWebSocketResponse] = None

        # Buffer used to store information in transport conpression.
        self.__buffer = bytearray()

        # The gateway can be disconnected from Discord. This variable stores if the
        # gateway should send a hello or reconnect.
        self.__should_reconnect: bool = False

        # The sequence number for the last received payload. This is used reconnecting.
        self.__sequence_number: int = 0

        # The heartbeat task
        self.__heartbeat_task: Optional[Task] = None

        # Keeps the Client waiting until the next heartbeat
        self.__wait_for_heartbeat: Optional[Task] = None

        # How long the client should wait between each Heartbeat.
        self.__heartbeat_interval: Optional[int] = None

        # Tracks whether the gateway has received an ack (opcode 11) since the last
        # heartbeat.
        #   True: An ack has been received
        #   False: No ack has been received. Attempt to reconnect with gateway,
        self.__has_received_ack: bool = True

        # Session ID received from `on_ready` event. It is set in the `on_ready`
        # middleware. This is used reconnecting.
        self.__session_id: Optional[str] = None

    def __del__(self):
        """Delete method ensures all connections are closed"""
        if self.__socket:
            create_task(self.__socket.close())
        if self.__session:
            create_task(self.__session.close())

    async def init_session(self):
        """
        Crates the ClientSession. ALWAYS run this function right after initializing
        a Gateway.
        """
        self.__session = ClientSession()

    def append_handlers(self, handlers: Dict[int, Handler]):
        """The Client that uses the handler can append their own methods. The gateway
        will run those methods when the specified opcode is received.
        """
        self.__dispatch_handlers = handlers | self.__dispatch_handlers

    def set_session_id(self, _id: str):
        """Session id is private for consistency"""
        self.__session_id = _id

    def decompress_msg(self, msg: bytes) -> Optional[str]:

        if GatewayConfig.compression == "zlib-payload":
            return inflator.decompress(msg)

        if GatewayConfig.compression == "zlib-stream":
            self.__buffer.extend(msg)

            if len(self.__buffer) < 4 or self.__buffer[-4:] != ZLIB_SUFFIX:
                return None

            msg = inflator.decompress(msg)
            self.__buffer = bytearray()
            return msg

        return None

    async def start_loop(self):
        """Instantiate the dispatcher, this will create a connection to the
        Discord websocket API on behalf of the client whose token has
        been passed.
        """
        for _try in count():
            try:
                self.__socket = await self.__session.ws_connect(
                    GatewayConfig.make_uri(self.url)
                )
                break
            except ClientConnectorError as e:
                if _try > GatewayConfig.MAX_RETRIES:
                    raise GatewayConnectionError from e

                _log.warning(
                    "%s Could not open websocket with Discord."
                    " Retrying in 15 seconds...",
                    self.shard_key
                )
                await sleep(15)

        _log.debug("%s Starting envent loop...", self.shard_key)
        await self.event_loop()

    async def event_loop(self):
        """Handles receiving messages and decompressing them if needed"""
        async for msg in self.__socket:
            if msg.type == WSMsgType.TEXT:
                await self.handle_data(msg.data)
            elif msg.type == WSMsgType.BINARY:
                # Message from transport compression that isn't complete returns None
                data = self.decompress_msg(msg.data)
                if data:
                    await self.handle_data(data)
            elif msg.type == WSMsgType.ERROR:
                raise GatewayError from self.__socket.exception()

        # The loop is broken when the gateway stops receiving messages.
        # The "error" op codes are in `self.__close_codes`. The rest of the
        # close codes are unknown issues (such as a unintended disconnect) so the
        # client should reconnect to the gateway.
        err = self.__close_codes.get(self.__socket.close_code)

        if err:
            raise err

        _log.debug(
            "%s Disconnected from Gateway due without any errors. Reconnecting.",
            self.shard_key
        )
        self.__should_reconnect = True
        self.start_loop()

    async def handle_data(self, data: Dict[Any]):
        """Method is run when a payload is received from the gateway.
        The message is expected to already have been decompressed.
        Handling the opcode is forked to the background so they aren't blocking.
        """
        payload = GatewayDispatch.from_string(data)

        # Op code -1 is activated on all payloads
        op_negative_one = self.__dispatch_handlers.get(-1)
        if op_negative_one:
            ensure_future(op_negative_one(payload))

        _log.debug(
            "%s %s GatewayDispatch with opcode %s received",
            self.shard_key,
            datetime.now(),
            payload.op
        )

        # Many events are sent with a `null` sequence. This sequence should not
        # be tracked.
        if payload.seq is not None:
            self.__sequence_number = payload.seq
            _log.debug("%s Set sequence number to %s", self.shard_key, payload.seq)

        handler = self.__dispatch_handlers.get(payload.op)

        if handler is None:
            raise UnhandledException(f"Opcode {payload.op} does not have a handler")

        ensure_future(handler(payload))

    async def handle_heartbeat_req(self, payload: GatewayDispatch):

        """
        Opcode 1 - Instantly send a heartbeat.
        """
        self.send_next_heartbeat()

    async def handle_reconnect(self, payload: GatewayDispatch):

        """
        Opcode 7 - Reconnect and resume immediately.
        """
        _log.debug(
            "%s Requested to reconnect to Discord. Closing session and attempting to"
            " resume...",
            self.shard_key
        )

        await self.__socket.close(code=1000)
        self.__should_reconnect = True
        await self.start_loop()

    async def handle_invalid_session(self, payload: GatewayDispatch):

        """
        Opcode 9 - Invalid connection
        Attempt to relog. This is probably because the session was already invalidated
        when we tried to reconnect.
        """
        _log.debug("%s Invalid session, attempting to relog...", self.shard_key)
        self.__should_reconnect = False
        await self.start_loop()

    async def identify_and_handle_hello(self, payload: GatewayDispatch):
        """
        Opcode 10 - Hello there general kenobi
        Runs when we connect to the gateway for the first time and every time after.
        If the client thinks it should reconnect, the opcode 6 resume payload is sent
        instead of the opcode 2 hello payload. A new session is only started after a
        reconnect if pcode 9 is received.

        Successful reconnects are handled in the `resumed` middleware.
        """
        if self.__should_reconnect:
            _log.debug("%s Resuming connection with Discord", self.shard_key)

            await self.send(str(GatewayDispatch(
                6, {
                    "token": self.token,
                    "session_id": self.__session_id,
                    "seq": self.__sequence_number
                }
            )))
            return

        await self.send(str(
            GatewayDispatch(
                2, {
                    "token": self.token,
                    "intents": self.intents,
                    "properties": {
                        "$os": system(),
                        "$browser": __package__,
                        "$device": __package__
                    },
                    "compress": GatewayConfig.compressed(),
                    "shard": self.shard_key
                }
            )
        ))
        self.__heartbeat_interval = payload.data["heartbeat_interval"]

        # This process should already be forked to the background so there is no need to
        # `ensure_future()` here.
        self.start_heartbeat()

    async def handle_heartbeat(self, payload: GatewayDispatch):

        """
        Opcode 11 - Heatbeat
        Track that the heartbeat has been received using shared state (Rustaceans would
        be very mad)
        """
        self.__has_received_ack = True

    async def send(self, payload: str):
        """
        Send a string object to the payload. Most of this method is just logging,
        the last line is the only one that matters for functionality.
        """
        safe_payload = payload.replace(self.token, "%s..." % self.token[:10])

        if self.__session_id:
            safe_payload = safe_payload.replace(
                self.__session_id, "%s..." % self.__session_id[:4]
            )

        _log.debug(
            "%s Sending payload: %s",
            self.shard_key,
            safe_payload
        )

        if self.__socket.closed:
            _log.debug(
                "%s Socket is closing. Payload not sent.",
                self.shard_key
            )
            return

        await self.__socket.send_str(payload)

    def start_heartbeat(self):
        """
        Starts the heartbeat if it is not already running.
        """
        if not self.__heartbeat_task or self.__heartbeat_task.cancelled():
            self.__heartbeat_task = ensure_future(self.__heartbeat_loop())

    def stop_heartbeat(self):

        self.__heartbeat_task.cancel()

    def send_next_heartbeat(self):
        """
        It is expected to always be waiting for a hearbeat. By canceling that task,
        a heartbeat can be sent.
        """
        self.__wait_for_heartbeat.cancel()

    async def __heartbeat_loop(self):
        """
        The heartbeat is responsible for keeping the connection to Discord alive.

        Jitter is only random for the first heartbeat. It should be 1 every other
        heartbeat.
        """
        _log.debug("%s Starting heartbeat loop...", self.shard_key)

        # When waiting for first heartbeat, there hasn't been an ack received yet.
        # Set to true so the ack received check doesn't incorrectly fail.
        self.__has_received_ack = True

        for jitter in chain((random(),), repeat(1)):
            duration = self.__heartbeat_interval * jitter

            _log.debug(
                "%s %s sending heartbeat in %sms",
                self.shard_key, datetime.now(),
                duration
            )

            # Task is needed so waiting can be cancelled by op code 1
            self.__wait_for_heartbeat = create_task(
                sleep(duration / 1000)
            )

            await self.__wait_for_heartbeat

            if not self.__has_received_ack:
                # Close code is specified to be anything that is not 1000 in the docs.
                _log.debug(
                    "%s %s ack not received. Attempting to reconnect."
                    " Closing socket with close code 1001. \U0001f480",
                    datetime.now(),
                    self.shard_key
                )
                await self.__socket.close(code=1001)
                self.__should_reconnect = True
                # A new loop is started in the background while this one is stopped.
                ensure_future(self.start_loop())
                self.stop_heartbeat()
                return

            self.__has_received_ack = False
            await self.send(str(GatewayDispatch(1, data=self.__sequence_number)))
            _log.debug("%s sent heartbeat", self.shard_key)


1			# Copyright Pincer 2021-Present
			0 ignored issues – show introduced 2021-09-06 08:54 UTC by Report Bug Copy Issue Report Missing module docstring Loading history...
2			# Full MIT License can be found in `LICENSE` at the project root.
3
4
5			from __future__ import annotations
6
7			from asyncio import create_task, Task, ensure_future, sleep
8			from dataclasses import dataclass
9			from datetime import datetime
10			from itertools import repeat, count, chain
11			import logging
12			from platform import system
13			from random import random
14			from typing import TYPE_CHECKING, Any, Dict, Callable, Optional
15			from zlib import decompressobj
16
17			from aiohttp import (
18			ClientSession, WSMsgType, ClientConnectorError, ClientWebSocketResponse
19			)
20
21			from . import __package__
22			from ..utils.api_object import APIObject
23			from .._config import GatewayConfig
24			from ..core.dispatch import GatewayDispatch
25			from ..exceptions import (
26			InvalidTokenError, GatewayConnectionError, GatewayError, UnhandledException
27			)
28
29			if TYPE_CHECKING:
30			from ..objects.app.intents import Intents
31			Handler = Callable[[GatewayDispatch], None]
32
33			_log = logging.getLogger(__package__)
34
35			ZLIB_SUFFIX = b'\x00\x00\xff\xff'
36			inflator = decompressobj()
37
38
39			@dataclass
40			class SessionStartLimit(APIObject):
41			"""Session start limit info returned from the `gateway/bot` endpoint"""
42			total: int
43			remaining: int
44			reset_after: int
45			max_concurrency: int
46
47
48			@dataclass
49			class GatewayInfo(APIObject):
50			"""Gateway info returned from the `gateway/bot` endpoint"""
51			url: str
52			shards: int
53			session_start_limit: SessionStartLimit
54
55
56			class Gateway:
			0 ignored issues – show best-practice introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report Too many instance attributes (18/7) Loading history...
57			"""The Gateway handles all interactions with the Discord Websocket API.
58			This also contains the main event loop, and handles the heartbeat.
59
60			Running the Gateway will create a connection with the
61			Discord Websocket API on behalf of the provided token.
62
63			This token must be a bot token.
64			(Which can be found on
65			`<https://discord.com/developers/applications/>`_)
66
67			Parameters
68			----------
69			token : str.
70			The token for this bot
71			intents : :class:`~pincer.objects.app.intents.Intents`
72			The itents to use. More information can be found at
73			`<https://discord.com/developers/docs/topics/gateway#gateway-intents>`_.
74			url : str
75			The gateway url.
76			shard : int
77			The ID of the shard to run.
78			num_shards : int
79			Number used to route traffic to the current. This should usually be the total
80			number of shards that will be run. More information at
81			`<https://discord.com/developers/docs/topics/gateway#sharding>`_.
82			"""
83			def __init__(
84			self,
			0 ignored issues – show Coding Style introduced 2021-11-24 09:33 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
85			token: str, *,
			0 ignored issues – show Coding Style introduced 2021-11-24 09:33 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
86			intents: Intents,
			0 ignored issues – show Coding Style introduced 2021-12-27 02:25 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
87			url: str,
			0 ignored issues – show Coding Style introduced 2021-12-27 02:25 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
88			shard: int,
			0 ignored issues – show Coding Style introduced 2021-12-27 02:25 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
89			num_shards: int
			0 ignored issues – show Coding Style introduced 2021-12-27 02:25 UTC by Report Bug Copy Issue Report Wrong hanging indentation before block (add 4 spaces). Loading history...
90			) -> None:
91			if len(token) != 59:
92			raise InvalidTokenError(
93			"Discord Token must have exactly 59 characters."
94			)
95
96			self.token = token
97			self.intents = intents
98			self.url = url
99			self.shard = shard
100			self.num_shards = num_shards
101			self.shard_key = [shard, num_shards]
102
103			self.__dispatch_handlers: Dict[int, Handler] = {
104			1: self.handle_heartbeat_req,
105			7: self.handle_reconnect,
106			9: self.handle_invalid_session,
107			10: self.identify_and_handle_hello,
108			11: self.handle_heartbeat
109			}
110
111			# 4000 and 4009 are not included. The client will reconnect when receiving
112			# either.
113			self.__close_codes: Dict[int, GatewayError] = {
114			4001: GatewayError("Invalid opcode was sent"),
115			4002: GatewayError("Invalid payload was sent."),
116			4003: GatewayError("Payload was sent prior to identifying"),
117			4004: GatewayError("Token is not valid"),
118			4005: GatewayError(
119			"Authentication was sent after client already authenticated"
120			),
121			4007: GatewayError("Invalid sequence sent when starting new session"),
122			4008: GatewayError("Client was rate limited"),
123			4010: GatewayError("Invalid shard"),
124			4011: GatewayError("Sharding required"),
125			4012: GatewayError("Invalid API version"),
126			4013: GatewayError("Invalid intents"),
127			4014: GatewayError("Disallowed intents")
128			}
129
130			# ClientSession to be used for this Dispatcher
131			self.__session: Optional[ClientSession] = None
132
133			# This type `_WSRequestContextManager` isn't exposed by aiohttp.
134			# `ClientWebSocketResponse` is a parent class.
135			self.__socket: Optional[ClientWebSocketResponse] = None
136
137			# Buffer used to store information in transport conpression.
138			self.__buffer = bytearray()
139
140			# The gateway can be disconnected from Discord. This variable stores if the
141			# gateway should send a hello or reconnect.
142			self.__should_reconnect: bool = False
143
144			# The sequence number for the last received payload. This is used reconnecting.
145			self.__sequence_number: int = 0
146
147			# The heartbeat task
148			self.__heartbeat_task: Optional[Task] = None
149
150			# Keeps the Client waiting until the next heartbeat
151			self.__wait_for_heartbeat: Optional[Task] = None
152
153			# How long the client should wait between each Heartbeat.
154			self.__heartbeat_interval: Optional[int] = None
155
156			# Tracks whether the gateway has received an ack (opcode 11) since the last
157			# heartbeat.
158			# True: An ack has been received
159			# False: No ack has been received. Attempt to reconnect with gateway,
160			self.__has_received_ack: bool = True
161
162			# Session ID received from `on_ready` event. It is set in the `on_ready`
163			# middleware. This is used reconnecting.
164			self.__session_id: Optional[str] = None
165
166			def __del__(self):
167			"""Delete method ensures all connections are closed"""
168			if self.__socket:
169			create_task(self.__socket.close())
170			if self.__session:
171			create_task(self.__session.close())
172
173			async def init_session(self):
174			"""
175			Crates the ClientSession. ALWAYS run this function right after initializing
176			a Gateway.
177			"""
178			self.__session = ClientSession()
179
180			def append_handlers(self, handlers: Dict[int, Handler]):
181			"""The Client that uses the handler can append their own methods. The gateway
182			will run those methods when the specified opcode is received.
183			"""
184			self.__dispatch_handlers = handlers \| self.__dispatch_handlers
185
186			def set_session_id(self, _id: str):
187			"""Session id is private for consistency"""
188			self.__session_id = _id
189
190			def decompress_msg(self, msg: bytes) -> Optional[str]:
			0 ignored issues – show introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report Missing function or method docstring Loading history...
191			if GatewayConfig.compression == "zlib-payload":
192			return inflator.decompress(msg)
193
194			if GatewayConfig.compression == "zlib-stream":
195			self.__buffer.extend(msg)
196
197			if len(self.__buffer) < 4 or self.__buffer[-4:] != ZLIB_SUFFIX:
198			return None
199
200			msg = inflator.decompress(msg)
201			self.__buffer = bytearray()
202			return msg
203
204			return None
205
206			async def start_loop(self):
207			"""Instantiate the dispatcher, this will create a connection to the
208			Discord websocket API on behalf of the client whose token has
209			been passed.
210			"""
211			for _try in count():
212			try:
213			self.__socket = await self.__session.ws_connect(
214			GatewayConfig.make_uri(self.url)
215			)
216			break
217			except ClientConnectorError as e:
218			if _try > GatewayConfig.MAX_RETRIES:
219			raise GatewayConnectionError from e
220
221			_log.warning(
222			"%s Could not open websocket with Discord."
223			" Retrying in 15 seconds...",
224			self.shard_key
225			)
226			await sleep(15)
227
228			_log.debug("%s Starting envent loop...", self.shard_key)
229			await self.event_loop()
230
231			async def event_loop(self):
232			"""Handles receiving messages and decompressing them if needed"""
233			async for msg in self.__socket:
234			if msg.type == WSMsgType.TEXT:
235			await self.handle_data(msg.data)
236			elif msg.type == WSMsgType.BINARY:
237			# Message from transport compression that isn't complete returns None
238			data = self.decompress_msg(msg.data)
239			if data:
240			await self.handle_data(data)
241			elif msg.type == WSMsgType.ERROR:
242			raise GatewayError from self.__socket.exception()
243
244			# The loop is broken when the gateway stops receiving messages.
245			# The "error" op codes are in `self.__close_codes`. The rest of the
246			# close codes are unknown issues (such as a unintended disconnect) so the
247			# client should reconnect to the gateway.
248			err = self.__close_codes.get(self.__socket.close_code)
249
250			if err:
251			raise err
252
253			_log.debug(
254			"%s Disconnected from Gateway due without any errors. Reconnecting.",
255			self.shard_key
256			)
257			self.__should_reconnect = True
258			self.start_loop()
259
260			async def handle_data(self, data: Dict[Any]):
261			"""Method is run when a payload is received from the gateway.
262			The message is expected to already have been decompressed.
263			Handling the opcode is forked to the background so they aren't blocking.
264			"""
265			payload = GatewayDispatch.from_string(data)
266
267			# Op code -1 is activated on all payloads
268			op_negative_one = self.__dispatch_handlers.get(-1)
269			if op_negative_one:
270			ensure_future(op_negative_one(payload))
271
272			_log.debug(
273			"%s %s GatewayDispatch with opcode %s received",
274			self.shard_key,
275			datetime.now(),
276			payload.op
277			)
278
279			# Many events are sent with a `null` sequence. This sequence should not
280			# be tracked.
281			if payload.seq is not None:
282			self.__sequence_number = payload.seq
283			_log.debug("%s Set sequence number to %s", self.shard_key, payload.seq)
284
285			handler = self.__dispatch_handlers.get(payload.op)
286
287			if handler is None:
288			raise UnhandledException(f"Opcode {payload.op} does not have a handler")
289
290			ensure_future(handler(payload))
291
292			async def handle_heartbeat_req(self, payload: GatewayDispatch):
			0 ignored issues – show Unused Code introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report The argument `payload` seems to be unused. Loading history...
293			"""
294			Opcode 1 - Instantly send a heartbeat.
295			"""
296			self.send_next_heartbeat()
297
298			async def handle_reconnect(self, payload: GatewayDispatch):
			0 ignored issues – show Unused Code introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report The argument `payload` seems to be unused. Loading history...
299			"""
300			Opcode 7 - Reconnect and resume immediately.
301			"""
302			_log.debug(
303			"%s Requested to reconnect to Discord. Closing session and attempting to"
304			" resume...",
305			self.shard_key
306			)
307
308			await self.__socket.close(code=1000)
309			self.__should_reconnect = True
310			await self.start_loop()
311
312			async def handle_invalid_session(self, payload: GatewayDispatch):
			0 ignored issues – show Unused Code introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report The argument `payload` seems to be unused. Loading history...
313			"""
314			Opcode 9 - Invalid connection
315			Attempt to relog. This is probably because the session was already invalidated
316			when we tried to reconnect.
317			"""
318			_log.debug("%s Invalid session, attempting to relog...", self.shard_key)
319			self.__should_reconnect = False
320			await self.start_loop()
321
322			async def identify_and_handle_hello(self, payload: GatewayDispatch):
323			"""
324			Opcode 10 - Hello there general kenobi
325			Runs when we connect to the gateway for the first time and every time after.
326			If the client thinks it should reconnect, the opcode 6 resume payload is sent
327			instead of the opcode 2 hello payload. A new session is only started after a
328			reconnect if pcode 9 is received.
329
330			Successful reconnects are handled in the `resumed` middleware.
331			"""
332			if self.__should_reconnect:
333			_log.debug("%s Resuming connection with Discord", self.shard_key)
334
335			await self.send(str(GatewayDispatch(
336			6, {
337			"token": self.token,
338			"session_id": self.__session_id,
339			"seq": self.__sequence_number
340			}
341			)))
342			return
343
344			await self.send(str(
345			GatewayDispatch(
346			2, {
347			"token": self.token,
348			"intents": self.intents,
349			"properties": {
350			"$os": system(),
351			"$browser": __package__,
352			"$device": __package__
353			},
354			"compress": GatewayConfig.compressed(),
355			"shard": self.shard_key
356			}
357			)
358			))
359			self.__heartbeat_interval = payload.data["heartbeat_interval"]
360
361			# This process should already be forked to the background so there is no need to
362			# `ensure_future()` here.
363			self.start_heartbeat()
364
365			async def handle_heartbeat(self, payload: GatewayDispatch):
			0 ignored issues – show Unused Code introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report The argument `payload` seems to be unused. Loading history...
366			"""
367			Opcode 11 - Heatbeat
368			Track that the heartbeat has been received using shared state (Rustaceans would
369			be very mad)
370			"""
371			self.__has_received_ack = True
372
373			async def send(self, payload: str):
374			"""
375			Send a string object to the payload. Most of this method is just logging,
376			the last line is the only one that matters for functionality.
377			"""
378			safe_payload = payload.replace(self.token, "%s..." % self.token[:10])
379
380			if self.__session_id:
381			safe_payload = safe_payload.replace(
382			self.__session_id, "%s..." % self.__session_id[:4]
383			)
384
385			_log.debug(
386			"%s Sending payload: %s",
387			self.shard_key,
388			safe_payload
389			)
390
391			if self.__socket.closed:
392			_log.debug(
393			"%s Socket is closing. Payload not sent.",
394			self.shard_key
395			)
396			return
397
398			await self.__socket.send_str(payload)
399
400			def start_heartbeat(self):
401			"""
402			Starts the heartbeat if it is not already running.
403			"""
404			if not self.__heartbeat_task or self.__heartbeat_task.cancelled():
405			self.__heartbeat_task = ensure_future(self.__heartbeat_loop())
406
407			def stop_heartbeat(self):
			0 ignored issues – show introduced 2021-12-26 23:24 UTC by Report Bug Copy Issue Report Missing function or method docstring Loading history...
408			self.__heartbeat_task.cancel()
409
410			def send_next_heartbeat(self):
411			"""
412			It is expected to always be waiting for a hearbeat. By canceling that task,
413			a heartbeat can be sent.
414			"""
415			self.__wait_for_heartbeat.cancel()
416
417			async def __heartbeat_loop(self):
418			"""
419			The heartbeat is responsible for keeping the connection to Discord alive.
420
421			Jitter is only random for the first heartbeat. It should be 1 every other
422			heartbeat.
423			"""
424			_log.debug("%s Starting heartbeat loop...", self.shard_key)
425
426			# When waiting for first heartbeat, there hasn't been an ack received yet.
427			# Set to true so the ack received check doesn't incorrectly fail.
428			self.__has_received_ack = True
429
430			for jitter in chain((random(),), repeat(1)):
431			duration = self.__heartbeat_interval * jitter
432
433			_log.debug(
434			"%s %s sending heartbeat in %sms",
435			self.shard_key, datetime.now(),
436			duration
437			)
438
439			# Task is needed so waiting can be cancelled by op code 1
440			self.__wait_for_heartbeat = create_task(
441			sleep(duration / 1000)
442			)
443
444			await self.__wait_for_heartbeat
445
446			if not self.__has_received_ack:
447			# Close code is specified to be anything that is not 1000 in the docs.
448			_log.debug(
449			"%s %s ack not received. Attempting to reconnect."
450			" Closing socket with close code 1001. \U0001f480",
451			datetime.now(),
452			self.shard_key
453			)
454			await self.__socket.close(code=1001)
455			self.__should_reconnect = True
456			# A new loop is started in the background while this one is stopped.
457			ensure_future(self.start_loop())
458			self.stop_heartbeat()
459			return
460
461			self.__has_received_ack = False
462			await self.send(str(GatewayDispatch(1, data=self.__sequence_number)))
463			_log.debug("%s sent heartbeat", self.shard_key)
464

Pincer-org / Pincer

Pull Request — main (#329)

pincer.core.gateway.Gateway.start_loop() A

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like