diff --git a/.vscode/launch.json b/.vscode/launch.json index 2fbab57..a4eac0e 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -12,7 +12,7 @@ "console": "integratedTerminal" }, { - "name": "Run App", + "name": "Run App (INFO)", "type": "python", "request": "launch", "program": "main.py", @@ -20,5 +20,14 @@ "console": "integratedTerminal", "justMyCode": false }, + { + "name": "Run App (DEBUG)", + "type": "python", + "request": "launch", + "program": "main.py", + "args": ["-vvv"], + "console": "integratedTerminal", + "justMyCode": false + }, ] } \ No newline at end of file diff --git a/channel.py b/channel.py index 9cc882b..a0ffa5a 100644 --- a/channel.py +++ b/channel.py @@ -7,11 +7,11 @@ import logging from copy import copy from base64 import b64encode from datetime import datetime, timezone -from typing import Any, Optional, Dict, TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING from inventory import Game from exceptions import MinerException -from constants import BASE_URL, GQL_OPERATIONS, ONLINE_DELAY, DROPS_ENABLED_TAG +from constants import JsonType, BASE_URL, GQL_OPERATIONS, ONLINE_DELAY, DROPS_ENABLED_TAG if TYPE_CHECKING: from twitch import Twitch @@ -21,7 +21,7 @@ logger = logging.getLogger("TwitchDrops") class Stream: - def __init__(self, channel: Channel, data: Dict[str, Any]): + def __init__(self, channel: Channel, data: JsonType): self._twitch = channel._twitch self.channel = channel stream = data["stream"] @@ -38,7 +38,7 @@ class Stream: self._timestamp = datetime.now(timezone.utc) @classmethod - def from_directory(cls, channel: Channel, data: Dict[str, Any]): + def from_directory(cls, channel: Channel, data: JsonType): self = super().__new__(cls) self._twitch = channel._twitch self.channel = channel @@ -72,7 +72,7 @@ class Channel: await self.get_stream() @classmethod - def from_directory(cls, twitch: Twitch, data: Dict[str, Any]): + def from_directory(cls, twitch: Twitch, data: JsonType): self = super().__new__(cls) self._twitch = twitch channel = data["broadcaster"] diff --git a/constants.py b/constants.py index 4ffdca5..17ba587 100644 --- a/constants.py +++ b/constants.py @@ -1,10 +1,17 @@ from __future__ import annotations +import logging from copy import copy from datetime import timedelta from typing import Any, Optional, Union, Dict, Callable - +# Typing +JsonType = Dict[str, Any] +TopicProcess = Callable[[JsonType], Any] +# Values +DEBUG_RAW = logging.DEBUG - 1 +MAX_WEBSOCKETS = 8 +WS_TOPICS_LIMIT = 50 # URLs BASE_URL = "https://twitch.tv" AUTH_URL = "https://passport.twitch.tv" @@ -28,8 +35,8 @@ DROPS_ENABLED_TAG = "c2542d6d-cd10-4532-919b-3d19f30a768b" TERMINATED_STR = "Application Terminated.\nClose the console window to exit the application." -class GQLOperation(Dict[str, Any]): - def __init__(self, name: str, sha256: str, *, variables: Optional[Dict[str, Any]] = None): +class GQLOperation(JsonType): + def __init__(self, name: str, sha256: str, *, variables: Optional[JsonType] = None): super().__init__( operationName=name, extensions={ @@ -42,10 +49,10 @@ class GQLOperation(Dict[str, Any]): if variables is not None: self.__setitem__("variables", variables) - def with_variables(self, variables: Dict[str, Any]): + def with_variables(self, variables: JsonType): modified = copy(self) if "variables" in self: - existing_variables: Dict[str, Any] = modified["variables"] + existing_variables: JsonType = modified["variables"] existing_variables.update(variables) else: modified["variables"] = variables @@ -121,42 +128,47 @@ GQL_OPERATIONS: Dict[str, GQLOperation] = { } -def get_topic( - topic_name: str, target_id: Union[str, int], process: Callable[[Dict[str, Any]], Any] -) -> WebsocketTopic: - return WebsocketTopic(f"{WEBSOCKET_TOPICS[topic_name]}.{target_id}", process) - - class WebsocketTopic: - def __init__(self, topic_id: str, process: Callable[[Dict[str, Any]], Any]): - self.id: str = topic_id - self.process: Callable[[Dict[str, Any]], Any] = process + def __init__( + self, + topic_category: str, + topic_name: str, + target_id: Union[int, str], + process: TopicProcess, + ): + self._id: str = f"{WEBSOCKET_TOPICS[topic_category][topic_name]}.{target_id}" + self._process: TopicProcess = process + + def __call__(self, *args): + return self._process(*args) def __str__(self) -> str: - return self.id + return self._id def __eq__(self, other): if isinstance(other, str): - return self.id == other + return self._id == other elif isinstance(other, WebsocketTopic): - return self.id == other.id + return self._id == other._id return NotImplemented def __hash__(self) -> int: - return hash((self.__class__.__name__, self.id)) + return hash((self.__class__.__name__, self._id)) -WEBSOCKET_TOPICS: Dict[str, str] = { - # Using user_id - "UserDrops": "user-drop-events", - "UserStreamState": "stream-change-v1", - "UserCommunityPoints": "community-points-user-v1", - "Presence": "presence", - "Notifications": "onsite-notifications", - # Using channel_id - "ChannelDrops": "channel-drop-events", - "ChannelStreamState": "stream-change-by-channel", - "ChannelCommunityPoints": "community-points-channel-v1", - "VideoPlayback": "video-playback-by-id", - "StreamUpdate": "broadcast-settings-update", +WEBSOCKET_TOPICS: Dict[str, Dict[str, str]] = { + "User": { # Using user_id + "Drops": "user-drop-events", + "StreamState": "stream-change-v1", + "CommunityPoints": "community-points-user-v1", + "Presence": "presence", + "Notifications": "onsite-notifications", + }, + "Channel": { # Using channel_id + "Drops": "channel-drop-events", + "StreamState": "stream-change-by-channel", + "CommunityPoints": "community-points-channel-v1", + "VideoPlayback": "video-playback-by-id", + "StreamUpdate": "broadcast-settings-update", + }, } diff --git a/inventory.py b/inventory.py index 3dc33f5..43828e5 100644 --- a/inventory.py +++ b/inventory.py @@ -1,16 +1,16 @@ from __future__ import annotations from datetime import datetime -from typing import Any, Optional, List, Dict, TYPE_CHECKING +from typing import Optional, List, Dict, TYPE_CHECKING -from constants import GQL_OPERATIONS +from constants import JsonType, GQL_OPERATIONS if TYPE_CHECKING: from twitch import Twitch class Game: - def __init__(self, data: Dict[str, Any]): + def __init__(self, data: JsonType): self.id: int = int(data["id"]) self.name: str = data["name"] @@ -24,7 +24,7 @@ class Game: class BaseDrop: - def __init__(self, campaign: DropsCampaign, data: Dict[str, Any]): + def __init__(self, campaign: DropsCampaign, data: JsonType): self._twitch: Twitch = campaign._twitch self.id: str = data["id"] self.name: str = data["name"] @@ -80,7 +80,7 @@ class BaseDrop: class TimedDrop(BaseDrop): - def __init__(self, campaign: DropsCampaign, data: Dict[str, Any]): + def __init__(self, campaign: DropsCampaign, data: JsonType): super().__init__(campaign, data) self.current_minutes: int = data["self"]["currentMinutesWatched"] self.required_minutes: int = data["requiredMinutesWatched"] @@ -96,7 +96,7 @@ class TimedDrop(BaseDrop): def progress(self) -> float: return self.current_minutes / self.required_minutes - def update(self, message: Dict[str, Any]): + def update(self, message: JsonType): # {"type": "drop-progress", data: {"current_progress_min": 3, "required_progress_min": 10}} # {"type": "drop-claim", data: {"drop_instance_id": ...}} msg_type = message["type"] @@ -108,7 +108,7 @@ class TimedDrop(BaseDrop): class DropsCampaign: - def __init__(self, twitch: Twitch, data: Dict[str, Any]): + def __init__(self, twitch: Twitch, data: JsonType): self._twitch: Twitch = twitch self.id: str = data["id"] self.name: str = data["name"] diff --git a/main.py b/main.py index ef10d87..2de0641 100644 --- a/main.py +++ b/main.py @@ -10,10 +10,10 @@ import asyncio import warnings import traceback import threading -from typing import Any, Dict, NoReturn +from typing import NoReturn from twitch import Twitch -from constants import SETTINGS_PATH, TERMINATED_STR +from constants import JsonType, SETTINGS_PATH, TERMINATED_STR, DEBUG_RAW try: import win32api @@ -41,6 +41,8 @@ def terminate() -> NoReturn: raise RuntimeError("Uh oh") # this will never run, solely for MyPy +# extra debug logging level +logging.addLevelName(DEBUG_RAW, "DEBUG_RAW") # handle extra stackable '-v' parameter, that switches the logging level logging_level = logging.ERROR if len(sys.argv) > 1: @@ -51,6 +53,9 @@ if len(sys.argv) > 1: logging_level = logging.INFO elif arg in ("-vvv", "-v3"): logging_level = logging.DEBUG + elif arg in ("-vvvv", "-v4"): + # enables debugging raw messages + logging_level = DEBUG_RAW # handle logging stuff root_logger = logging.getLogger() root_logger.addHandler(logging.NullHandler()) @@ -62,7 +67,7 @@ logger.setLevel(logging_level) # handle settings try: with open(SETTINGS_PATH, 'r', encoding="utf8") as file: - settings: Dict[str, Any] = json.load(file) + settings: JsonType = json.load(file) except json.JSONDecodeError as exc: print(f"Error while reading the settings file:\n{str(exc)}") terminate() @@ -99,10 +104,6 @@ except (asyncio.CancelledError, KeyboardInterrupt): # Because we don't want anything from there to actually run during cleanup, # we need to explicitly cancel the task ourselves here. main_task.cancel() - # cancel all other tasks - for task in asyncio.all_tasks(loop): - if not task.done(): - task.cancel() # main_task was cancelled due to program shutting down - do the cleanup loop.run_until_complete(client.close()) # notify we're free to exit diff --git a/twitch.py b/twitch.py index 4a56371..4127901 100644 --- a/twitch.py +++ b/twitch.py @@ -14,11 +14,14 @@ except ImportError: raise ImportError("You have to run 'python -m pip install aiohttp' first") from channel import Channel -from websocket import Websocket +from websocket import WebsocketPool from inventory import DropsCampaign, Game from exceptions import LoginException, CaptchaRequired from constants import ( + JsonType, + WebsocketTopic, CLIENT_ID, + DEBUG_RAW, USER_AGENT, COOKIES_PATH, AUTH_URL, @@ -27,8 +30,6 @@ from constants import ( DROPS_ENABLED_TAG, TERMINATED_STR, GQLOperation, - WebsocketTopic, - get_topic, ) @@ -55,8 +56,6 @@ class Twitch: self._access_token: Optional[str] = None self._user_id: Optional[int] = None self._is_logged_in = asyncio.Event() - # Websocket - self.websocket = Websocket(self) # Storing, watching and changing channels self.channels: Dict[int, Channel] = {} self._watching_channel: Optional[Channel] = None @@ -65,6 +64,8 @@ class Twitch: # Inventory self.inventory: List[DropsCampaign] = [] self._campaign_change = asyncio.Event() + # Websocket + self.websocket = WebsocketPool(self) def wait_until_login(self): return self._is_logged_in.wait() @@ -77,7 +78,7 @@ class Twitch: self._session.cookie_jar.save(COOKIES_PATH) # type: ignore self.stop_watching() await self._session.close() - await self.websocket.close() + await self.websocket.stop() await asyncio.sleep(1) # allows aiohttp to safely close the session def is_currently_watching(self, channel: Channel) -> bool: @@ -111,7 +112,7 @@ class Twitch: print(f"No active campaigns to farm drops for.\n\n{TERMINATED_STR}") await asyncio.Future() # Start our websocket connection, only after we confirm that there are drops to mine - self.websocket.start() + await self.websocket.start() if not channel_names: # get a list of all channels with drops enabled print("Fetching suitable live channels to watch...") @@ -129,14 +130,16 @@ class Twitch: channel: Channel = await Channel(self, channel_name) # type: ignore self.channels[channel.id] = channel # Sub to these channel updates - topics: List[WebsocketTopic] = [] - for channel_id in self.channels: - topics.append( - get_topic( - "VideoPlayback", channel_id, partial(self.process_stream_state, channel_id) - ) + topics: List[WebsocketTopic] = [ + WebsocketTopic( + "Channel", + "VideoPlayback", + channel_id, + partial(self.process_stream_state, channel_id), ) - await self.websocket.add_topics(topics) + for channel_id in self.channels + ] + self.websocket.add_topics(topics) # Repeat: Change into a channel we can watch, then reset the flag self._channel_change.set() @@ -152,10 +155,10 @@ class Twitch: ) if self._campaign_change.is_set(): # we need to reevaluate all campaigns - # first of all, stop watching + # stop watching self.stop_watching() - # reconnect the websocket - await self.websocket.close() + # close the websocket + await self.websocket.stop() break # cycle the outer loop # otherwise, it was the channel change one for channel in self.channels.values(): @@ -197,8 +200,8 @@ class Twitch: if i == 0: # ensure every 30 minutes that we don't have unclaimed points bonus response = await self.gql_request(op) - channel_data: Dict[str, Any] = response["data"]["community"]["channel"] - claim_available: Dict[str, Any] = ( + channel_data: JsonType = response["data"]["community"]["channel"] + claim_available: JsonType = ( channel_data["self"]["communityPoints"]["availableClaim"] ) if claim_available: @@ -222,7 +225,7 @@ class Twitch: self._watching_task = None self._watching_channel = None - async def process_stream_state(self, channel_id: int, message: Dict[str, Any]): + async def process_stream_state(self, channel_id: int, message: JsonType): msg_type = message["type"] channel = self.channels.get(channel_id) if channel is None: @@ -308,7 +311,7 @@ class Twitch: print("\nNote: Password can be pasted in by pressing right click inside the window.\n") self.password = await self.get_password() - payload: Dict[str, Any] = { + payload: JsonType = { "username": self.username, "password": self.password, "client_id": CLIENT_ID, @@ -407,16 +410,16 @@ class Twitch: jar.update_cookies(cookie, URL("https://twitch.tv")) jar.save(COOKIES_PATH) - async def gql_request(self, op: GQLOperation) -> Dict[str, Any]: + async def gql_request(self, op: GQLOperation) -> JsonType: await self.check_login() headers = { "Authorization": f"OAuth {self._access_token}", "Client-Id": CLIENT_ID, } - logger.debug(f"GQL Request: {op}") + logger.log(DEBUG_RAW, f"GQL Request: {op}") async with self._session.post(GQL_URL, json=op, headers=headers) as response: response_json = await response.json() - logger.debug(f"GQL Response: {response_json}") + logger.log(DEBUG_RAW, f"GQL Response: {response_json}") return response_json async def get_inventory(self) -> List[DropsCampaign]: @@ -427,7 +430,7 @@ class Twitch: async def get_live_streams( self, games: Collection[Game], tag_ids: List[str] ) -> Dict[Game, List[Channel]]: - limit = int(45 / len(games)) + limit = 100 live_streams = {} for game in games: response = await self.gql_request( diff --git a/websocket.py b/websocket.py index 6091476..b29b414 100644 --- a/websocket.py +++ b/websocket.py @@ -5,22 +5,36 @@ import random import string import asyncio import logging +from time import time from functools import wraps -from collections import deque -from typing import Any, Optional, Union, Dict, Tuple, Set, Iterable, cast, TYPE_CHECKING +from contextlib import suppress +from typing import Any, Optional, List, Dict, Set, Iterable, TYPE_CHECKING from websockets.exceptions import ConnectionClosed, ConnectionClosedOK from websockets.client import WebSocketClientProtocol, connect as websocket_connect from inventory import TimedDrop from exceptions import MinerException -from constants import WEBSOCKET_URL, PING_INTERVAL, WebsocketTopic, get_topic +from constants import ( + JsonType, + WebsocketTopic, + DEBUG_RAW, + WEBSOCKET_URL, + PING_INTERVAL, + MAX_WEBSOCKETS, + WS_TOPICS_LIMIT, +) if TYPE_CHECKING: from twitch import Twitch logger = logging.getLogger("TwitchDrops") +NONCE_CHARS = string.ascii_letters + string.digits + + +def create_nonce(length: int = 30) -> str: + return ''.join(random.choices(NONCE_CHARS, k=length)) def task_wrapper(afunc): @@ -30,116 +44,95 @@ def task_wrapper(afunc): await afunc(self, *args, **kwargs) except Exception: logger.exception("Exception in websocket task") + raise # raise up to the wrapping task return wrapper class Websocket: - def __init__(self, twitch: Twitch): - self._twitch = twitch - self._ws: Optional[WebSocketClientProtocol] = None - self.connected = asyncio.Event() # set when there's an active websocket connection - self.reconnect = asyncio.Event() # set when the websocket needs to reconnect - self._send_queue: deque[Tuple[str, Dict[str, Any]]] = deque() - self._recv_dict: Dict[str, asyncio.Future[Any]] = {} - self._topics: Set[WebsocketTopic] = set() - self._ping_task: Optional[asyncio.Task[Any]] = None - self._connect_task: Optional[asyncio.Task[Any]] = None + def __init__(self, pool: WebsocketPool, index: int): + self._pool = pool + self._twitch = pool._twitch + # websocket index + self._idx: int = index + # current websocket connection + self._ws: WebSocketClientProtocol + # set when there's an active websocket connection + self._connected_flag = asyncio.Event() + # set when the websocket needs to reconnect + self._reconnect_requested = asyncio.Event() + # set when the topics changed + self._topics_changed = asyncio.Event() + # ping timestamps + self._next_ping = time() + self._max_pong = time() + # main task, responsible for receiving messages, sending them, and websocket ping + self._handle_task: Optional[asyncio.Task[Any]] = None + # topics stuff + self.topics: Dict[str, WebsocketTopic] = {} + self._submitted: Set[WebsocketTopic] = set() - async def _ping_loop(self): - await self.connected.wait() - ping_every = PING_INTERVAL.total_seconds() - while self.connected.is_set(): - try: - await asyncio.wait_for(self.send({"type": "PING"}), timeout=10) - except asyncio.TimeoutError: - # per documentation, if there's no response for a PING, reconnect to the websocket - logger.warning("Websocket got no response to PING - reconnect") - self.reconnect.set() - break - await asyncio.sleep(ping_every) + @property + def connected(self) -> bool: + return self._connected_flag.is_set() - def change_connection_state(self, state: bool): - if state: - # websocket is considered connected - logger.info("Websocket Connected") - self.connected.set() - self._ping_task = asyncio.create_task(self._ping_loop()) + def wait_until_connected(self): + return self._connected_flag.wait() + + def request_reconnect(self): + self._reconnect_requested.set() + + async def start(self): + if self.connected: + return + if self._handle_task is None: + self._handle_task = asyncio.create_task(self._handle()) else: - # websocket is considered disconnected - self.connected.clear() - if self._ping_task is not None: - self._ping_task.cancel() - self._ping_task = None + self.request_reconnect() + await self.wait_until_connected() - def start(self): - self._connect_task = asyncio.create_task(self.connect()) + async def stop(self): + if self._ws is not None: + await self._ws.close() + if self._handle_task is not None: + await self._handle_task + self._handle_task = None - def stop(self): - self.change_connection_state(False) - if self._connect_task is not None: - self._connect_task.cancel() - self._connect_task = None + def start_nowait(self): + if self.connected: + return + if self._handle_task is None: + self._handle_task = asyncio.create_task(self._handle()) + else: + self.request_reconnect() + + def stop_nowait(self): + if self._ws is not None: + asyncio.create_task(self._ws.close()) + self._handle_task = None @task_wrapper - async def connect(self): + async def _handle(self): # ensure we're logged in before connecting await self._twitch.wait_until_login() logger.info("Connecting to Websocket") - # Listen to our events of choice - user_id = cast(int, self._twitch._user_id) - # Add default topics - self.add_topics([ - get_topic("UserDrops", user_id, self.process_drops), - get_topic("UserCommunityPoints", user_id, self.process_points), - ]) # Connect/Reconnect loop async for websocket in websocket_connect(WEBSOCKET_URL, ssl=True, ping_interval=None): websocket.BACKOFF_MAX = 3 * 60 # type: ignore # 3 minutes self._ws = websocket - self.reconnect.clear() - self.change_connection_state(True) - # Send all our chosen topics - topics_list = list(map(str, self._topics)) - logger.debug(f"Listening for: {', '.join(topics_list)}") - self.send( - { - "type": "LISTEN", - "data": { - "topics": topics_list, - "auth_token": self._twitch._access_token, - } - } - ) try: - while not self.reconnect.is_set(): - # Process receive - try: - # Wait up to 0.5s for a message we're supposed to receive - raw_message = await asyncio.wait_for(websocket.recv(), timeout=0.5) - except asyncio.TimeoutError: - # nothing - skip handling - pass - else: - # we've got something to process - # separate method solely because the indent was getting rather ridiculus - await self.process_message(raw_message) - - # Early exit if needed - if self.reconnect.is_set(): - break - - # Process send - while self._send_queue: - nonce, message = self._send_queue.popleft() - if nonce != "PING": - message["nonce"] = nonce - await websocket.send(json.dumps(message, separators=(',', ':'))) - logger.debug(f"Websocket sent: {message}") + try: + self._reconnect_requested.clear() + self._connected_flag.set() + while not self._reconnect_requested.is_set(): + await self._handle_ping() + await self._handle_topics() + await self._handle_recv() + finally: + self._submitted.clear() + self._connected_flag.clear() # A reconnect was requested - self.change_connection_state(False) continue except ConnectionClosed as exc: - self.change_connection_state(False) if isinstance(exc, ConnectionClosedOK): if exc.rcvd_then_sent: # server closed the connection, not us - reconnect @@ -154,72 +147,42 @@ class Websocket: logger.exception("Exception in Websocket - Reconnecting") continue - async def process_message(self, raw_message: Union[bytes, str]): - message = json.loads(raw_message) - logger.debug(f"Websocket received: {message}") - msg_type = message["type"] - # handle the simple PING case - if msg_type == "PONG": - ping_future = self._recv_dict.pop("PING", None) - if ping_future is not None and not ping_future.done(): - ping_future.set_result(message) - elif msg_type == "RESPONSE": - try: - self._recv_dict.pop(message["nonce"]).set_result(message) - except KeyError: - logger.exception("Received response for a request we didn't send") - elif msg_type == "RECONNECT": - # We've received a reconnect request - logger.warning("Received a Websocket Reconnect Request") - self.reconnect.set() - elif msg_type == "MESSAGE": - # request the assigned topic to process the response - target_topic = message["data"]["topic"] - for topic in self._topics: - if target_topic == topic: - # use a task to not block the websocket - asyncio.create_task(topic.process(json.loads(message["data"]["message"]))) - break - else: - logger.error(f"Received unknown websocket payload: {message}") + async def _handle_ping(self): + now = time() + if now >= self._next_ping: + self._next_ping = now + PING_INTERVAL.total_seconds() + self._max_pong = now + 10 # wait for a PONG for up to 10s + await self.send({"type": "PING"}) + elif now >= self._max_pong: + # it's been more than 10s and there was no PONG + self.request_reconnect() - async def close(self): - self.stop() - if self._ws is not None: - await self._ws.close() - - def create_nonce(self, length: int = 30) -> str: - available_chars = string.ascii_letters + string.digits - return ''.join(random.choices(available_chars, k=length)) - - def send(self, message: Dict[str, Any]) -> asyncio.Future[Dict[str, Any]]: - logger.debug(f"Websocket sending: {message}") - msg_type = message["type"] - if msg_type == "PING": - nonce = "PING" - else: - nonce = self.create_nonce() - self._send_queue.append((nonce, message)) - future: asyncio.Future[Dict[str, Any]] = asyncio.get_running_loop().create_future() - self._recv_dict[nonce] = future - return future - - def add_topics(self, topics: Iterable[WebsocketTopic]): - # ensure no topics end up duplicated - topics = set(topics) - topics.difference_update(self._topics) - if not topics: - # none left to add + async def _handle_topics(self): + if not self._topics_changed.is_set(): + # nothing to do return - self._topics.update(topics) - if len(self._topics) >= 50: - # TODO: Handle multiple connections (up to 10) since one allows only up to 50 topics - raise MinerException("Too many topics") - if self.connected.is_set(): - # we're already connected, so we have to send the topics list ourselves - topics_list = list(map(str, topics)) - logger.debug(f"Listening for: {', '.join(topics_list)}") - return self.send( + current: Set[WebsocketTopic] = set(self.topics.values()) + # handle removed topics + removed = self._submitted.difference(current) + if removed: + topics_list = list(map(str, removed)) + logger.debug(f"Websocket[{self._idx}]: Removing topics: {', '.join(topics_list)}") + await self.send( + { + "type": "UNLISTEN", + "data": { + "topics": topics_list, + "auth_token": self._twitch._access_token, + } + } + ) + self._submitted.difference_update(removed) + # handle added topics + added = current.difference(self._submitted) + if added: + topics_list = list(map(str, added)) + logger.debug(f"Websocket[{self._idx}]: Adding topics: {', '.join(topics_list)}") + await self.send( { "type": "LISTEN", "data": { @@ -228,12 +191,158 @@ class Websocket: } } ) - else: - # no connection is made, so let it wait until there is one - return self.connected.wait() + self._submitted.update(added) + self._topics_changed.clear() + + async def _gather_recv(self, messages: List[JsonType]): + """ + Gather incoming messages over the timeout specified. + Note that there's no return value - this modifies `messages` in-place. + """ + while True: + raw_message = await self._ws.recv() + message = json.loads(raw_message) + logger.log(DEBUG_RAW, f"Websocket received: {message}") + messages.append(message) + + def _handle_message(self, message): + # request the assigned topic to process the response + topic_process = self.topics.get(message["data"]["topic"]) + if topic_process is not None: + # use a task to not block the websocket + asyncio.create_task(topic_process(json.loads(message["data"]["message"]))) + + async def _handle_recv(self): + """ + Handle receiving messages from the websocket. + """ + # listen over 0.5s for incoming messages + messages: List[JsonType] = [] + with suppress(asyncio.TimeoutError): + await asyncio.wait_for(self._gather_recv(messages), timeout=0.5) + # process them + for message in messages: + msg_type = message["type"] + if msg_type == "MESSAGE": + self._handle_message(message) + elif msg_type == "PONG": + # move the timestamp to something much later + self._max_pong = self._next_ping + elif msg_type == "RESPONSE": + # no special handling for these (for now) + pass + elif msg_type == "RECONNECT": + # We've received a reconnect request + logger.warning("Received a Websocket Reconnect Request") + self.request_reconnect() + else: + logger.error(f"Received unknown websocket payload: {message}") + + def add_topics(self, topics_set: Set[WebsocketTopic]): + while topics_set and len(self.topics) < WS_TOPICS_LIMIT: + topic = topics_set.pop() + self.topics[str(topic)] = topic + self._topics_changed.set() + + def remove_topics(self, topics_set: Set[WebsocketTopic]): + existing = topics_set.intersection(self.topics.values()) + if not existing: + # nothing to remove from here + return + topics_set.difference_update(existing) + for topic in existing: + del self.topics[str(topic)] + self._topics_changed.set() + + async def send(self, message: JsonType): + if self._ws is None: + return + if message["type"] != "PING": + message["nonce"] = create_nonce() + await self._ws.send(json.dumps(message, separators=(',', ':'))) + logger.log(DEBUG_RAW, f"Websocket sent: {message}") + + +class WebsocketPool: + def __init__(self, twitch: Twitch): + self._twitch: Twitch = twitch + self._running = asyncio.Event() + self.websockets: List[Websocket] = [] + + @property + def running(self) -> bool: + return self._running.is_set() + + def wait_until_connected(self): + return self._running.wait() + + async def start(self): + await self._twitch.wait_until_login() + self._running.set() + # Add default topics + assert self._twitch._user_id is not None + user_id = self._twitch._user_id + self.add_topics([ + WebsocketTopic("User", "Drops", user_id, self.process_drops), + WebsocketTopic("User", "CommunityPoints", user_id, self.process_points), + ]) + await asyncio.gather(*(ws.start() for ws in self.websockets)) + + async def stop(self): + self._running.clear() + await asyncio.gather(*(ws.stop() for ws in self.websockets)) + + def add_topics(self, topics: Iterable[WebsocketTopic]): + # ensure no topics end up duplicated + topics_set = set(topics) + if not topics_set: + # nothing to add + return + topics_set.difference_update(ws.topics.values() for ws in self.websockets) + if not topics_set: + # none left to add + return + for ws_idx in range(MAX_WEBSOCKETS): + if ws_idx < len(self.websockets): + # just read it back + ws = self.websockets[ws_idx] + else: + # create new + ws = Websocket(self, len(self.websockets)) + if self._running: + ws.start_nowait() + self.websockets.append(ws) + # ask websocket to take any topics it can - this modifies the set in-place + ws.add_topics(topics_set) + # see if there's any leftover topics for the next websocket connection + if not topics_set: + return + # if we're here, there were leftover topics after filling up all websockets + raise MinerException("Maximum topics limit has been reached") + + def remove_topics(self, topics: Iterable[WebsocketTopic]): + topics_set = set(topics) + if not topics_set: + # nothing to remove + return + for ws in self.websockets: + ws.remove_topics(topics_set) + # count up all the topics - if we happen to have more websockets connected than needed, + # stop the last one and recycle topics from it - repeat until we have enough + topics = [] + while True: + count = sum(len(ws.topics) for ws in self.websockets) + if count <= (len(self.websockets) - 1) * WS_TOPICS_LIMIT: + ws = self.websockets.pop() + topics.extend(ws.topics.values()) + ws.stop_nowait() + else: + break + if topics: + self.add_topics(topics) @task_wrapper - async def process_drops(self, message: Dict[str, Any]): + async def process_drops(self, message: JsonType): drop_id = message["data"]["drop_id"] drop: Optional[TimedDrop] = None for campaign in self._twitch.inventory: @@ -261,7 +370,7 @@ class Websocket: self._twitch.reevaluate_campaigns() @task_wrapper - async def process_points(self, message: Dict[str, Any]): + async def process_points(self, message: JsonType): msg_type = message["type"] if msg_type == "points-earned": points = message["data"]["point_gain"]["total_points"]