From 8cd77bae2a93d06d76997694f5d734ed9e7251cc Mon Sep 17 00:00:00 2001 From: DevilXD Date: Fri, 8 Apr 2022 22:04:53 +0200 Subject: [PATCH] Change websocket internals to use aiohttp instead --- channel.py | 6 +-- exceptions.py | 17 +++++++++ main.py | 4 -- requirements.txt | 1 - twitch.py | 18 ++++----- utils.py | 76 +++++++++++++++++++++++++++++++++---- websocket.py | 99 +++++++++++++++++++++++++++++++----------------- 7 files changed, 160 insertions(+), 61 deletions(-) diff --git a/channel.py b/channel.py index 351c307..752bd66 100644 --- a/channel.py +++ b/channel.py @@ -1,14 +1,13 @@ from __future__ import annotations import re -import json import asyncio import logging from base64 import b64encode from functools import cached_property from typing import Any, SupportsInt, TYPE_CHECKING -from utils import Game, invalidate_cache +from utils import invalidate_cache, json_minify, Game from exceptions import MinerException, RequestException from constants import BASE_URL, GQL_OPERATIONS, ONLINE_DELAY, DROPS_ENABLED_TAG, URLType @@ -326,8 +325,7 @@ class Channel: } } ] - json_event = json.dumps(payload, separators=(",", ":")) - return {"data": (b64encode(json_event.encode("utf8"))).decode("utf8")} + return {"data": (b64encode(json_minify(payload).encode("utf8"))).decode("utf8")} async def send_watch(self) -> bool: """ diff --git a/exceptions.py b/exceptions.py index 08b9927..fb11e65 100644 --- a/exceptions.py +++ b/exceptions.py @@ -33,6 +33,23 @@ class RequestException(MinerException): super().__init__("Unknown error during request") +class WebsocketClosed(RequestException): + """ + Raised when the websocket connection has been closed. + + Attributes: + ----------- + received: bool + `True` if the closing was caused by our side receiving a close frame, `False` otherwise. + """ + def __init__(self, *args: object, received: bool = False): + if args: + super().__init__(*args) + else: + super().__init__("Websocket has been closed") + self.received: bool = received + + class LoginException(RequestException): """ Raised when an exception occurs during login phase. diff --git a/main.py b/main.py index a58d567..6a05715 100644 --- a/main.py +++ b/main.py @@ -17,10 +17,6 @@ try: import aiohttp # noqa except ModuleNotFoundError as exc: raise ImportError("You have to run 'pip install aiohttp' first") from exc -try: - import websockets # noqa -except ModuleNotFoundError as exc: - raise ImportError("You have to run 'pip install websockets' first") from exc try: import pystray # noqa except ModuleNotFoundError as exc: diff --git a/requirements.txt b/requirements.txt index b2280b9..b3cbb08 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,3 @@ aiohttp>2.0,<4.0 Pillow pystray -websockets>10.0 diff --git a/twitch.py b/twitch.py index 89d0fd8..920bda9 100644 --- a/twitch.py +++ b/twitch.py @@ -72,7 +72,9 @@ class Twitch: # Maintenance task self._mnt_task: asyncio.Task[None] | None = None - def initialize(self) -> None: + async def get_session(self) -> aiohttp.ClientSession: + if self._session is not None: + return self._session cookie_jar = aiohttp.CookieJar() if COOKIES_PATH.exists(): cookie_jar.load(COOKIES_PATH) @@ -81,6 +83,7 @@ class Twitch: headers={"User-Agent": USER_AGENT}, timeout=aiohttp.ClientTimeout(connect=5, total=10), ) + return self._session async def shutdown(self) -> None: start_time = time() @@ -135,7 +138,7 @@ class Twitch: def prevent_close(self): """ Called when the application window has to be prevented from closing, even after the user - closes it with X. Usually used solely to display tracebacks drom the closing sequence. + closes it with X. Usually used solely to display tracebacks from the closing sequence. """ self.gui.prevent_close() @@ -776,12 +779,10 @@ class Twitch: logger.debug("Checking login") login_form.update("Logging in...", None) # NOTE: We need this here because of the jar being accessed - if self._session is None: - self.initialize() - assert self._session is not None + session = await self.get_session() url = URL(BASE_URL) assert url.host is not None - jar = cast(aiohttp.CookieJar, self._session.cookie_jar) + jar = cast(aiohttp.CookieJar, session.cookie_jar) for attempt in range(2): cookie = jar.filter_cookies(url) if not cookie: @@ -823,10 +824,7 @@ class Twitch: async def request( self, method: str, url: str, *, attempts: int = 5, **kwargs ) -> abc.AsyncIterator[aiohttp.ClientResponse]: - session = self._session - if session is None: - self.initialize() - assert session is not None + session = await self.get_session() method = method.upper() cause: Exception | None = None for attempt in range(attempts): diff --git a/utils.py b/utils.py index 3bf6105..7d1cece 100644 --- a/utils.py +++ b/utils.py @@ -6,6 +6,7 @@ import random import string import asyncio import logging +import traceback from enum import Enum from pathlib import Path from functools import wraps @@ -32,11 +33,21 @@ _D = TypeVar("_D") # default _P = ParamSpec("_P") # params _JSON_T = TypeVar("_JSON_T", bound=Mapping[Any, Any]) logger = logging.getLogger("TwitchDrops") -NONCE_CHARS = string.ascii_letters + string.digits -serialize_env: dict[str, Callable[[Any], object]] = { - "set": set, - "datetime": lambda d: datetime.fromtimestamp(d, timezone.utc), -} + + +def format_traceback(exc: BaseException, **kwargs: Any) -> str: + """ + Like `traceback.print_exc` but returns a string. Uses the passed-in exception. + Any additional `**kwargs` are passed to the underlaying `traceback.format_exception`. + """ + return ''.join(traceback.format_exception(type(exc), exc, **kwargs)) + + +def json_minify(data: JsonType | list[JsonType]) -> str: + """ + Returns minified JSON for payload usage. + """ + return json.dumps(data, separators=(',', ':')) def resource_path(relative_path: Path | str) -> Path: @@ -58,6 +69,9 @@ def timestamp(string: str) -> datetime: return datetime.strptime(string, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc) +NONCE_CHARS = string.ascii_letters + string.digits + + def create_nonce(length: int = 30) -> str: return ''.join(random.choices(NONCE_CHARS, k=length)) @@ -110,6 +124,10 @@ def _serialize(obj: Any) -> Any: _MISSING = object() +SERIALIZE_ENV: dict[str, Callable[[Any], object]] = { + "set": set, + "datetime": lambda d: datetime.fromtimestamp(d, timezone.utc), +} def _remove_missing(obj: JsonType) -> JsonType: @@ -125,8 +143,8 @@ def _remove_missing(obj: JsonType) -> JsonType: def _deserialize(obj: JsonType) -> Any: if "__type" in obj: obj_type = obj["__type"] - if obj_type in serialize_env: - return serialize_env[obj_type](obj["data"]) + if obj_type in SERIALIZE_ENV: + return SERIALIZE_ENV[obj_type](obj["data"]) else: return _MISSING return obj @@ -145,6 +163,50 @@ def json_save(path: Path, contents: Mapping[Any, Any]) -> None: json.dump(contents, file, default=_serialize, sort_keys=True, indent=4) +class ExponentialBackoff: + def __init__( + self, + *, + base: float = 2, + variance: float | tuple[float, float] = 0.1, + shift: float = -1, + maximum: float = 300, + ): + if base <= 1: + raise ValueError("base has to be greater than 1") + self.exp: int = 0 + self.base: float = float(base) + self.shift: float = float(shift) + self.maximum: float = float(maximum) + self.variance_min: float + self.variance_max: float + if isinstance(variance, tuple): + self.variance_min, self.variance_max = variance + else: + self.variance_min = 1 - variance + self.variance_max = 1 + variance + + def __iter__(self) -> abc.Iterator[float]: + return self + + def __next__(self) -> float: + value: float = ( + pow(self.base, self.exp) + * random.uniform(self.variance_min, self.variance_max) + + self.shift + ) + if value > self.maximum: + return self.maximum + # NOTE: variance can cause the returned value to be lower than the previous one already, + # so this should be safe to move past the first return, + # to prevent the exponent from getting very big after reaching max and many iterations + self.exp += 1 + return value + + def reset(self) -> None: + self.exp = 0 + + class OrderedSet(MutableSet[_T]): """ Implementation of a set that preserves insertion order, diff --git a/websocket.py b/websocket.py index ea73c8d..540f177 100644 --- a/websocket.py +++ b/websocket.py @@ -7,32 +7,36 @@ from time import time from contextlib import suppress from typing import TYPE_CHECKING -from websockets.exceptions import ConnectionClosed, ConnectionClosedOK -from websockets.client import WebSocketClientProtocol, connect as websocket_connect +import aiohttp -from exceptions import MinerException -from utils import task_wrapper, create_nonce, AwaitableValue +from exceptions import MinerException, WebsocketClosed from constants import PING_INTERVAL, PING_TIMEOUT, MAX_WEBSOCKETS, WS_TOPICS_LIMIT +from utils import ( + task_wrapper, create_nonce, json_minify, format_traceback, AwaitableValue, ExponentialBackoff +) if TYPE_CHECKING: from collections import abc from twitch import Twitch + from gui import WebsocketStatus from constants import JsonType, WebsocketTopic +WSMsgType = aiohttp.WSMsgType logger = logging.getLogger("TwitchDrops") ws_logger = logging.getLogger("TwitchDrops.websocket") class Websocket: def __init__(self, pool: WebsocketPool, index: int): - self._pool = pool - self._twitch = pool._twitch + self._pool: WebsocketPool = pool + self._twitch: Twitch = pool._twitch + self._ws_gui: WebsocketStatus = self._twitch.gui.websockets # websocket index self._idx: int = index # current websocket connection - self._ws: AwaitableValue[WebSocketClientProtocol] = AwaitableValue() + self._ws: AwaitableValue[aiohttp.ClientWebSocketResponse] = AwaitableValue() # set when the websocket needs to reconnect self._reconnect_requested = asyncio.Event() # set when the topics changed @@ -61,7 +65,6 @@ class Websocket: ) def request_reconnect(self): - ws_logger.warning(f"Websocket[{self._idx}] requested reconnect.") # reset our ping interval, so we send a PING after reconnect right away self._next_ping = time() self._reconnect_requested.set() @@ -104,20 +107,37 @@ class Websocket: self._twitch.gui.websockets.remove(self._idx) asyncio.create_task(remover()) + async def _backoff_connect( + self, ws_url: str, **kwargs + ) -> abc.AsyncGenerator[aiohttp.ClientWebSocketResponse, None]: + session = await self._twitch.get_session() + backoff = ExponentialBackoff(**kwargs) + for delay in backoff: + try: + async with session.ws_connect(ws_url, ssl=True) as websocket: + backoff.reset() + yield websocket + except aiohttp.ClientConnectionError: + ws_logger.info( + f"Websocket[{self._idx}] connection error (sleep: {delay:.3}s)", exc_info=True + ) + await asyncio.sleep(delay) + @task_wrapper async def _handle(self): # ensure we're logged in before connecting + self.set_status("Initializing...") await self._twitch.wait_until_login() ws_logger.info(f"Websocket[{self._idx}] connecting...") self.set_status("Connecting...") # Connect/Reconnect loop - async for websocket in websocket_connect( - "wss://pubsub-edge.twitch.tv/v1", ssl=True, ping_interval=None + async for websocket in self._backoff_connect( + "wss://pubsub-edge.twitch.tv/v1", maximum=3*60 # 3 minutes maximum backoff time ): - # 3 minutes of max backoff - websocket.BACKOFF_MAX = 3 * 60 # type: ignore self._ws.set(websocket) self._reconnect_requested.clear() + # NOTE: _topics_changed doesn't start set, + # because there's no initial topics we can sub to right away self.set_status("Connected") try: try: @@ -128,26 +148,20 @@ class Websocket: finally: self._ws.clear() self._submitted.clear() - self._topics_changed.set() # lets the next WS connection resub to the topics + # set _topics_changed to let the next WS connection resub to the topics + self._topics_changed.set() # A reconnect was requested - except ConnectionClosed as exc: - if isinstance(exc, ConnectionClosedOK): - if exc.rcvd_then_sent: - # server closed the connection, not us - reconnect - ws_logger.warning(f"Websocket[{self._idx}] got disconnected.") - else: - # we closed it - exit - ws_logger.info(f"Websocket[{self._idx}] stopped.") - self.set_status("Disconnected") - return + except WebsocketClosed as exc: + if exc.received: + # server closed the connection, not us - reconnect + ws_logger.warning( + f"Websocket[{self._idx}] closed unexpectedly: {websocket.close_code}" + ) else: - if exc.rcvd is not None: - code = exc.rcvd.code - elif exc.sent is not None: - code = exc.sent.code - else: - code = -1 - ws_logger.warning(f"Websocket[{self._idx}] closed unexpectedly: {code}") + # we closed it - exit + ws_logger.info(f"Websocket[{self._idx}] stopped.") + self.set_status("Disconnected") + return except Exception: ws_logger.exception(f"Exception in Websocket[{self._idx}]") self.set_status("Reconnecting...") @@ -161,6 +175,7 @@ class Websocket: await self.send({"type": "PING"}) elif now >= self._max_pong: # it's been more than 10s and there was no PONG + ws_logger.warning(f"Websocket[{self._idx}] didn't receive a PONG, reconnecting...") self.request_reconnect() async def _handle_topics(self): @@ -208,10 +223,23 @@ class Websocket: ws = self._ws.get_with_default(None) assert ws is not None while True: - raw_message = await ws.recv() - message = json.loads(raw_message) - ws_logger.debug(f"Websocket[{self._idx}] received: {message}") - messages.append(message) + raw_message: aiohttp.WSMessage = await ws.receive() + ws_logger.debug(f"Websocket[{self._idx}] received: {raw_message}") + if raw_message.type is WSMsgType.TEXT: + message: JsonType = json.loads(raw_message.data) + messages.append(message) + continue # shortcut to avoid checking all other elifs if not necessary + elif raw_message.type is WSMsgType.CLOSE: + raise WebsocketClosed(received=True) + elif raw_message.type is WSMsgType.CLOSED: + raise WebsocketClosed(received=False) + elif raw_message.type is WSMsgType.CLOSING: + pass # skip these + elif raw_message.type is WSMsgType.ERROR: + logger.error(f"Websocket[{self._idx}] error: {format_traceback(raw_message.data)}") + raise WebsocketClosed() + else: + logger.error(f"Websocket[{self._idx}] error: Unknown message: {raw_message}") def _handle_message(self, message): # request the assigned topic to process the response @@ -241,6 +269,7 @@ class Websocket: pass elif msg_type == "RECONNECT": # We've received a reconnect request + ws_logger.warning(f"Websocket[{self._idx}] requested reconnect.") self.request_reconnect() else: ws_logger.warning(f"Websocket[{self._idx}] received unknown payload: {message}") @@ -271,7 +300,7 @@ class Websocket: assert ws is not None if message["type"] != "PING": message["nonce"] = create_nonce() - await ws.send(json.dumps(message, separators=(',', ':'))) + await ws.send_json(message, dumps=json_minify) ws_logger.debug(f"Websocket[{self._idx}] sent: {message}")