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        """
290        Close/terminate the process and wait for exit.
291
292        An enclosing timeout is not a reliable bound on this call: the cleanup may
293        swallow the cancellation and run to completion, and a cancellation that does
294        land leaves the process unreaped.
295        """
296        if self._close_called and self.returncode is not None:
297            # Already closed and reaped, so there is nothing left to signal or
298            # drain. The stream locks below are still held by that first call
299            # and would only be waited out again (5s each).
300            return
301        self._close_called = True
302        if not self.proc:
303            return
304
305        if self._stdin_feeder_task:
306            await self._cancel_and_await(self._stdin_feeder_task, "stdin feeder")
307
308        # close stdin to signal we're done sending data
309        with suppress(TimeoutError, asyncio.CancelledError):
310            await asyncio.wait_for(self._stdin_lock.acquire(), 5)
311        if self.proc.stdin and not self.proc.stdin.is_closing():
312            self.proc.stdin.close()
313        elif not self.proc.stdin and self.proc.returncode is None:
314            # the process may exit between the returncode check and the signal; guard the
315            # race the same way the SIGKILL delivery below does
316            with suppress(ProcessLookupError, OSError):
317                self.proc.send_signal(SIGINT)
318
319        # ensure we have no more readers active and stdout is drained
320        with suppress(TimeoutError, asyncio.CancelledError):
321            await asyncio.wait_for(self._stdout_lock.acquire(), 5)
322        if self.proc.stdout and not self.proc.stdout.at_eof():
323            with suppress(Exception):
324                await asyncio.wait_for(self.proc.stdout.read(-1), PIPE_DRAIN_TIMEOUT)
325        # if we have a stderr task active, allow it to finish
326        if self._stderr_reader_task:
327            with suppress(TimeoutError, asyncio.CancelledError):
328                await asyncio.wait_for(self._stderr_reader_task, 5)
329        elif self.proc.stderr and not self.proc.stderr.at_eof():
330            with suppress(TimeoutError, asyncio.CancelledError):
331                await asyncio.wait_for(self._stderr_lock.acquire(), 5)
332            # drain stderr
333            with suppress(Exception):
334                await asyncio.wait_for(self.proc.stderr.read(-1), PIPE_DRAIN_TIMEOUT)
335
336        # make sure the process is really cleaned up.
337        # especially with pipes this can cause deadlocks if not properly guarded
338        # we need to ensure stdout and stderr are flushed and stdin closed
339        pid = self.proc.pid
340        terminate_attempts = 0
341        while self.returncode is None:
342            try:
343                # use communicate to flush all pipe buffers
344                await asyncio.wait_for(self.proc.communicate(), 2)
345            except TimeoutError:
346                terminate_attempts += 1
347                self.logger.debug(
348                    "Process %s with PID %s did not stop in time (attempt %d). Sending SIGKILL...",
349                    self.name,
350                    pid,
351                    terminate_attempts,
352                )
353                # Use os.kill for more direct signal delivery
354                with suppress(ProcessLookupError, OSError):
355                    os.kill(pid, 9)  # SIGKILL = 9
356                # Give up after 5 attempts - process may be zombie
357                if terminate_attempts >= 5:
358                    self.logger.warning(
359                        "Process %s (PID %s) did not terminate after %d SIGKILL attempts",
360                        self.name,
361                        pid,
362                        terminate_attempts,
363                    )
364                    break
365        self.logger.log(
366            VERBOSE_LOG_LEVEL,
367            "Process %s with PID %s stopped with returncode %s",
368            self.name,
369            self.proc.pid,
370            self.returncode,
371        )
372
373    async def kill(self) -> None:
374        """
375        Immediately kill the process with SIGKILL.
376
377        Use this for forceful termination when the process doesn't respond to
378        normal termination signals. Unlike close(), this doesn't attempt graceful
379        shutdown - it immediately sends SIGKILL.
380        """
381        self._close_called = True
382        if not self.proc or self.returncode is not None:
383            return
384
385        pid = self.proc.pid
386
387        if self._stdin_feeder_task:
388            await self._cancel_and_await(self._stdin_feeder_task, "stdin feeder")
389        if self._stderr_reader_task:
390            await self._cancel_and_await(self._stderr_reader_task, "stderr reader")
391
392        # Close stdin to signal we're done sending data
393        # Note: Don't manually call feed_eof() on stdout/stderr - this causes
394        # "feed_data after feed_eof" assertion errors when the subprocess transport
395        # still has buffered data to deliver. Let the process termination naturally
396        # close the streams.
397        if self.proc.stdin and not self.proc.stdin.is_closing():
398            self.proc.stdin.close()
399
400        # Send SIGKILL immediately using os.kill for more direct signal delivery
401        self.logger.debug("Killing process %s with PID %s", self.name, pid)
402        with suppress(ProcessLookupError, OSError):
403            os.kill(pid, 9)  # SIGKILL = 9
404
405        # SIGKILL leaves whatever the child already wrote in the pipes, and the reap
406        # below only completes once they disconnect - so drain them here rather than
407        # waiting that out for output nobody is going to read
408        try:
409            await asyncio.wait_for(self.proc.communicate(), 2)
410        except TimeoutError:
411            pass  # the escalation below takes over
412        except Exception as err:
413            self.logger.warning("Failed to drain the pipes of PID %s: %s", pid, err)
414
415        # Wait for process to actually terminate
416        try:
417            await asyncio.wait_for(self.proc.wait(), 2)
418        except TimeoutError:
419            # Try one more time with os.kill
420            with suppress(ProcessLookupError, OSError):
421                os.kill(pid, 9)
422            try:
423                await asyncio.wait_for(self.proc.wait(), 2)
424            except TimeoutError:
425                self.logger.warning(
426                    "Process %s with PID %s did not terminate after SIGKILL - may be zombie",
427                    self.name,
428                    pid,
429                )
430
431        self.logger.log(
432            VERBOSE_LOG_LEVEL,
433            "Process %s with PID %s killed with returncode %s",
434            self.name,
435            pid,
436            self.returncode,
437        )
438
439    async def wait(self) -> int:
440        """Wait for the process and return the returncode."""
441        if self._returncode is None:
442            assert self.proc is not None
443            self._returncode = await self.proc.wait()
444        return self._returncode
445
446    async def wait_with_timeout(self, timeout: int) -> int:
447        """Wait for the process and return the returncode with a timeout."""
448        return await asyncio.wait_for(self.wait(), timeout)
449
450    def attach_stderr_reader(self, task: asyncio.Task[None]) -> None:
451        """Attach a stderr reader task to this process."""
452        self._stderr_reader_task = task
453
454    async def _readline(self, stream: asyncio.StreamReader, lock: asyncio.Lock) -> bytes:
455        """
456        Read a single line from one of the process' output streams.
457
458        :param stream: The stream to read the line from.
459        :param lock: The lock guarding that stream's readers.
460        """
461        async with lock:
462            try:
463                return await stream.readline()
464            except ValueError as err:
465                # we're waiting for a line (separator found), but the line was too big
466                # this may happen with ffmpeg during a long (radio) stream where progress
467                # gets outputted to the stderr but no newline
468                # https://stackoverflow.com/questions/55457370/how-to-avoid-valueerror-separator-is-not-found-and-chunk-exceed-the-limit
469                # NOTE: this consumes the line that was too big
470                if "chunk exceed the limit" in str(err):
471                    return await stream.readline()
472                # raise for all other (value) errors
473                raise
474
475    async def _iter_lines(
476        self, read_line: Callable[[], Coroutine[Any, Any, bytes]]
477    ) -> AsyncGenerator[str]:
478        """
479        Yield decoded, non-empty lines until the underlying stream reaches EOF.
480
481        :param read_line: Coroutine function returning the next raw line.
482        """
483        while True:
484            raw = await read_line()
485            if raw == b"":
486                break
487            if line := raw.decode("utf-8", errors="ignore").strip():
488                yield line
489
490    async def _drain_stdin_locked(self, timeout: float) -> bool:
491        """
492        Empty the stdin write buffer, with the write lock already held.
493
494        :param timeout: Seconds to wait for the buffer to empty.
495        :return: True once the buffer is empty, False when the wait timed out.
496        """
497        assert self.proc is not None  # for type checking
498        assert self.proc.stdin is not None  # for type checking
499        transport = self.proc.stdin.transport
500        low, high = transport.get_write_buffer_limits()
501        try:
502            # Pausing the transport at a zero high-water mark is what makes
503            # drain() resolve only once the buffer is completely empty: it
504            # otherwise resolves as soon as the transport is not paused.
505            transport.set_write_buffer_limits(high=0)
506            await asyncio.wait_for(self.proc.stdin.drain(), timeout)
507        except TimeoutError:
508            return False
509        except BrokenPipeError, RuntimeError, ConnectionResetError:
510            # already exited, race condition: nothing is left to arrive
511            return True
512        finally:
513            # Restore what this process was configured with rather than the
514            # asyncio defaults a bare call would reinstate.
515            with suppress(RuntimeError):
516                transport.set_write_buffer_limits(high=high, low=low)
517        return True
518
519    async def _cancel_and_await(self, task: asyncio.Task[None], description: str) -> None:
520        """
521        Cancel one of this process' helper tasks and wait for it to end.
522
523        A task that already finished is awaited too, so it is never left with an
524        unretrieved exception. An unexpected error is logged rather than raised.
525
526        :param task: The task to cancel and await.
527        :param description: How the task is named when logging an unexpected error.
528        """
529        if not task.done():
530            task.cancel()
531        try:
532            await task
533        except asyncio.CancelledError:
534            pass  # expected when we cancel the task
535        except Exception as err:
536            # retrieving the failure is what keeps asyncio from reporting it as
537            # unhandled once the task is collected, so log it here rather than
538            # dropping the only trace of it
539            self.logger.warning("The %s task ended with error: %s", description, err)
540
541
542async def check_output(
543    *args: str, env: dict[str, str] | None = None, timeout: float | None = None
544) -> tuple[int, bytes]:
545    """
546    Run subprocess and return returncode and output.
547
548    :param env: Optional environment overrides for the subprocess.
549    :param timeout: Maximum seconds to wait for the process to exit. On expiry the
550        process is killed and TimeoutError is raised; None (default) waits forever.
551    """
552    proc = await asyncio.create_subprocess_exec(
553        *args,
554        stderr=asyncio.subprocess.STDOUT,
555        stdout=asyncio.subprocess.PIPE,
556        env=get_subprocess_env(env),
557    )
558    try:
559        async with asyncio.timeout(timeout):
560            stdout, _ = await proc.communicate()
561    except TimeoutError:
562        proc.kill()
563        with suppress(ProcessLookupError):
564            await proc.wait()
565        raise
566    assert proc.returncode is not None  # for type checking
567    return (proc.returncode, stdout)
568
569
570async def communicate(
571    args: list[str],
572    input: bytes | None = None,  # noqa: A002
573) -> tuple[int, bytes, bytes]:
574    """Communicate with subprocess and return returncode, stdout and stderr output."""
575    proc = await asyncio.create_subprocess_exec(
576        *args,
577        stderr=asyncio.subprocess.PIPE,
578        stdout=asyncio.subprocess.PIPE,
579        stdin=asyncio.subprocess.PIPE if input is not None else None,
580        env=get_subprocess_env(),
581    )
582    stdout, stderr = await proc.communicate(input)
583    assert proc.returncode is not None  # for type checking
584    return (proc.returncode, stdout, stderr)
585