Add a separate timeout for Lavalink download (Cog-Creators#6461)

Co-authored-by: chovin <[email protected]>
Co-authored-by: Jakub Kuczys <[email protected]>

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 -

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)