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