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