diff --git a/cogs/listeners.py b/cogs/listeners.py index 4db2495..a5c4a4d 100644 --- a/cogs/listeners.py +++ b/cogs/listeners.py @@ -65,7 +65,7 @@ class Listeners(commands.Cog): async def on_voicelink_track_exception(self, player: voicelink.Player, track, error: dict): try: player._track_is_stuck = True - await player.context.send(f"{error['message']}! The next song will begin in the next 5 seconds.", delete_after=10) + await player.context.send(f"{error['message']} The next song will begin in the next 5 seconds.", delete_after=10) except: pass diff --git a/requirements.txt b/requirements.txt index 386fa6e..26abb7c 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,4 @@ -discord.py==2.4.0 +discord.py==2.5.0 motor==3.6.0 dnspython==2.2.1 tldextract==3.2.1 @@ -6,4 +6,4 @@ validators==0.18.2 humanize==4.0.0 beautifulsoup4==4.11.1 psutil==5.9.8 -aiohttp==3.9.5 \ No newline at end of file +aiohttp==3.11.12 \ No newline at end of file diff --git a/settings Example.json b/settings Example.json index 5a26452..875626c 100644 --- a/settings Example.json +++ b/settings Example.json @@ -12,7 +12,15 @@ "port": 2333, "password": "youshallnotpass", "secure": false, - "identifier": "DEFAULT" + "identifier": "DEFAULT", + "yt_ratelimit": { + "tokens": [], + "config": { + "retry_time": 10800, + "max_requests": 30 + }, + "strategy": "LoadBalance" + } } }, "prefix": "?", diff --git a/update.py b/update.py index efae7f2..5b84a2b 100644 --- a/update.py +++ b/update.py @@ -25,7 +25,7 @@ import requests, zipfile, os, shutil, argparse from io import BytesIO ROOT_DIR = os.path.dirname(os.path.abspath(__file__)) -__version__ = "v2.7.0b3" +__version__ = "v2.7.0b4" GITHUB_API_URL = "https://api.github.com/repos/ChocoMeow/Vocard/releases/latest" VOCARD_URL = "https://github.com/ChocoMeow/Vocard/archive/" diff --git a/voicelink/player.py b/voicelink/player.py index 8b6f51d..b17ff76 100644 --- a/voicelink/player.py +++ b/voicelink/player.py @@ -43,9 +43,10 @@ from discord import ( ) from discord.ext import commands + from . import events from .enums import SearchType, LoopType, RequestMethod -from .events import VoicelinkEvent, TrackEndEvent, TrackStartEvent +from .events import VoicelinkEvent, TrackEndEvent, TrackStartEvent, TrackExceptionEvent from .exceptions import VoicelinkException, FilterInvalidArgument, TrackInvalidPosition, TrackLoadError, FilterTagAlreadyInUse, DuplicateTrack from .filters import Filter, Filters from .objects import Track, Playlist @@ -350,6 +351,10 @@ class Player(VoiceProtocol): if isinstance(event, TrackEndEvent) and event.reason != "replaced": self._current = None + + if isinstance(event, TrackExceptionEvent) and event.exception["message"] == "This content isn’t available.": + if self._node.yt_ratelimit: + await self._node.yt_ratelimit.flag_active_token() event.dispatch(self._bot) @@ -435,16 +440,14 @@ class Player(VoiceProtocol): self.controller = await self.context.channel.send(embed=embed, view=view) elif not await self.is_position_fresh(): - try: - await self.controller.delete() - except Exception as e: - self._logger.warning( - f"Failed to delete outdated controller in {self.guild.name}({self.guild.id}): {e}" - ) + await self.controller.delete() self.controller = await self.context.channel.send(embed=embed, view=view) else: await self.controller.edit(embed=embed, view=view) + + except errors.Forbidden: + self._logger.warning(f"Missing permission to update the music controller on {self.guild.name}({self.guild.id})") except Exception as e: self._logger.error(f"Something went wrong while sending music controller to {self.guild.name}({self.guild.id})", exc_info=e) @@ -572,8 +575,10 @@ class Player(VoiceProtocol): if end or track.end_time: data["endTime"] = str(end if end else track.end_time) - + await self.send(method=RequestMethod.PATCH, query=f"noReplace={ignore_if_playing}", data=data) + if self._node.yt_ratelimit: + await self._node.yt_ratelimit.handle_request() self._current = track diff --git a/voicelink/pool.py b/voicelink/pool.py index f5bba50..593695e 100644 --- a/voicelink/pool.py +++ b/voicelink/pool.py @@ -52,6 +52,7 @@ from .exceptions import ( from .objects import Playlist, Track from .utils import ExponentialBackoff, NodeStats, NodeInfo, Ping from .enums import RequestMethod +from .ratelimit import YTRatelimit, YTToken, STRATEGY if TYPE_CHECKING: from .player import Player @@ -88,6 +89,7 @@ class Node: identifier: str, secure: bool = False, heartbeat: int = 30, + yt_ratelimit: dict = None, session: Optional[aiohttp.ClientSession] = None, spotify_client_id: Optional[str] = None, spotify_client_secret: Optional[str] = None, @@ -103,7 +105,7 @@ class Node: self._heartbeat: int = heartbeat self._secure: bool = secure self._logger: Optional[logging.Logger] = logger - + self._websocket_uri: str = f"{'wss' if self._secure else 'ws'}://{self._host}:{self._port}/" + NODE_VERSION + "/websocket" self._rest_uri: str = f"{'https' if self._secure else 'http'}://{self._host}:{self._port}" @@ -129,6 +131,8 @@ class Node: self._spotify_client_secret: Optional[str] = spotify_client_secret self._spotify_client: Optional[spotify.Client] = None + self.yt_ratelimit: Optional[YTRatelimit] = STRATEGY.get(yt_ratelimit.get("strategy"))(self, yt_ratelimit) if yt_ratelimit else None + self._bot.add_listener(self._update_handler, "on_socket_response") def __repr__(self): @@ -287,11 +291,15 @@ class Node: return await resp.json(content_type=None) return await resp.json() - + async def connect(self) -> Node: """Initiates a connection with a Lavalink node and adds it to the node pool.""" try: + if self._available: + self._logger.info(f"Node [{self._identifier}] already connected.") + return + self._websocket = await self._session.ws_connect( self._websocket_uri, headers=self._headers, heartbeat=self._heartbeat ) @@ -364,18 +372,8 @@ class Node: Context object on the track it builds. """ - async with self._session.get( - f"{self._rest_uri}/" + NODE_VERSION + "/decodetrack?", - headers={"Authorization": self._password}, - params={"track": identifier} - ) as resp: - if not resp.status == 200: - raise TrackLoadError( - f"Failed to build track. Check if the identifier is correct and try again." - ) - - data: dict = await resp.json() - return Track(track_id=identifier, info=data, requester=requester) + data = await self.send(RequestMethod.GET, f"decodetrack?encodedTrack={identifier}") + return Track(track_id=identifier, info=data, requester=requester) async def get_tracks( self, @@ -436,11 +434,7 @@ class Node: ) elif DISCORD_MP3_URL_REGEX.match(query): - async with self._session.get( - url=f"{self._rest_uri}/" + NODE_VERSION + f"/loadtracks?identifier={quote(query)}", - headers={"Authorization": self._password} - ) as response: - data: dict = await response.json() + data = await self.send(RequestMethod.GET, f"loadtracks?identifier={quote(query)}") try: track: dict = data["data"] @@ -455,11 +449,7 @@ class Node: ) ] else: - async with self._session.get( - url=f"{self._rest_uri}/" + NODE_VERSION + f"/loadtracks?identifier={quote(query)}", - headers={"Authorization": self._password} - ) as response: - data = await response.json() + data = await self.send(RequestMethod.GET, f"loadtracks?identifier={quote(query)}") load_type = data.get("loadType") @@ -556,7 +546,21 @@ class Node: tracks = tracks.tracks return tracks[:limit] if limit else tracks + + async def update_refresh_yt_access_token(self, token: YTToken) -> dict: + if not self._available: + raise NodeNotAvailable(f"The node '{self._identifier}' is unavailable.") + uri: str = f"{self._rest_uri}/youtube" + async with self._session.request( + method="POST", + url=uri, + headers={"Authorization": self._password}, + json={"refreshToken": token.token} + ) as resp: + if resp.status >= 300: + raise NodeException(f"Getting errors from Lavalink REST api") + class NodePool: """The base class for the node pool. This holds all the nodes that are to be used by the bot. @@ -634,6 +638,7 @@ class NodePool: identifier: str, secure: bool = False, heartbeat: int = 30, + yt_ratelimit: dict = None, spotify_client_id: Optional[str] = None, spotify_client_secret: Optional[str] = None, session: Optional[aiohttp.ClientSession] = None, @@ -651,8 +656,8 @@ class NodePool: node = Node( pool=cls, bot=bot, host=host, port=port, password=password, - identifier=identifier, secure=secure, heartbeat=heartbeat, spotify_client_id=spotify_client_id, - session=session, spotify_client_secret=spotify_client_secret, + identifier=identifier, secure=secure, heartbeat=heartbeat, yt_ratelimit=yt_ratelimit, + session=session, spotify_client_id=spotify_client_id, spotify_client_secret=spotify_client_secret, resume_key=resume_key, logger=logger ) diff --git a/voicelink/ratelimit.py b/voicelink/ratelimit.py new file mode 100644 index 0000000..171a512 --- /dev/null +++ b/voicelink/ratelimit.py @@ -0,0 +1,120 @@ +"""MIT License + +Copyright (c) 2023 - present Vocard Development + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. +""" + +import time +from abc import ABC, abstractmethod +from typing import List, Optional, Dict, TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .pool import Node + +class YTToken: + def __init__(self, token: str): + self.token: str = token + self.allow_retry_time: float = 0.0 + self.requested_times: int = 0 + self.is_flagged: bool = False + self.flagged_time: float = 0.0 + + @property + def allow_retry(self) -> bool: + """Determine if the token can be used again.""" + return time.time() >= self.allow_retry_time + +class YTRatelimit(ABC): + """ + Abstract base class for YouTube rate limit strategies. + """ + def __init__(self, node: "Node", tokens: List[str]) -> None: + self.node: "Node" = node + self.tokens: List[YTToken] = [YTToken(token) for token in tokens] + self.active_token: Optional[YTToken] = self.tokens[0] if self.tokens else None + + @abstractmethod + async def flag_active_token(self) -> None: + """ + Mark the current active token as flagged when a rate-limit is encountered. + """ + pass + + @abstractmethod + async def handle_request(self) -> None: + """ + Update usage count or perform any necessary pre-request operations. + """ + pass + + async def swap_token(self) -> Optional[YTToken]: + """ + Swap the active token with another token that is either not flagged or ready to retry. + If a new token is found, update it via the node and return it. + """ + for token in self.tokens: + if token != self.active_token and (not token.is_flagged or token.allow_retry): + try: + await self.node.update_refresh_yt_access_token(token) + self.active_token = token + return token + except Exception as e: + self.node._logger.error("Something wrong while updating the youtube access token.", exc_info=e) + + self.node._logger.warning("No active token available for processing the request.") + return None + +class LoadBalance(YTRatelimit): + """ + A rate limiting strategy that load balances requests across tokens. + """ + def __init__(self, node: "Node", config: Dict[str, Any]): + super().__init__(node, tokens=config.get("tokens", [])) + self._config: Dict[str, Any] = config.get("config", {}) + self._retry_time: int = self._config.get("retry_time", 10_800) + self._max_requests: int = self._config.get("max_requests", 30) + + async def flag_active_token(self) -> None: + """ + Flag the active token and set a delay (e.g., 3 hours) until it can be retried. + """ + if self.active_token: + self.active_token.is_flagged = True + self.active_token.flagged_time = time.time() + self.active_token.allow_retry_time = self.active_token.flagged_time + self._retry_time + await self.swap_token() + + async def handle_request(self) -> None: + """ + Increment the active token's usage counter and swap tokens if a threshold is reached. + """ + if not self.active_token: + return await self.swap_token() + + self.active_token.requested_times += 1 + if self.active_token.requested_times >= self._max_requests: + self.active_token.requested_times = 0 + swapped_token = await self.swap_token() + if swapped_token is None: + return self.node._logger.warning("No available token found after swapping.") + +STRATEGY = { + "LoadBalance": LoadBalance +} \ No newline at end of file