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