music-assistant-server

1"""
2AsyncProcess.
3
4Wrapper around asyncio subprocess to help with using pipe streams and
5taking care of properly closing the process in case of exit (on both success and failures),
6without deadlocking.
7"""
8
9from __future__ import annotations
10
11import asyncio
12import logging
13import os
14
15# if TYPE_CHECKING:
16from collections.abc import AsyncGenerator
17from contextlib import suppress
18from signal import SIGINT
19from types import TracebackType
20from typing import Self
21
22from music_assistant.constants import MASS_LOGGER_NAME, VERBOSE_LOG_LEVEL
23
24LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.helpers.process")
25
26DEFAULT_CHUNKSIZE = 64000
27
28
29def get_subprocess_env(env: dict[str, str] | None = None) -> dict[str, str]:
30    """Get environment for subprocess, stripping LD_PRELOAD to avoid jemalloc warnings."""
31    result = dict(os.environ)
32    result.pop("LD_PRELOAD", None)
33    if env:
34        result.update(env)
35    return result
36
37
38class AsyncProcess:
39    """
40    AsyncProcess.
41
42    Wrapper around asyncio subprocess to help with using pipe streams and
43    taking care of properly closing the process in case of exit (on both success and failures),
44    without deadlocking.
45    """
46
47    _stdin_feeder_task: asyncio.Task[None] | None = None  # used for ffmpeg
48    _stderr_reader_task: asyncio.Task[None] | None = None  # used for ffmpeg
49
50    def __init__(
51        self,
52        args: list[str],
53        stdin: bool | int | None = None,
54        stdout: bool | int | None = None,
55        stderr: bool | int | None = False,
56        name: str | None = None,
57        env: dict[str, str] | None = None,
58    ) -> None:
59        """
60        Initialize AsyncProcess.
61
62        :param args: Command and arguments to execute.
63        :param stdin: Stdin configuration (True for PIPE, False for None, or custom).
64        :param stdout: Stdout configuration (True for PIPE, False for None, or custom).
65        :param stderr: Stderr configuration (True for PIPE, False for DEVNULL, or custom).
66        :param name: Process name for logging.
67        :param env: Environment variables for the subprocess (None inherits parent env).
68        """
69        self.proc: asyncio.subprocess.Process | None = None
70        if name is None:
71            name = args[0].split(os.sep)[-1]
72        self.name = name
73        self.logger = LOGGER.getChild(name)
74        self._args = args
75        self._stdin = None if stdin is False else stdin
76        self._stdout = None if stdout is False else stdout
77        self._stderr = asyncio.subprocess.DEVNULL if stderr is False else stderr
78        self._env = get_subprocess_env(env)
79        self._stderr_lock = asyncio.Lock()
80        self._stdout_lock = asyncio.Lock()
81        self._stdin_lock = asyncio.Lock()
82        self._close_called = False
83        self._returncode: int | None = None
84
85    @property
86    def closed(self) -> bool:
87        """Return if the process was closed."""
88        return self._close_called or self.returncode is not None
89
90    @property
91    def returncode(self) -> int | None:
92        """Return the erturncode of the process."""
93        if self._returncode is not None:
94            return self._returncode
95        if self.proc is None:
96            return None
97        if (ret_code := self.proc.returncode) is not None:
98            self._returncode = ret_code
99        return ret_code
100
101    async def __aenter__(self) -> Self:
102        """Enter context manager."""
103        await self.start()
104        return self
105
106    async def __aexit__(
107        self,
108        exc_type: type[BaseException] | None,
109        exc_val: BaseException | None,
110        exc_tb: TracebackType | None,
111    ) -> bool | None:
112        """Exit context manager."""
113        # make sure we close and cleanup the process
114        await self.close()
115        self._returncode = self.returncode
116        return None
117
118    async def start(self) -> None:
119        """Perform Async init of process."""
120        self.proc = await asyncio.create_subprocess_exec(
121            *self._args,
122            stdin=asyncio.subprocess.PIPE if self._stdin is True else self._stdin,
123            stdout=asyncio.subprocess.PIPE if self._stdout is True else self._stdout,
124            stderr=asyncio.subprocess.PIPE if self._stderr is True else self._stderr,
125            env=self._env,
126            bufsize=0,
127        )
128        self.logger.log(
129            VERBOSE_LOG_LEVEL, "Process %s started with PID %s", self.name, self.proc.pid
130        )
131
132    async def iter_chunked(self, n: int = DEFAULT_CHUNKSIZE) -> AsyncGenerator[bytes]:
133        """Yield chunks of n size from the process stdout."""
134        while True:
135            chunk = await self.readexactly(n)
136            if len(chunk) == 0:
137                break
138            yield chunk
139
140    async def iter_any(self, n: int = DEFAULT_CHUNKSIZE) -> AsyncGenerator[bytes]:
141        """Yield chunks as they come in from process stdout."""
142        while True:
143            chunk = await self.read(n)
144            if len(chunk) == 0:
145                break
146            yield chunk
147
148    async def readexactly(self, n: int) -> bytes:
149        """Read exactly n bytes from the process stdout (or less if eof)."""
150        if self._close_called:
151            return b""
152        assert self.proc is not None  # for type checking
153        assert self.proc.stdout is not None  # for type checking
154        async with self._stdout_lock:
155            try:
156                return await self.proc.stdout.readexactly(n)
157            except asyncio.IncompleteReadError as err:
158                return err.partial
159
160    async def read(self, n: int) -> bytes:
161        """
162        Read up to n bytes from the stdout stream.
163
164        If n is positive, this function try to read n bytes,
165        and may return less or equal bytes than requested, but at least one byte.
166        If EOF was received before any byte is read, this function returns empty byte object.
167        """
168        if self._close_called:
169            return b""
170        assert self.proc is not None  # for type checking
171        assert self.proc.stdout is not None  # for type checking
172        async with self._stdout_lock:
173            return await self.proc.stdout.read(n)
174
175    async def write(self, data: bytes) -> None:
176        """Write data to process stdin."""
177        if self._close_called or self.proc is None:
178            return
179        if self.proc.stdin is None:
180            return
181        async with self._stdin_lock:
182            self.proc.stdin.write(data)
183            await self.proc.stdin.drain()
184
185    async def write_eof(self) -> None:
186        """Write end of file to to process stdin."""
187        if self._close_called or self.proc is None:
188            return
189        if self.proc.stdin is None:
190            return
191        async with self._stdin_lock:
192            try:
193                if self.proc.stdin.can_write_eof():
194                    self.proc.stdin.write_eof()
195                await self.proc.stdin.drain()
196            except (
197                AttributeError,
198                AssertionError,
199                BrokenPipeError,
200                RuntimeError,
201                ConnectionResetError,
202            ):
203                # already exited, race condition
204                pass
205
206    async def read_stderr(self) -> bytes:
207        """Read line from stderr."""
208        if self.returncode is not None:
209            return b""
210        assert self.proc is not None  # for type checking
211        assert self.proc.stderr is not None  # for type checking
212        async with self._stderr_lock:
213            try:
214                return await self.proc.stderr.readline()
215            except ValueError as err:
216                # we're waiting for a line (separator found), but the line was too big
217                # this may happen with ffmpeg during a long (radio) stream where progress
218                # gets outputted to the stderr but no newline
219                # https://stackoverflow.com/questions/55457370/how-to-avoid-valueerror-separator-is-not-found-and-chunk-exceed-the-limit
220                # NOTE: this consumes the line that was too big
221                if "chunk exceed the limit" in str(err):
222                    return await self.proc.stderr.readline()
223                # raise for all other (value) errors
224                raise
225
226    async def iter_stderr(self) -> AsyncGenerator[str]:
227        """Iterate lines from the stderr stream as string."""
228        line: str | bytes
229        while True:
230            line = await self.read_stderr()
231            if line == b"":
232                break
233            line = line.decode("utf-8", errors="ignore").strip()
234            if not line:
235                continue
236            yield line
237
238    async def communicate(
239        self,
240        input: bytes | None = None,  # noqa: A002
241        timeout: float | None = None,
242    ) -> tuple[bytes, bytes]:
243        """Communicate with the process and return stdout and stderr."""
244        if self.closed:
245            raise RuntimeError("communicate called while process already done")
246        # abort existing readers on stderr/stdout first before we send communicate
247        await self._stderr_lock.acquire()
248        await self._stdout_lock.acquire()
249        assert self.proc is not None  # for type checking
250        stdout, stderr = await asyncio.wait_for(self.proc.communicate(input), timeout)
251        return (stdout, stderr)
252
253    async def close(self) -> None:
254        """Close/terminate the process and wait for exit."""
255        self._close_called = True
256        if not self.proc:
257            return
258
259        # cancel existing stdin feeder task if any
260        if self._stdin_feeder_task:
261            if not self._stdin_feeder_task.done():
262                self._stdin_feeder_task.cancel()
263            # Always await the task to consume any exception and prevent
264            # "Task exception was never retrieved" errors.
265            try:
266                await self._stdin_feeder_task
267            except asyncio.CancelledError:
268                pass  # Expected when we cancel the task
269            except Exception as err:
270                # Log unexpected exceptions from the stdin feeder before suppressing
271                LOGGER.warning(
272                    "Process stdin feeder task ended with error: %s",
273                    err,
274                )
275
276        # close stdin to signal we're done sending data
277        with suppress(TimeoutError, asyncio.CancelledError):
278            await asyncio.wait_for(self._stdin_lock.acquire(), 5)
279        if self.proc.stdin and not self.proc.stdin.is_closing():
280            self.proc.stdin.close()
281        elif not self.proc.stdin and self.proc.returncode is None:
282            # the process may exit between the returncode check and the signal; guard the
283            # race the same way the SIGKILL delivery below does
284            with suppress(ProcessLookupError, OSError):
285                self.proc.send_signal(SIGINT)
286
287        # ensure we have no more readers active and stdout is drained
288        with suppress(TimeoutError, asyncio.CancelledError):
289            await asyncio.wait_for(self._stdout_lock.acquire(), 5)
290        if self.proc.stdout and not self.proc.stdout.at_eof():
291            with suppress(Exception):
292                await self.proc.stdout.read(-1)
293        # if we have a stderr task active, allow it to finish
294        if self._stderr_reader_task:
295            with suppress(TimeoutError, asyncio.CancelledError):
296                await asyncio.wait_for(self._stderr_reader_task, 5)
297        elif self.proc.stderr and not self.proc.stderr.at_eof():
298            with suppress(TimeoutError, asyncio.CancelledError):
299                await asyncio.wait_for(self._stderr_lock.acquire(), 5)
300            # drain stderr
301            with suppress(Exception):
302                await self.proc.stderr.read(-1)
303
304        # make sure the process is really cleaned up.
305        # especially with pipes this can cause deadlocks if not properly guarded
306        # we need to ensure stdout and stderr are flushed and stdin closed
307        pid = self.proc.pid
308        terminate_attempts = 0
309        while self.returncode is None:
310            try:
311                # use communicate to flush all pipe buffers
312                await asyncio.wait_for(self.proc.communicate(), 2)
313            except TimeoutError:
314                terminate_attempts += 1
315                self.logger.debug(
316                    "Process %s with PID %s did not stop in time (attempt %d). Sending SIGKILL...",
317                    self.name,
318                    pid,
319                    terminate_attempts,
320                )
321                # Use os.kill for more direct signal delivery
322                with suppress(ProcessLookupError, OSError):
323                    os.kill(pid, 9)  # SIGKILL = 9
324                # Give up after 5 attempts - process may be zombie
325                if terminate_attempts >= 5:
326                    self.logger.warning(
327                        "Process %s (PID %s) did not terminate after %d SIGKILL attempts",
328                        self.name,
329                        pid,
330                        terminate_attempts,
331                    )
332                    break
333        self.logger.log(
334            VERBOSE_LOG_LEVEL,
335            "Process %s with PID %s stopped with returncode %s",
336            self.name,
337            self.proc.pid,
338            self.returncode,
339        )
340
341    async def kill(self) -> None:
342        """
343        Immediately kill the process with SIGKILL.
344
345        Use this for forceful termination when the process doesn't respond to
346        normal termination signals. Unlike close(), this doesn't attempt graceful
347        shutdown - it immediately sends SIGKILL.
348        """
349        self._close_called = True
350        if not self.proc or self.returncode is not None:
351            return
352
353        pid = self.proc.pid
354
355        # Cancel stdin feeder task if any
356        if self._stdin_feeder_task and not self._stdin_feeder_task.done():
357            self._stdin_feeder_task.cancel()
358            with suppress(asyncio.CancelledError, Exception):
359                await self._stdin_feeder_task
360
361        # Cancel stderr reader task if any
362        if self._stderr_reader_task and not self._stderr_reader_task.done():
363            self._stderr_reader_task.cancel()
364            with suppress(asyncio.CancelledError, Exception):
365                await self._stderr_reader_task
366
367        # Close stdin to signal we're done sending data
368        # Note: Don't manually call feed_eof() on stdout/stderr - this causes
369        # "feed_data after feed_eof" assertion errors when the subprocess transport
370        # still has buffered data to deliver. Let the process termination naturally
371        # close the streams.
372        if self.proc.stdin and not self.proc.stdin.is_closing():
373            self.proc.stdin.close()
374
375        # Send SIGKILL immediately using os.kill for more direct signal delivery
376        self.logger.debug("Killing process %s with PID %s", self.name, pid)
377        with suppress(ProcessLookupError, OSError):
378            os.kill(pid, 9)  # SIGKILL = 9
379
380        # Wait for process to actually terminate
381        try:
382            await asyncio.wait_for(self.proc.wait(), 2)
383        except TimeoutError:
384            # Try one more time with os.kill
385            with suppress(ProcessLookupError, OSError):
386                os.kill(pid, 9)
387            try:
388                await asyncio.wait_for(self.proc.wait(), 2)
389            except TimeoutError:
390                self.logger.warning(
391                    "Process %s with PID %s did not terminate after SIGKILL - may be zombie",
392                    self.name,
393                    pid,
394                )
395
396        self.logger.log(
397            VERBOSE_LOG_LEVEL,
398            "Process %s with PID %s killed with returncode %s",
399            self.name,
400            pid,
401            self.returncode,
402        )
403
404    async def wait(self) -> int:
405        """Wait for the process and return the returncode."""
406        if self._returncode is None:
407            assert self.proc is not None
408            self._returncode = await self.proc.wait()
409        return self._returncode
410
411    async def wait_with_timeout(self, timeout: int) -> int:
412        """Wait for the process and return the returncode with a timeout."""
413        return await asyncio.wait_for(self.wait(), timeout)
414
415    def attach_stderr_reader(self, task: asyncio.Task[None]) -> None:
416        """Attach a stderr reader task to this process."""
417        self._stderr_reader_task = task
418
419
420async def check_output(
421    *args: str, env: dict[str, str] | None = None, timeout: float | None = None
422) -> tuple[int, bytes]:
423    """
424    Run subprocess and return returncode and output.
425
426    :param env: Optional environment overrides for the subprocess.
427    :param timeout: Maximum seconds to wait for the process to exit. On expiry the
428        process is killed and TimeoutError is raised; None (default) waits forever.
429    """
430    proc = await asyncio.create_subprocess_exec(
431        *args,
432        stderr=asyncio.subprocess.STDOUT,
433        stdout=asyncio.subprocess.PIPE,
434        env=get_subprocess_env(env),
435    )
436    try:
437        async with asyncio.timeout(timeout):
438            stdout, _ = await proc.communicate()
439    except TimeoutError:
440        proc.kill()
441        with suppress(ProcessLookupError):
442            await proc.wait()
443        raise
444    assert proc.returncode is not None  # for type checking
445    return (proc.returncode, stdout)
446
447
448async def communicate(
449    args: list[str],
450    input: bytes | None = None,  # noqa: A002
451) -> tuple[int, bytes, bytes]:
452    """Communicate with subprocess and return returncode, stdout and stderr output."""
453    proc = await asyncio.create_subprocess_exec(
454        *args,
455        stderr=asyncio.subprocess.PIPE,
456        stdout=asyncio.subprocess.PIPE,
457        stdin=asyncio.subprocess.PIPE if input is not None else None,
458        env=get_subprocess_env(),
459    )
460    stdout, stderr = await proc.communicate(input)
461    assert proc.returncode is not None  # for type checking
462    return (proc.returncode, stdout, stderr)
463