From 1f489190055ef4fac41443b1d113f069f87369af Mon Sep 17 00:00:00 2001 From: Chovin <17536161+Chovin@users.noreply.github.com> Date: Tue, 24 Dec 2024 13:08:49 +1000 Subject: [PATCH] Add a separate timeout for Lavalink download (#6461) Co-authored-by: chovin Co-authored-by: Jakub Kuczys --- redbot/cogs/audio/core/tasks/lavalink.py | 4 +++- redbot/cogs/audio/manager.py | 21 +++++++++++++++++++-- 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/redbot/cogs/audio/core/tasks/lavalink.py b/redbot/cogs/audio/core/tasks/lavalink.py index b952cc130c7..b94f36eed2a 100644 --- a/redbot/cogs/audio/core/tasks/lavalink.py +++ b/redbot/cogs/audio/core/tasks/lavalink.py @@ -56,7 +56,9 @@ async def lavalink_attempt_connect(self, timeout: int = 50, manual: bool = False password = configs["yaml"]["lavalink"]["server"]["password"] secured = False # Make this timeout customizable for lower powered machines? - self.managed_node_controller = ServerManager(self.config, timeout=60, cog=self) + self.managed_node_controller = ServerManager( + self.config, timeout=60, download_timeout=60 * 3, cog=self + ) try: await self.managed_node_controller.start(java_exec) # timeout is the same as ServerManager.timeout - diff --git a/redbot/cogs/audio/manager.py b/redbot/cogs/audio/manager.py index 93ad89e3d6a..4e750da26b7 100644 --- a/redbot/cogs/audio/manager.py +++ b/redbot/cogs/audio/manager.py @@ -125,13 +125,21 @@ class ServerManager: _buildtime: ClassVar[Optional[str]] = None _java_exc: ClassVar[str] = "java" - def __init__(self, config: Config, cog: "Audio", timeout: Optional[int] = None) -> None: + def __init__( + self, + config: Config, + cog: "Audio", + timeout: Optional[int] = None, + download_timeout: Optional[int] = None, + ) -> None: self.ready: asyncio.Event = asyncio.Event() + self.downloaded: asyncio.Event = asyncio.Event() self._config = config self._proc: Optional[asyncio.subprocess.Process] = None # pylint:disable=no-member self._shutdown: bool = False self.start_monitor_task = None self.timeout = timeout + self.download_timeout = download_timeout self.cog = cog self._args = [] self._pipe_task = None @@ -357,6 +365,7 @@ async def shutdown(self) -> None: async def _partial_shutdown(self) -> None: self.ready.clear() + self.downloaded.clear() if self._shutdown is True: # For convenience, calling this method more than once or calling it before starting it # does nothing. @@ -414,6 +423,7 @@ async def _download_jar(self) -> None: log.info("Successfully downloaded Lavalink.jar (%s bytes written)", format(nbytes, ",")) await self._is_up_to_date() + self.downloaded.set() async def _is_up_to_date(self): if self._up_to_date is True: @@ -487,8 +497,14 @@ async def maybe_download_jar(self): managed_node.JAR_VERSION, ) await self._download_jar() + else: + self.downloaded.set() - async def wait_until_ready(self, timeout: Optional[float] = None): + async def wait_until_ready( + self, timeout: Optional[float] = None, download_timeout: Optional[float] = None + ): + download_timeout = download_timeout or self.download_timeout + await asyncio.wait_for(self.downloaded.wait(), timeout=download_timeout) await asyncio.wait_for(self.ready.wait(), timeout=timeout or self.timeout) async def start_monitor(self, java_path: str): @@ -499,6 +515,7 @@ async def start_monitor(self, java_path: str): self._shutdown = False if self._proc is None or self._proc.returncode is not None: self.ready.clear() + self.downloaded.clear() await self._start(java_path=java_path) while True: await self.wait_until_ready(timeout=self.timeout)