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, AsyncIterator, Callable, Coroutine
17from contextlib import asynccontextmanager, suppress
18from pathlib import Path
19from signal import SIGINT
20from types import TracebackType
21from typing import Any, Self
22
23from music_assistant.constants import MASS_LOGGER_NAME, VERBOSE_LOG_LEVEL
24
25LOGGER = logging.getLogger(f"{MASS_LOGGER_NAME}.helpers.process")
26
27DEFAULT_CHUNKSIZE = 64000
28
29
30def get_subprocess_env(env: dict[str, str] | None = None) -> dict[str, str]:
31    """Get environment for subprocess, stripping LD_PRELOAD to avoid jemalloc warnings."""
32    result = dict(os.environ)
33    result.pop("LD_PRELOAD", None)
34    if env:
35        result.update(env)
36    return result
37
38
39class AsyncProcess:
40    """
41    AsyncProcess.
42
43    Wrapper around asyncio subprocess to help with using pipe streams and
44    taking care of properly closing the process in case of exit (on both success and failures),
45    without deadlocking.
46    """
47
48    _stdin_feeder_task: asyncio.Task[None] | None = None  # used for ffmpeg
49    _stderr_reader_task: asyncio.Task[None] | None = None  # used for ffmpeg
50
51    def __init__(
52        self,
53        args: list[str],
54        stdin: bool | int | None = None,
55        stdout: bool | int | None = None,
56        stderr: bool | int | None = False,
57        name: str | None = None,
58        env: dict[str, str] | None = None,
59    ) -> None:
60        """
61        Initialize AsyncProcess.
62
63        :param args: Command and arguments to execute.
64        :param stdin: Stdin configuration (True for PIPE, False for None, or custom).
65        :param stdout: Stdout configuration (True for PIPE, False for None, or custom).
66        :param stderr: Stderr configuration (True for PIPE, False for DEVNULL, or custom).
67        :param name: Process name for logging.
68        :param env: Environment variables for the subprocess (None inherits parent env).
69        """
70        self.proc: asyncio.subprocess.Process | None = None
71        if name is None:
72            name = Path(args[0]).name
73        self.name = name
74        self.logger = LOGGER.getChild(name)
75        self._args = args
76        self._stdin = None if stdin is False else stdin
77        self._stdout = None if stdout is False else stdout
78        self._stderr = asyncio.subprocess.DEVNULL if stderr is False else stderr
79        self._env = get_subprocess_env(env)
80        self._stderr_lock = asyncio.Lock()
81        self._stdout_lock = asyncio.Lock()
82        self._stdin_lock = asyncio.Lock()
83        self._close_called = False
84        self._returncode: int | None = None
85
86    @property
87    def closed(self) -> bool:
88        """Return if the process was closed."""
89        return self._close_called or self.returncode is not None
90
91    @property
92    def returncode(self) -> int | None:
93        """Return the erturncode of the process."""
94        if self._returncode is not None:
95            return self._returncode
96        if self.proc is None:
97            return None
98        if (ret_code := self.proc.returncode) is not None:
99            self._returncode = ret_code
100        return ret_code
101
102    async def __aenter__(self) -> Self:
103        """Enter context manager."""
104        await self.start()
105        return self
106
107    async def __aexit__(
108        self,
109        exc_type: type[BaseException] | None,
110        exc_val: BaseException | None,
111        exc_tb: TracebackType | None,
112    ) -> bool | None:
113        """Exit context manager."""
114        # make sure we close and cleanup the process
115        await self.close()
116        self._returncode = self.returncode
117        return None
118
119    async def start(self) -> None:
120        """Perform Async init of process."""
121        self.proc = await asyncio.create_subprocess_exec(
122            *self._args,
123            stdin=asyncio.subprocess.PIPE if self._stdin is True else self._stdin,
124            stdout=asyncio.subprocess.PIPE if self._stdout is True else self._stdout,
125            stderr=asyncio.subprocess.PIPE if self._stderr is True else self._stderr,
126            env=self._env,
127            bufsize=0,
128        )
129        self.logger.log(
130            VERBOSE_LOG_LEVEL, "Process %s started with PID %s", self.name, self.proc.pid
131        )
132
133    async def iter_chunked(self, n: int = DEFAULT_CHUNKSIZE) -> AsyncGenerator[bytes]:
134        """Yield chunks of n size from the process stdout."""
135        while True:
136            chunk = await self.readexactly(n)
137            if len(chunk) == 0:
138                break
139            yield chunk
140
141    async def iter_any(self, n: int = DEFAULT_CHUNKSIZE) -> AsyncGenerator[bytes]:
142        """Yield chunks as they come in from process stdout."""
143        while True:
144            chunk = await self.read(n)
145            if len(chunk) == 0:
146                break
147            yield chunk
148
149    async def readexactly(self, n: int) -> bytes:
150        """Read exactly n bytes from the process stdout (or less if eof)."""
151        if self._close_called:
152            return b""
153        assert self.proc is not None  # for type checking
154        assert self.proc.stdout is not None  # for type checking
155        async with self._stdout_lock:
156            try:
157                return await self.proc.stdout.readexactly(n)
158            except asyncio.IncompleteReadError as err:
159                return err.partial
160
161    async def read(self, n: int) -> bytes:
162        """
163        Read up to n bytes from the stdout stream.
164
165        If n is positive, this function try to read n bytes,
166        and may return less or equal bytes than requested, but at least one byte.
167        If EOF was received before any byte is read, this function returns empty byte object.
168        """
169        if self._close_called:
170            return b""
171        assert self.proc is not None  # for type checking
172        assert self.proc.stdout is not None  # for type checking
173        async with self._stdout_lock:
174            return await self.proc.stdout.read(n)
175
176    async def write(self, data: bytes) -> None:
177        """Write data to process stdin."""
178        if self._close_called or self.proc is None:
179            return
180        if self.proc.stdin is None:
181            return
182        async with self._stdin_lock:
183            self.proc.stdin.write(data)
184            await self.proc.stdin.drain()
185
186    @asynccontextmanager
187    async def stdin_quiesced(self, timeout: float = 5.0) -> AsyncIterator[bool]:
188        """
189        Hold stdin quiet for a block, with what was already written seen through to the pipe.
190
191        :meth:`write` only waits while the transport is paused, which it is only
192        above the high-water mark, so it returns with up to that much still queued
193        locally (64 KiB by default). This first sees those bytes through to the
194        kernel pipe -- as far as it can guarantee; whether the process has read
195        them is its own business -- and then keeps the write lock for the body, so
196        no :meth:`write` or :meth:`write_eof` can interleave. For a caller telling
197        the process something about the bytes it has been handed -- out of band,
198        and in a sequence the process must not see a write inside -- that turns
199        "we happen to have stopped writing" into something the block enforces.
200
201        Yields True when stdin was emptied, False when it could not be: the
202        process is then still owed bytes, so a caller whose message depends on it
203        having received everything must give up rather than send it.
204
205        :param timeout: Seconds to wait for the buffer to empty.
206        """
207        if self._close_called or self.proc is None or self.proc.stdin is None:
208            yield True
209            return
210        async with self._stdin_lock:
211            yield await self._drain_stdin_locked(timeout)
212
213    async def write_eof(self) -> None:
214        """Write end of file to to process stdin."""
215        if self._close_called or self.proc is None:
216            return
217        if self.proc.stdin is None:
218            return
219        async with self._stdin_lock:
220            try:
221                if self.proc.stdin.can_write_eof():
222                    self.proc.stdin.write_eof()
223                await self.proc.stdin.drain()
224            except (
225                AttributeError,
226                AssertionError,
227                BrokenPipeError,
228                RuntimeError,
229                ConnectionResetError,
230            ):
231                # already exited, race condition
232                pass
233
234    async def read_stderr(self) -> bytes:
235        """Read line from stderr."""
236        if self.returncode is not None:
237            return b""
238        assert self.proc is not None  # for type checking
239        assert self.proc.stderr is not None  # for type checking
240        return await self._readline(self.proc.stderr, self._stderr_lock)
241
242    async def read_stdout(self) -> bytes:
243        """Read line from stdout."""
244        # keyed on the close flag rather than the returncode (like read() and
245        # readexactly()): a process that already exited still has its last
246        # lines sitting in the stream buffer, and those must still be readable
247        if self._close_called:
248            return b""
249        assert self.proc is not None  # for type checking
250        assert self.proc.stdout is not None  # for type checking
251        return await self._readline(self.proc.stdout, self._stdout_lock)
252
253    async def iter_stderr(self) -> AsyncGenerator[str]:
254        """Iterate lines from the stderr stream as string."""
255        async for line in self._iter_lines(self.read_stderr):
256            yield line
257
258    async def iter_stdout(self) -> AsyncGenerator[str]:
259        """Iterate lines from the stdout stream as string."""
260        async for line in self._iter_lines(self.read_stdout):
261            yield line
262
263    async def communicate(
264        self,
265        input: bytes | None = None,  # noqa: A002
266        timeout: float | None = None,
267    ) -> tuple[bytes, bytes]:
268        """Communicate with the process and return stdout and stderr."""
269        if self.closed:
270            raise RuntimeError("communicate called while process already done")
271        # abort existing readers on stderr/stdout first before we send communicate
272        await self._stderr_lock.acquire()
273        await self._stdout_lock.acquire()
274        assert self.proc is not None  # for type checking
275        stdout, stderr = await asyncio.wait_for(self.proc.communicate(input), timeout)
276        return (stdout, stderr)
277
278    async def close(self) -> None:
279        """Close/terminate the process and wait for exit."""
280        if self._close_called and self.returncode is not None:
281            # Already closed and reaped, so there is nothing left to signal or
282            # drain. The stream locks below are still held by that first call
283            # and would only be waited out again (5s each).
284            return
285        self._close_called = True
286        if not self.proc:
287            return
288
289        # cancel existing stdin feeder task if any
290        if self._stdin_feeder_task:
291            if not self._stdin_feeder_task.done():
292                self._stdin_feeder_task.cancel()
293            # Always await the task to consume any exception and prevent
294            # "Task exception was never retrieved" errors.
295            try:
296                await self._stdin_feeder_task
297            except asyncio.CancelledError:
298                pass  # Expected when we cancel the task
299            except Exception as err:
300                # Log unexpected exceptions from the stdin feeder before suppressing
301                LOGGER.warning(
302                    "Process stdin feeder task ended with error: %s",
303                    err,
304                )
305
306        # close stdin to signal we're done sending data
307        with suppress(TimeoutError, asyncio.CancelledError):
308            await asyncio.wait_for(self._stdin_lock.acquire(), 5)
309        if self.proc.stdin and not self.proc.stdin.is_closing():
310            self.proc.stdin.close()
311        elif not self.proc.stdin and self.proc.returncode is None:
312            # the process may exit between the returncode check and the signal; guard the
313            # race the same way the SIGKILL delivery below does
314            with suppress(ProcessLookupError, OSError):
315                self.proc.send_signal(SIGINT)
316
317        # ensure we have no more readers active and stdout is drained
318        with suppress(TimeoutError, asyncio.CancelledError):
319            await asyncio.wait_for(self._stdout_lock.acquire(), 5)
320        if self.proc.stdout and not self.proc.stdout.at_eof():
321            with suppress(Exception):
322                await self.proc.stdout.read(-1)
323        # if we have a stderr task active, allow it to finish
324        if self._stderr_reader_task:
325            with suppress(TimeoutError, asyncio.CancelledError):
326                await asyncio.wait_for(self._stderr_reader_task, 5)
327        elif self.proc.stderr and not self.proc.stderr.at_eof():
328            with suppress(TimeoutError, asyncio.CancelledError):
329                await asyncio.wait_for(self._stderr_lock.acquire(), 5)
330            # drain stderr
331            with suppress(Exception):
332                await self.proc.stderr.read(-1)
333
334        # make sure the process is really cleaned up.
335        # especially with pipes this can cause deadlocks if not properly guarded
336        # we need to ensure stdout and stderr are flushed and stdin closed
337        pid = self.proc.pid
338        terminate_attempts = 0
339        while self.returncode is None:
340            try:
341                # use communicate to flush all pipe buffers
342                await asyncio.wait_for(self.proc.communicate(), 2)
343            except TimeoutError:
344                terminate_attempts += 1
345                self.logger.debug(
346                    "Process %s with PID %s did not stop in time (attempt %d). Sending SIGKILL...",
347                    self.name,
348                    pid,
349                    terminate_attempts,
350                )
351                # Use os.kill for more direct signal delivery
352                with suppress(ProcessLookupError, OSError):
353                    os.kill(pid, 9)  # SIGKILL = 9
354                # Give up after 5 attempts - process may be zombie
355                if terminate_attempts >= 5:
356                    self.logger.warning(
357                        "Process %s (PID %s) did not terminate after %d SIGKILL attempts",
358                        self.name,
359                        pid,
360                        terminate_attempts,
361                    )
362                    break
363        self.logger.log(
364            VERBOSE_LOG_LEVEL,
365            "Process %s with PID %s stopped with returncode %s",
366            self.name,
367            self.proc.pid,
368            self.returncode,
369        )
370
371    async def kill(self) -> None:
372        """
373        Immediately kill the process with SIGKILL.
374
375        Use this for forceful termination when the process doesn't respond to
376        normal termination signals. Unlike close(), this doesn't attempt graceful
377        shutdown - it immediately sends SIGKILL.
378        """
379        self._close_called = True
380        if not self.proc or self.returncode is not None:
381            return
382
383        pid = self.proc.pid
384
385        # Cancel stdin feeder task if any
386        if self._stdin_feeder_task and not self._stdin_feeder_task.done():
387            self._stdin_feeder_task.cancel()
388            with suppress(asyncio.CancelledError, Exception):
389                await self._stdin_feeder_task
390
391        # Cancel stderr reader task if any
392        if self._stderr_reader_task and not self._stderr_reader_task.done():
393            self._stderr_reader_task.cancel()
394            with suppress(asyncio.CancelledError, Exception):
395                await self._stderr_reader_task
396
397        # Close stdin to signal we're done sending data
398        # Note: Don't manually call feed_eof() on stdout/stderr - this causes
399        # "feed_data after feed_eof" assertion errors when the subprocess transport
400        # still has buffered data to deliver. Let the process termination naturally
401        # close the streams.
402        if self.proc.stdin and not self.proc.stdin.is_closing():
403            self.proc.stdin.close()
404
405        # Send SIGKILL immediately using os.kill for more direct signal delivery
406        self.logger.debug("Killing process %s with PID %s", self.name, pid)
407        with suppress(ProcessLookupError, OSError):
408            os.kill(pid, 9)  # SIGKILL = 9
409
410        # Wait for process to actually terminate
411        try:
412            await asyncio.wait_for(self.proc.wait(), 2)
413        except TimeoutError:
414            # Try one more time with os.kill
415            with suppress(ProcessLookupError, OSError):
416                os.kill(pid, 9)
417            try:
418                await asyncio.wait_for(self.proc.wait(), 2)
419            except TimeoutError:
420                self.logger.warning(
421                    "Process %s with PID %s did not terminate after SIGKILL - may be zombie",
422                    self.name,
423                    pid,
424                )
425
426        self.logger.log(
427            VERBOSE_LOG_LEVEL,
428            "Process %s with PID %s killed with returncode %s",
429            self.name,
430            pid,
431            self.returncode,
432        )
433
434    async def wait(self) -> int:
435        """Wait for the process and return the returncode."""
436        if self._returncode is None:
437            assert self.proc is not None
438            self._returncode = await self.proc.wait()
439        return self._returncode
440
441    async def wait_with_timeout(self, timeout: int) -> int:
442        """Wait for the process and return the returncode with a timeout."""
443        return await asyncio.wait_for(self.wait(), timeout)
444
445    def attach_stderr_reader(self, task: asyncio.Task[None]) -> None:
446        """Attach a stderr reader task to this process."""
447        self._stderr_reader_task = task
448
449    async def _readline(self, stream: asyncio.StreamReader, lock: asyncio.Lock) -> bytes:
450        """
451        Read a single line from one of the process' output streams.
452
453        :param stream: The stream to read the line from.
454        :param lock: The lock guarding that stream's readers.
455        """
456        async with lock:
457            try:
458                return await stream.readline()
459            except ValueError as err:
460                # we're waiting for a line (separator found), but the line was too big
461                # this may happen with ffmpeg during a long (radio) stream where progress
462                # gets outputted to the stderr but no newline
463                # https://stackoverflow.com/questions/55457370/how-to-avoid-valueerror-separator-is-not-found-and-chunk-exceed-the-limit
464                # NOTE: this consumes the line that was too big
465                if "chunk exceed the limit" in str(err):
466                    return await stream.readline()
467                # raise for all other (value) errors
468                raise
469
470    async def _iter_lines(
471        self, read_line: Callable[[], Coroutine[Any, Any, bytes]]
472    ) -> AsyncGenerator[str]:
473        """
474        Yield decoded, non-empty lines until the underlying stream reaches EOF.
475
476        :param read_line: Coroutine function returning the next raw line.
477        """
478        while True:
479            raw = await read_line()
480            if raw == b"":
481                break
482            if line := raw.decode("utf-8", errors="ignore").strip():
483                yield line
484
485    async def _drain_stdin_locked(self, timeout: float) -> bool:
486        """
487        Empty the stdin write buffer, with the write lock already held.
488
489        :param timeout: Seconds to wait for the buffer to empty.
490        :return: True once the buffer is empty, False when the wait timed out.
491        """
492        assert self.proc is not None  # for type checking
493        assert self.proc.stdin is not None  # for type checking
494        transport = self.proc.stdin.transport
495        low, high = transport.get_write_buffer_limits()
496        try:
497            # Pausing the transport at a zero high-water mark is what makes
498            # drain() resolve only once the buffer is completely empty: it
499            # otherwise resolves as soon as the transport is not paused.
500            transport.set_write_buffer_limits(high=0)
501            await asyncio.wait_for(self.proc.stdin.drain(), timeout)
502        except TimeoutError:
503            return False
504        except BrokenPipeError, RuntimeError, ConnectionResetError:
505            # already exited, race condition: nothing is left to arrive
506            return True
507        finally:
508            # Restore what this process was configured with rather than the
509            # asyncio defaults a bare call would reinstate.
510            with suppress(RuntimeError):
511                transport.set_write_buffer_limits(high=high, low=low)
512        return True
513
514
515async def check_output(
516    *args: str, env: dict[str, str] | None = None, timeout: float | None = None
517) -> tuple[int, bytes]:
518    """
519    Run subprocess and return returncode and output.
520
521    :param env: Optional environment overrides for the subprocess.
522    :param timeout: Maximum seconds to wait for the process to exit. On expiry the
523        process is killed and TimeoutError is raised; None (default) waits forever.
524    """
525    proc = await asyncio.create_subprocess_exec(
526        *args,
527        stderr=asyncio.subprocess.STDOUT,
528        stdout=asyncio.subprocess.PIPE,
529        env=get_subprocess_env(env),
530    )
531    try:
532        async with asyncio.timeout(timeout):
533            stdout, _ = await proc.communicate()
534    except TimeoutError:
535        proc.kill()
536        with suppress(ProcessLookupError):
537            await proc.wait()
538        raise
539    assert proc.returncode is not None  # for type checking
540    return (proc.returncode, stdout)
541
542
543async def communicate(
544    args: list[str],
545    input: bytes | None = None,  # noqa: A002
546) -> tuple[int, bytes, bytes]:
547    """Communicate with subprocess and return returncode, stdout and stderr output."""
548    proc = await asyncio.create_subprocess_exec(
549        *args,
550        stderr=asyncio.subprocess.PIPE,
551        stdout=asyncio.subprocess.PIPE,
552        stdin=asyncio.subprocess.PIPE if input is not None else None,
553        env=get_subprocess_env(),
554    )
555    stdout, stderr = await proc.communicate(input)
556    assert proc.returncode is not None  # for type checking
557    return (proc.returncode, stdout, stderr)
558