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