diff --git a/redbot/cogs/audio/__init__.py b/redbot/cogs/audio/__init__.py index 19711b903..3f912df3d 100644 --- a/redbot/cogs/audio/__init__.py +++ b/redbot/cogs/audio/__init__.py @@ -1,10 +1,8 @@ from pathlib import Path -from aiohttp import ClientSession -import shutil import logging from .audio import Audio -from .manager import start_lavalink_server +from .manager import start_lavalink_server, maybe_download_lavalink from redbot.core import commands from redbot.core.data_manager import cog_data_path import redbot.core @@ -22,30 +20,6 @@ APP_YML_FILE = LAVALINK_DOWNLOAD_DIR / "application.yml" BUNDLED_APP_YML_FILE = Path(__file__).parent / "data/application.yml" -async def download_lavalink(session): - with LAVALINK_JAR_FILE.open(mode="wb") as f: - async with session.get(LAVALINK_DOWNLOAD_URL) as resp: - while True: - chunk = await resp.content.read(512) - if not chunk: - break - f.write(chunk) - - -async def maybe_download_lavalink(loop, cog): - jar_exists = LAVALINK_JAR_FILE.exists() - current_build = redbot.VersionInfo.from_json(await cog.config.current_version()) - - if not jar_exists or current_build < redbot.core.version_info: - log.info("Downloading Lavalink.jar") - LAVALINK_DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) - async with ClientSession(loop=loop) as session: - await download_lavalink(session) - await cog.config.current_version.set(redbot.core.version_info.to_json()) - - shutil.copyfile(str(BUNDLED_APP_YML_FILE), str(APP_YML_FILE)) - - async def setup(bot: commands.Bot): cog = Audio(bot) if not await cog.config.use_external_lavalink(): diff --git a/redbot/cogs/audio/audio.py b/redbot/cogs/audio/audio.py index bc4471d39..a567ebf66 100644 --- a/redbot/cogs/audio/audio.py +++ b/redbot/cogs/audio/audio.py @@ -25,7 +25,7 @@ from redbot.core.utils.menus import ( ) from redbot.core.utils.predicates import MessagePredicate, ReactionPredicate from urllib.parse import urlparse -from .manager import shutdown_lavalink_server +from .manager import shutdown_lavalink_server, start_lavalink_server, maybe_download_lavalink _ = Translator("Audio", __file__) @@ -74,26 +74,47 @@ class Audio(commands.Cog): self.config.register_global(**default_global) self.skip_votes = {} self.session = aiohttp.ClientSession() + self._connect_task = None self._disconnect_task = None self._cleaned_up = False async def initialize(self): - host = await self.config.host() - password = await self.config.password() - rest_port = await self.config.rest_port() - ws_port = await self.config.ws_port() - - await lavalink.initialize( - bot=self.bot, - host=host, - password=password, - rest_port=rest_port, - ws_port=ws_port, - timeout=60, - ) + self._restart_connect() + self._disconnect_task = self.bot.loop.create_task(self.disconnect_timer()) lavalink.register_event_listener(self.event_handler) - self._disconnect_task = self.bot.loop.create_task(self.disconnect_timer()) + def _restart_connect(self): + if self._connect_task: + self._connect_task.cancel() + + self._connect_task = self.bot.loop.create_task(self.attempt_connect()) + + async def attempt_connect(self, timeout: int = 30): + while True: # run until success + external = await self.config.use_external_lavalink() + if not external: + shutdown_lavalink_server() + await maybe_download_lavalink(self.bot.loop, self) + await start_lavalink_server(self.bot.loop) + try: + host = await self.config.host() + password = await self.config.password() + rest_port = await self.config.rest_port() + ws_port = await self.config.ws_port() + + await lavalink.initialize( + bot=self.bot, + host=host, + password=password, + rest_port=rest_port, + ws_port=ws_port, + timeout=timeout, + ) + return # break infinite loop + except Exception: + if not external: + shutdown_lavalink_server() + await asyncio.sleep(1) # prevent busylooping async def event_handler(self, player, event_type, extra): notify = await self.config.guild(player.channel.guild).notify() @@ -903,6 +924,10 @@ class Audio(commands.Cog): player.store("connect", datetime.datetime.utcnow()) except AttributeError: return await self._embed_msg(ctx, _("Connect to a voice channel first.")) + except IndexError: + return await self._embed_msg( + ctx, _("Connection to Lavalink has not yet been established.") + ) if dj_enabled: if not await self._can_instaskip(ctx, ctx.author): return await self._embed_msg(ctx, _("You need the DJ role to queue tracks.")) @@ -1416,6 +1441,10 @@ class Audio(commands.Cog): except AttributeError: await self._embed_msg(ctx, _("Connect to a voice channel first.")) return False + except IndexError: + return await self._embed_msg( + ctx, _("Connection to Lavalink has not yet been established.") + ) player = lavalink.get_player(ctx.guild.id) player.store("channel", ctx.channel.id) player.store("guild", ctx.guild.id) @@ -1793,6 +1822,10 @@ class Audio(commands.Cog): player.store("connect", datetime.datetime.utcnow()) except AttributeError: return await self._embed_msg(ctx, _("Connect to a voice channel first.")) + except IndexError: + return await self._embed_msg( + ctx, _("Connection to Lavalink has not yet been established.") + ) player = lavalink.get_player(ctx.guild.id) shuffle = await self.config.guild(ctx.guild).shuffle() player.store("channel", ctx.channel.id) @@ -1877,6 +1910,10 @@ class Audio(commands.Cog): player.store("connect", datetime.datetime.utcnow()) except AttributeError: return await self._embed_msg(ctx, _("Connect to a voice channel first.")) + except IndexError: + return await self._embed_msg( + ctx, _("Connection to Lavalink has not yet been established.") + ) player = lavalink.get_player(ctx.guild.id) jukebox_price = await self.config.guild(ctx.guild).jukebox_price() shuffle = await self.config.guild(ctx.guild).shuffle() @@ -1897,7 +1934,6 @@ class Audio(commands.Cog): except IndexError: search_choice = tracks[-1] try: - search_check = search_choice.uri if "localtracks" in search_choice.uri: if search_choice.title == "Unknown title": description = "**{} - {}**\n{}".format( @@ -2333,6 +2369,7 @@ class Audio(commands.Cog): """Toggle using external lavalink servers.""" external = await self.config.use_external_lavalink() await self.config.use_external_lavalink.set(not external) + if external: await self.config.host.set("localhost") await self.config.password.set("youshallnotpass") @@ -2345,13 +2382,15 @@ class Audio(commands.Cog): ), ) embed.set_footer(text=_("Defaults reset.")) - return await ctx.send(embed=embed) + await ctx.send(embed=embed) else: await self._embed_msg( ctx, _("External lavalink server: {true_or_false}.").format(true_or_false=not external), ) + self._restart_connect() + @llsetup.command() async def host(self, ctx, host): """Set the lavalink server host.""" @@ -2365,6 +2404,8 @@ class Audio(commands.Cog): else: await self._embed_msg(ctx, _("Host set to {host}.").format(host=host)) + self._restart_connect() + @llsetup.command() async def password(self, ctx, password): """Set the lavalink server password.""" @@ -2381,6 +2422,8 @@ class Audio(commands.Cog): ctx, _("Server password set to {password}.").format(password=password) ) + self._restart_connect() + @llsetup.command() async def restport(self, ctx, rest_port: int): """Set the lavalink REST server port.""" @@ -2395,6 +2438,8 @@ class Audio(commands.Cog): else: await self._embed_msg(ctx, _("REST port set to {port}.").format(port=rest_port)) + self._restart_connect() + @llsetup.command() async def wsport(self, ctx, ws_port: int): """Set the lavalink websocket server port.""" @@ -2409,6 +2454,8 @@ class Audio(commands.Cog): else: await self._embed_msg(ctx, _("Websocket port set to {port}.").format(port=ws_port)) + self._restart_connect() + async def _check_external(self): external = await self.config.use_external_lavalink() if not external: @@ -2555,7 +2602,7 @@ class Audio(commands.Cog): try: query_url = urlparse(url) return all([query_url.scheme, query_url.netloc, query_url.path]) - except: + except Exception: return False @staticmethod @@ -2659,8 +2706,13 @@ class Audio(commands.Cog): def __unload(self): if not self._cleaned_up: self.session.detach() + if self._disconnect_task: self._disconnect_task.cancel() + + if self._connect_task: + self._connect_task.cancel() + lavalink.unregister_event_listener(self.event_handler) self.bot.loop.create_task(lavalink.close()) shutdown_lavalink_server() diff --git a/redbot/cogs/audio/manager.py b/redbot/cogs/audio/manager.py index b1ca4eef3..10cc7d4d9 100644 --- a/redbot/cogs/audio/manager.py +++ b/redbot/cogs/audio/manager.py @@ -8,6 +8,10 @@ import re from subprocess import Popen, DEVNULL from typing import Optional, Tuple +from aiohttp import ClientSession + +import redbot.core + _JavaVersion = Tuple[int, int] log = logging.getLogger("red.audio.manager") @@ -111,6 +115,10 @@ async def start_lavalink_server(loop): start_cmd = "java {} -jar {}".format(extra_flags, LAVALINK_JAR_FILE.resolve()) global proc + + if proc and proc.poll() is None: + return # already running + proc = Popen( shlex.split(start_cmd, posix=os.name == "posix"), cwd=str(LAVALINK_DOWNLOAD_DIR), @@ -126,11 +134,39 @@ async def start_lavalink_server(loop): def shutdown_lavalink_server(): - log.info("Shutting down lavalink server.") global shutdown shutdown = True global proc if proc is not None: + log.info("Shutting down lavalink server.") proc.terminate() proc.wait() proc = None + + +async def download_lavalink(session): + from . import LAVALINK_DOWNLOAD_URL, LAVALINK_JAR_FILE + + with LAVALINK_JAR_FILE.open(mode="wb") as f: + async with session.get(LAVALINK_DOWNLOAD_URL) as resp: + while True: + chunk = await resp.content.read(512) + if not chunk: + break + f.write(chunk) + + +async def maybe_download_lavalink(loop, cog): + from . import LAVALINK_DOWNLOAD_DIR, LAVALINK_JAR_FILE, BUNDLED_APP_YML_FILE, APP_YML_FILE + + jar_exists = LAVALINK_JAR_FILE.exists() + current_build = redbot.VersionInfo.from_json(await cog.config.current_version()) + + if not jar_exists or current_build < redbot.core.version_info: + log.info("Downloading Lavalink.jar") + LAVALINK_DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True) + async with ClientSession(loop=loop) as session: + await download_lavalink(session) + await cog.config.current_version.set(redbot.core.version_info.to_json()) + + shutil.copyfile(str(BUNDLED_APP_YML_FILE), str(APP_YML_FILE))