Huge websocket redesign, to support more than a single connection

This commit is contained in:
DevilXD committed 2021-12-15 18:37:30 +01:00
1 parent 106831cc9a
commit a45838e088
7 files changed
+364 -230

No files matched your search

+10 -1
View File
@@ -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
},
]
}
+5 -5
View File
@@ -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"]
+39 -27
View File
@@ -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",
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",
# Using channel_id
"ChannelDrops": "channel-drop-events",
"ChannelStreamState": "stream-change-by-channel",
"ChannelCommunityPoints": "community-points-channel-v1",
},
"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",
},
}
+7 -7
View File
@@ -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"]
+8 -7
View File
@@ -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
+28 -25
View File
@@ -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(
+267 -158
View File
@@ -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}")
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:
}
}
)
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:
# no connection is made, so let it wait until there is one
return self.connected.wait()
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"]