Red-DiscordBot/redbot/cogs/audio/manager.py

import itertools
import pathlib
import platform
import shutil
import asyncio
import asyncio.subprocess
import logging
import re
import tempfile
from typing import Optional, Tuple, ClassVar, List

import aiohttp

from redbot.core import data_manager

JAR_VERSION = "3.2.0.3"
JAR_BUILD = 796
LAVALINK_DOWNLOAD_URL = (
    f"https://github.com/Cog-Creators/Lavalink-Jars/releases/download/{JAR_VERSION}_{JAR_BUILD}/"
    f"Lavalink.jar"
)
LAVALINK_DOWNLOAD_DIR = data_manager.cog_data_path(raw_name="Audio")
LAVALINK_JAR_FILE = LAVALINK_DOWNLOAD_DIR / "Lavalink.jar"

BUNDLED_APP_YML = pathlib.Path(__file__).parent / "data" / "application.yml"
LAVALINK_APP_YML = LAVALINK_DOWNLOAD_DIR / "application.yml"

READY_LINE_RE = re.compile(rb"Started Launcher in \S+ seconds")
BUILD_LINE_RE = re.compile(rb"Build:\s+(?P<build>\d+)")

log = logging.getLogger("red.audio.manager")


class ServerManager:

    _java_available: ClassVar[Optional[bool]] = None
    _java_version: ClassVar[Optional[Tuple[int, int]]] = None
    _up_to_date: ClassVar[Optional[bool]] = None

    _blacklisted_archs = ["armv6l", "aarch32", "aarch64"]

    def __init__(self) -> None:
        self.ready = asyncio.Event()

        self._proc: Optional[asyncio.subprocess.Process] = None
        self._monitor_task: Optional[asyncio.Task] = None
        self._shutdown: bool = False

    async def start(self) -> None:
        arch_name = platform.machine()
        if arch_name in self._blacklisted_archs:
            raise asyncio.CancelledError(
                "You are attempting to run Lavalink audio on an unsupported machine architecture."
            )

        if self._proc is not None:
            if self._proc.returncode is None:
                raise RuntimeError("Internal Lavalink server is already running")
            else:
                raise RuntimeError("Server manager has already been used - create another one")

        await self.maybe_download_jar()

        # Copy the application.yml across.
        # For people to customise their Lavalink server configuration they need to run it
        # externally
        shutil.copyfile(BUNDLED_APP_YML, LAVALINK_APP_YML)

        args = await self._get_jar_args()
        self._proc = await asyncio.subprocess.create_subprocess_exec(
            *args,
            cwd=str(LAVALINK_DOWNLOAD_DIR),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )

        log.info("Internal Lavalink server started. PID: %s", self._proc.pid)

        try:
            await asyncio.wait_for(self._wait_for_launcher(), timeout=120)
        except asyncio.TimeoutError:
            log.warning("Timeout occurred whilst waiting for internal Lavalink server to be ready")

        self._monitor_task = asyncio.create_task(self._monitor())

    @classmethod
    async def _get_jar_args(cls) -> List[str]:
        java_available, java_version = await cls._has_java()
        if not java_available:
            raise RuntimeError("You must install Java 1.8+ for Lavalink to run.")

        if java_version == (1, 8):
            extra_flags = ["-Dsun.zip.disableMemoryMapping=true"]
        elif java_version >= (11, 0):
            extra_flags = ["-Djdk.tls.client.protocols=TLSv1.2"]
        else:
            extra_flags = []

        return ["java", *extra_flags, "-jar", str(LAVALINK_JAR_FILE)]

    @classmethod
    async def _has_java(cls) -> Tuple[bool, Optional[Tuple[int, int]]]:
        if cls._java_available is not None:
            # Return cached value if we've checked this before
            return cls._java_available, cls._java_version
        java_available = shutil.which("java") is not None
        if not java_available:
            cls.java_available = False
            cls.java_version = None
        else:
            cls._java_version = version = await cls._get_java_version()
            cls._java_available = (2, 0) > version >= (1, 8) or version >= (8, 0)
        return cls._java_available, cls._java_version

    @staticmethod
    async def _get_java_version() -> Tuple[int, int]:
        """
        This assumes we've already checked that java exists.
        """
        _proc: asyncio.subprocess.Process = await asyncio.create_subprocess_exec(
            "java", "-version", stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
        )
        # java -version outputs to stderr
        _, err = await _proc.communicate()

        version_info: str = err.decode("utf-8")
        # We expect the output to look something like:
        #     $ java -version
        #     ...
        #     ... version "MAJOR.MINOR.PATCH[_BUILD]" ...
        #     ...
        # We only care about the major and minor parts though.
        version_line_re = re.compile(
            r'version "(?P<major>\d+).(?P<minor>\d+).\d+(?:_\d+)?(?:-[A-Za-z0-9]+)?"'
        )
        short_version_re = re.compile(r'version "(?P<major>\d+)"')

        lines = version_info.splitlines()
        for line in lines:
            match = version_line_re.search(line)
            short_match = short_version_re.search(line)
            if match:
                return int(match["major"]), int(match["minor"])
            elif short_match:
                return int(short_match["major"]), 0

        raise RuntimeError(
            "The output of `java -version` was unexpected. Please report this issue on Red's "
            "issue tracker."
        )

    async def _wait_for_launcher(self) -> None:
        log.debug("Waiting for Lavalink server to be ready")
        for i in itertools.cycle(range(50)):
            line = await self._proc.stdout.readline()
            if READY_LINE_RE.search(line):
                self.ready.set()
                break
            if self._proc.returncode is not None:
                log.critical("Internal lavalink server exited early")
            if i == 49:
                # Sleep after 50 lines to prevent busylooping
                await asyncio.sleep(0.1)

    async def _monitor(self) -> None:
        while self._proc.returncode is None:
            await asyncio.sleep(0.5)

        # This task hasn't been cancelled - Lavalink was shut down by something else
        log.info("Internal Lavalink jar shutdown unexpectedly")
        if not self._has_java_error():
            log.info("Restarting internal Lavalink server")
            await self.start()
        else:
            log.critical(
                "Your Java is borked. Please find the hs_err_pid{}.log file"
                " in the Audio data folder and report this issue.",
                self._proc.pid,
            )

    def _has_java_error(self) -> bool:
        poss_error_file = LAVALINK_DOWNLOAD_DIR / "hs_err_pid{}.log".format(self._proc.pid)
        return poss_error_file.exists()

    async def shutdown(self) -> None:
        if self._shutdown is True or self._proc is None:
            # For convenience, calling this method more than once or calling it before starting it
            # does nothing.
            return
        log.info("Shutting down internal Lavalink server")
        if self._monitor_task is not None:
            self._monitor_task.cancel()
        self._proc.terminate()
        await self._proc.wait()
        self._shutdown = True

    @staticmethod
    async def _download_jar() -> None:
        log.info("Downloading Lavalink.jar...")
        async with aiohttp.ClientSession() as session:
            async with session.get(LAVALINK_DOWNLOAD_URL) as response:
                if response.status == 404:
                    raise RuntimeError(
                        f"Lavalink jar version {JAR_VERSION}_{JAR_BUILD} hasn't been published"
                    )
                fd, path = tempfile.mkstemp()
                file = open(fd, "wb")
                try:
                    chunk = await response.content.read(1024)
                    while chunk:
                        file.write(chunk)
                        chunk = await response.content.read(1024)
                    file.flush()
                finally:
                    file.close()
                pathlib.Path(path).replace(LAVALINK_JAR_FILE)

    @classmethod
    async def _is_up_to_date(cls):
        if cls._up_to_date is True:
            # Return cached value if we've checked this before
            return True
        args = await cls._get_jar_args()
        args.append("--version")
        _proc = await asyncio.subprocess.create_subprocess_exec(
            *args,
            cwd=str(LAVALINK_DOWNLOAD_DIR),
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.STDOUT,
        )
        stdout = (await _proc.communicate())[0]
        match = BUILD_LINE_RE.search(stdout)
        if not match:
            # Output is unexpected, suspect corrupted jarfile
            return False
        build = int(match["build"])
        cls._up_to_date = build == JAR_BUILD
        return cls._up_to_date

    @classmethod
    async def maybe_download_jar(cls):
        if not (LAVALINK_JAR_FILE.exists() and await cls._is_up_to_date()):
            await cls._download_jar()