| # Copyright 2023 The Fuchsia Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Asynchronously run commands and read output as events. |
| |
| This module provides the AsyncCommand class, which runs a command in a separate |
| process and converts stdout, stderr, and termination into typed events that |
| can be iterated over asynchronously. |
| |
| asyncio supports spawning and monitoring subprocesses asynchronously, but |
| the official docs (https://docs.python.org/3.8/library/asyncio-subprocess.html) |
| include an ominous warning about awaiting on output streams directly. |
| Directly waiting on output streams can lead to a deadlock due to |
| output buffering, while relying on communicate() to read large |
| buffers data is not recommended. |
| |
| AsyncCommand handles spawning individual async tasks to read stderr |
| and stdout concurrently. Because asyncio internally connects these |
| pipes using non-blocking async readers, we do not hang. |
| |
| """ |
| |
| import asyncio |
| import atexit |
| import os |
| import signal |
| import time |
| import typing |
| from dataclasses import dataclass, field |
| from io import BytesIO, StringIO |
| |
| |
| @dataclass |
| class StdoutEvent: |
| """This event represents data emitted to the program's stdout pipe.""" |
| |
| # The text, as bytes. |
| text: bytes |
| |
| # The monotonic timestamp this program received the content. |
| timestamp: float = field(default_factory=lambda: time.monotonic()) |
| |
| |
| @dataclass |
| class StderrEvent: |
| """This event represents data emitted to the program's stderr pipe.""" |
| |
| # The text, as bytes. |
| text: bytes |
| |
| # The monotonic timestamp this program received the content. |
| timestamp: float = field(default_factory=lambda: time.monotonic()) |
| |
| |
| @dataclass |
| class TerminationEvent: |
| """This event represents the termination of a subcommand.""" |
| |
| # The return code for the program. |
| return_code: int |
| |
| # The program's runtime, in seconds. |
| runtime: float |
| |
| # If the output of the program was wrapped in another, this |
| # value indicates the return code of the wrapper program |
| wrapper_return_code: int | None = None |
| |
| # The monotonic timestamp this program received the termination event. |
| timestamp: float = field(default_factory=lambda: time.monotonic()) |
| |
| # If true, this command is terminating due to hitting its timeout. |
| was_timeout: bool = False |
| |
| |
| class AsyncCommandError(Exception): |
| """An error occurred starting or running a command asynchronously.""" |
| |
| |
| @dataclass |
| class CommandOutput: |
| """Summary information for an entire command execution.""" |
| |
| # The stdout contents, as a UTF-8 string. |
| stdout: str |
| |
| # The stderr contents, as a UTF-8 string. |
| stderr: str |
| |
| # The return code for the program. |
| return_code: int |
| |
| # The runtime for the program, in seconds. |
| runtime: float |
| |
| # If this program's output was wrapped by another program, this is |
| # the return code of that wrapper. |
| wrapper_return_code: int | None |
| |
| # If True, this command was forced to terminate due to timeout. |
| was_timeout: bool = False |
| |
| |
| CommandEvent = StdoutEvent | StderrEvent | TerminationEvent |
| |
| |
| @dataclass |
| class _InputWriter: |
| input_bytes: bytes |
| writer: asyncio.StreamWriter |
| |
| async def write_stream(self) -> None: |
| self.writer.write(self.input_bytes) |
| self.writer.write_eof() |
| await self.writer.drain() |
| |
| |
| class AsyncCommand: |
| """Asynchronously executed subprocess command. |
| |
| This class implements an iterator over CommandEvent, which represents |
| activity of the command over time. Note that at most one iterator |
| may exist over the command events. |
| """ |
| |
| def __init__( |
| self, |
| process: asyncio.subprocess.Process, |
| wrapped_process: asyncio.subprocess.Process | None = None, |
| timeout: float | None = None, |
| input_writer: _InputWriter | None = None, |
| ): |
| """Create an AsyncCommand that wraps a subprocess. |
| |
| The subprocess must have been created with both stdout and stderr |
| set to "PIPE". |
| |
| This method should not be used directly, use AsyncCommand.create instead. |
| |
| Args: |
| process (asyncio.subprocess.Process): The process to wrap. |
| wrapped_process (asyncio.subprocess.Process): A secondary process to wrap, which must also be closed. |
| timeout (float, optional): Terminate the process after this number of seconds. |
| """ |
| self._process = process |
| self._wrapped_process = wrapped_process |
| self._event_queue: asyncio.Queue[CommandEvent] = asyncio.Queue(128) |
| self._done_iterating = False |
| self._start = time.time() |
| self._timeout = timeout |
| self._input_writer = input_writer |
| self._runner_task = asyncio.create_task(self._task()) |
| |
| @classmethod |
| async def create( |
| cls: typing.Type["AsyncCommand"], |
| program: str, |
| *args: str, |
| symbolizer_args: list[str] | None = None, |
| env: dict[str, str] | None = None, |
| timeout: float | None = None, |
| input_bytes: bytes | None = None, |
| ) -> "AsyncCommand": |
| """Create a new AsyncCommand that runs the given program. |
| |
| Args: |
| program (str): Name of the program to run. |
| args (List[str]): Arguments to pass to the program. |
| symbolizer_args (List[str], optional): If set, pipe |
| output from the program through this program. |
| env (Dict[str,str], optional): If set, use this dict |
| to populate the command's environment. |
| Note: the CWD environment variable is handled specially to ensure |
| that the command runs in the given working directory. |
| Note: fx test's own environment is inherited and overridden by |
| the values in this dict, if set. |
| timeout (float, optional): Terminate the command after the given number of seconds. |
| input_bytes (bytes, optional): If set, send this input to the running command. |
| |
| Returns: |
| AsyncCommand: A new AsyncCommand executing the process. |
| |
| Raises: |
| AsyncCommandError: If there is a problem executing the process. |
| """ |
| cwd = env["CWD"] if env and "CWD" in env else None |
| if cwd and env: |
| env.pop("CWD") |
| if not env: |
| env = None |
| if env: |
| # Ensure that we inherit the incoming environment and |
| # extend it with the arguments to env, only if set. |
| new_env = dict(os.environ.items()) |
| new_env.update(env) |
| env = new_env |
| if cwd is not None and "CWD" in env: |
| # CWD is already handled above, do not inherit. |
| env.pop("CWD") |
| |
| try: |
| |
| async def make_process( |
| output_pipe: int | typing.TextIO, input_pipe: int | None = None |
| ) -> asyncio.subprocess.Process: |
| return await asyncio.subprocess.create_subprocess_exec( |
| program, |
| *args, |
| stdout=output_pipe, |
| stderr=output_pipe, |
| stdin=input_pipe, |
| env=env, |
| cwd=cwd, |
| preexec_fn=os.setpgrp, # To support killing by pgid. |
| ) |
| |
| base_command = None |
| input_pipe = None if input is None else asyncio.subprocess.PIPE |
| input_writer: _InputWriter | None = None |
| if symbolizer_args: |
| # Wrap the base command we want to run in another |
| # command that executes the symbolizer. We need to explicitly |
| # close our side of the pipes after passing to the relevant |
| # subprocesses. |
| |
| read_pipe, write_pipe = os.pipe() |
| |
| base_command = await make_process(write_pipe, input_pipe) |
| os.close(write_pipe) |
| command = await asyncio.subprocess.create_subprocess_exec( |
| *symbolizer_args, |
| stdout=asyncio.subprocess.PIPE, |
| stderr=asyncio.subprocess.PIPE, |
| stdin=read_pipe, |
| preexec_fn=os.setpgrp, # To support killing by pgid. |
| ) |
| if input_bytes is not None and base_command.stdin is not None: |
| input_writer = _InputWriter(input_bytes, base_command.stdin) |
| os.close(read_pipe) |
| else: |
| command = await make_process( |
| asyncio.subprocess.PIPE, input_pipe |
| ) |
| if input_bytes is not None and command.stdin is not None: |
| input_writer = _InputWriter(input_bytes, command.stdin) |
| |
| return AsyncCommand( |
| command, |
| base_command, |
| timeout=timeout, |
| input_writer=input_writer, |
| ) |
| except FileNotFoundError as e: |
| raise AsyncCommandError( |
| f"The program '{e.filename}' was not found." |
| ) |
| except Exception as e: |
| raise AsyncCommandError(f"An unknown error occurred: {e}") |
| |
| def send_signal(self, sig: int) -> None: |
| """Send the given signal to the underlying process(es) and their |
| children in the same process group. |
| |
| Args: |
| sig (int): The signal to send. |
| """ |
| # We need to signal the entire process group we created, otherwise |
| # we run the risk of signallying only a shell script and not |
| # the processes run by that shell script (e.g. `fx`). |
| pg = os.getpgid(self._process.pid) |
| os.killpg(pg, sig) |
| if self._wrapped_process is not None: |
| pg = os.getpgid(self._wrapped_process.pid) |
| os.killpg(pg, sig) |
| |
| def terminate(self) -> None: |
| """Terminate the underlying process(es) by sending SIGTERM.""" |
| self.send_signal(signal.SIGTERM) |
| |
| def kill(self) -> None: |
| """Immediately kill the underlying process(es) by sending SIGKILL.""" |
| self.send_signal(signal.SIGKILL) |
| |
| async def run_to_completion( |
| self, |
| callback: typing.Callable[[CommandEvent], None] | None = None, |
| ) -> CommandOutput: |
| """Runs the program to completion, collecting the resulting outputs. |
| |
| Args: |
| callback (typing.Callable[[CommandEvent], None], optional): If set, send all |
| CommandEvents to this callback function as they are produced. |
| Defaults to None. |
| |
| Raises: |
| AsyncCommandError: If an internal error causes the task queue to be cancelled. |
| |
| Returns: |
| CommandOutput: Summary of the execution of the program. |
| """ |
| with StringIO() as stdout, StringIO() as stderr: |
| async for e in self: |
| if callback: |
| callback(e) |
| if isinstance(e, StdoutEvent): |
| text = e.text.decode("utf-8", errors="replace") |
| stdout.write(text) |
| elif isinstance(e, StderrEvent): |
| text = e.text.decode("utf-8", errors="replace") |
| stderr.write(text) |
| elif isinstance(e, TerminationEvent): |
| return CommandOutput( |
| stdout=stdout.getvalue(), |
| stderr=stderr.getvalue(), |
| return_code=e.return_code, |
| wrapper_return_code=e.wrapper_return_code, |
| runtime=e.runtime, |
| was_timeout=e.was_timeout, |
| ) |
| |
| raise AsyncCommandError("Event stream unexpectedly stopped") |
| |
| async def _task(self) -> None: |
| # Make sure that the child process group is killed when this process |
| # exits. Otherwise the child process group keeps running if a user |
| # does Ctrl+C on this script. |
| def kill_async_process() -> None: |
| self.kill() |
| |
| atexit.register(kill_async_process) |
| |
| tasks = [] |
| |
| async def read_stream( |
| input_stream: asyncio.StreamReader, |
| event_type: typing.Type[typing.Union[StderrEvent, StdoutEvent]], |
| ) -> None: |
| """Wrap reading from a stream and emitting a specific type of event. |
| |
| Args: |
| input_stream (asyncio.StreamReader): Reader to read from. |
| event_type: Either StdoutEvent or StderrEvent type, to wrap each line. |
| """ |
| |
| # Iteratively read blocks of data from the stream, splitting on newlines. |
| |
| # buf is used to contain the line in progress, it will contain only |
| # one line and will never contain anything after a newline character. |
| buf = BytesIO() |
| |
| while not input_stream.at_eof(): |
| more: bytes = await input_stream.read(128 * 1024) |
| # Process the current buffer, splitting by lines. |
| while True: |
| # Find the end of the line. |
| if (idx := more.find(b"\n")) == -1: |
| # No newline yet, buffer this batch and fetch more. |
| buf.write(more) |
| break |
| |
| # Append the rest of the line to the current line, then send as an event. |
| buf.write(more[: idx + 1]) |
| await self._event_queue.put(event_type(text=buf.getvalue())) |
| |
| # Reset the buffer and continue with the next line. |
| buf = BytesIO() |
| more = more[idx + 1 :] |
| |
| # Publish any remaining data as a line. |
| if len(buf.getvalue()) > 0: |
| await self._event_queue.put(event_type(text=buf.getvalue())) |
| |
| async def timeout_handler( |
| timeout: float, timeout_signal: asyncio.Event |
| ) -> None: |
| """Handle timeouts. |
| |
| We first send SIGTERM after the timeout, and if the program has |
| still not terminated after an additional delay we send SIGKILL. |
| |
| Args: |
| timeout (float): The initial wait time before termination. |
| timeout_signal (asyncio.Event): An event to set on timeout. |
| """ |
| await asyncio.sleep(timeout) |
| timeout_signal.set() |
| self.terminate() |
| # Give the process at most 5 seconds to terminate. |
| await asyncio.sleep(5) |
| self.kill() |
| |
| timeout_task: asyncio.Task[None] | None = None |
| did_timeout = asyncio.Event() |
| if self._timeout is not None: |
| timeout_task = asyncio.create_task( |
| timeout_handler(self._timeout, did_timeout) |
| ) |
| |
| # Read stdout and stderr in parallel, forwarding the appropriate events. |
| assert self._process.stdout |
| assert self._process.stderr |
| tasks.append( |
| asyncio.create_task(read_stream(self._process.stdout, StdoutEvent)) |
| ) |
| tasks.append( |
| asyncio.create_task(read_stream(self._process.stderr, StderrEvent)) |
| ) |
| if self._input_writer is not None: |
| tasks.append(asyncio.create_task(self._input_writer.write_stream())) |
| |
| # Wait for the process to exit and get its return code. |
| return_code = await self._process.wait() |
| if timeout_task and not timeout_task.done(): |
| # Cancel the timeout task so we do not try to send signals to a terminated process. |
| timeout_task.cancel() |
| wrapper_return_code: int | None = None |
| if self._wrapped_process: |
| # Also ensure we wait for the wrapped process, and use it's return code as the canonical code. |
| wrapper_return_code = return_code |
| return_code = await self._wrapped_process.wait() |
| |
| # There is no longer a need to cleanup this process, since it already |
| # terminated. |
| atexit.unregister(kill_async_process) |
| |
| runtime = time.time() - self._start |
| |
| if timeout_task is not None: |
| # Make sure the timeout task is cancelled and awaited |
| # so it is cleaned up. |
| timeout_task.cancel() |
| tasks.append(timeout_task) |
| |
| # Wait for all output to drain before termination event. |
| await asyncio.wait(tasks, return_when="FIRST_EXCEPTION") |
| await self._event_queue.put( |
| TerminationEvent( |
| return_code, |
| runtime, |
| wrapper_return_code=wrapper_return_code, |
| was_timeout=did_timeout.is_set(), |
| ) |
| ) |
| |
| async def next_event(self) -> CommandEvent | None: |
| """Return the next event from the process execution, if one exists. |
| |
| Returns: |
| Event | None: The next event on the command, or None if the command is terminated. |
| """ |
| if self._done_iterating: |
| return None |
| |
| ret = await self._event_queue.get() |
| if isinstance(ret, TerminationEvent): |
| self._done_iterating = True |
| return ret |
| |
| def __aiter__(self) -> typing.Self: |
| return self |
| |
| async def __anext__(self) -> CommandEvent: |
| ret = await self.next_event() |
| if ret is None: |
| raise StopAsyncIteration() |
| else: |
| return ret |
