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