Change websocket internals to use aiohttp instead

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:
         """

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.

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:

requirements.txt

@@ -1,4 +1,3 @@
 aiohttp>2.0,<4.0
 Pillow
 pystray
-websockets>10.0

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):

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,

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}")