From 8e6db0829cba1d4d8a1053a13e9dbf5d16ba7e18 Mon Sep 17 00:00:00 2001 From: Caleb Johnson Date: Sat, 16 Feb 2019 17:22:55 -0500 Subject: [PATCH] [Audio] Connect to lavalink in the background (#2460) Also: - restart and reconnect if connection settings change - shutdown and restart if not configured to use external - show a message in [p]play et al. when the connection hasn't been made - move the JAR download to manager so audio.py can access it - only start if no process exists - bump red-lavalink to 0.2.3 Resolves #2306 --- Pipfile.lock | 26 +++++++--- redbot/cogs/audio/__init__.py | 28 +---------- redbot/cogs/audio/audio.py | 90 ++++++++++++++++++++++++++++------- redbot/cogs/audio/manager.py | 38 ++++++++++++++- setup.cfg | 2 +- 5 files changed, 131 insertions(+), 53 deletions(-) diff --git a/Pipfile.lock b/Pipfile.lock index 4e38150fb..d181d11fb 100644 --- a/Pipfile.lock +++ b/Pipfile.lock @@ -83,6 +83,13 @@ ], "version": "==0.4.1" }, + "distro": { + "hashes": [ + "sha256:362dde65d846d23baee4b5c058c8586f219b5a54be1cf5fc6ff55c4578392f57", + "sha256:eedf82a470ebe7d010f1872c17237c79ab04097948800029994fa458e52fb4b4" + ], + "version": "==1.4.0" + }, "dnspython": { "hashes": [ "sha256:36c5e8e38d4369a08b6780b7f27d790a292b2b08eea01607865bf0936c558e01", @@ -254,10 +261,10 @@ }, "red-lavalink": { "hashes": [ - "sha256:10a07b2b5736f52a0f5c3eeab3fbc3bf6a242ca6ee284a29ad79d6d1673ddfc3", - "sha256:9df0ddaa92d0d7294a4e236c4069765a0b8e5f258c18075dedf28f4b64a1aab5" + "sha256:13e1a3f91b990be9582cba039d9a32ec4cef760da1e7e6952143116ec83d4302", + "sha256:3dd0d73b4a908bbe9cfb703d2563dad1d1a58f8eea5896a0dacdf37d54a39d9c" ], - "version": "==0.2.1" + "version": "==0.2.3" }, "schema": { "hashes": [ @@ -421,6 +428,13 @@ ], "version": "==0.4.1" }, + "distro": { + "hashes": [ + "sha256:362dde65d846d23baee4b5c058c8586f219b5a54be1cf5fc6ff55c4578392f57", + "sha256:eedf82a470ebe7d010f1872c17237c79ab04097948800029994fa458e52fb4b4" + ], + "version": "==1.4.0" + }, "docutils": { "hashes": [ "sha256:02aec4bd92ab067f6ff27a38a38a41173bf01bed8f89157768c1573f53e474a6", @@ -664,10 +678,10 @@ }, "red-lavalink": { "hashes": [ - "sha256:10a07b2b5736f52a0f5c3eeab3fbc3bf6a242ca6ee284a29ad79d6d1673ddfc3", - "sha256:9df0ddaa92d0d7294a4e236c4069765a0b8e5f258c18075dedf28f4b64a1aab5" + "sha256:13e1a3f91b990be9582cba039d9a32ec4cef760da1e7e6952143116ec83d4302", + "sha256:3dd0d73b4a908bbe9cfb703d2563dad1d1a58f8eea5896a0dacdf37d54a39d9c" ], - "version": "==0.2.1" + "version": "==0.2.3" }, "requests": { "hashes": [ 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..6abfe1a97 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.")) @@ -1413,9 +1438,15 @@ class Audio(commands.Cog): await lavalink.connect(ctx.author.voice.channel) player = lavalink.get_player(ctx.guild.id) player.store("connect", datetime.datetime.utcnow()) + except IndexError: + await self._embed_msg( + ctx, _("Connection to Lavalink has not yet been established.") + ) + return False except AttributeError: await self._embed_msg(ctx, _("Connect to a voice channel first.")) return False + player = lavalink.get_player(ctx.guild.id) player.store("channel", ctx.channel.id) player.store("guild", ctx.guild.id) @@ -1793,6 +1824,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 +1912,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 +1936,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 +2371,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 +2384,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 +2406,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 +2424,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 +2440,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 +2456,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 +2604,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 +2708,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)) diff --git a/setup.cfg b/setup.cfg index e602a7e79..96ebc869d 100644 --- a/setup.cfg +++ b/setup.cfg @@ -40,7 +40,7 @@ install_requires = multidict==4.5.2 python-levenshtein-wheels==0.13.1 pyyaml==3.13 - red-lavalink==0.2.2 + red-lavalink==0.2.3 schema==0.6.8 websockets==7.0 yarl==1.3.0