async def _get_proc_output(
proc: asyncio.subprocess.Process,
in_data: Optional[bytes],
timeout: Optional[int],
text: Union[bool, FormatType],
) -> Tuple[ProcessData, ProcessData, Optional[int]]:
stdout: Any
stderr: Any
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(in_data),
timeout)
except asyncio.TimeoutError:
try:
proc.kill()
except ProcessLookupError:
pass
raise
if text:
if text is not StderrOnly and stdout is not None:
stdout = stdout.decode(errors="replace").strip()
if stderr is not None:
stderr = stderr.decode(errors="replace").strip()
return stdout, stderr, proc.returncode
async def stream_subprocess_output(
process: asyncio.subprocess.Process,
*,
timeout: Timeout = None,
stdout_callback: Optional[OutputStreamCallback] = None,
stderr_callback: Optional[OutputStreamCallback] = None,
) -> int:
"""
Asynchronously read the stdout and stderr output streams of a subprocess and
and optionally invoke a callback with each line of text read.
:param process: An asyncio subprocess created with `create_subprocess_exec` or `create_subprocess_shell`.
:param timeout: An optional timeout in seconds for how long to read the streams before giving up.
:param stdout_callback: An optional callable invoked with each line read from stdout. Must accept a single string positional argument and returns nothing.
:param stderr_callback: An optional callable invoked with each line read from stderr. Must accept a single string positional argument and returns nothing.
:raises asyncio.TimeoutError: Raised if the timeout expires before the subprocess exits.
:return: The exit status of the subprocess.
"""
tasks = []
if process.stdout:
tasks.append(
asyncio.create_task(
_read_lines_from_output_stream(process.stdout,
stdout_callback),
name="stdout",
))
if process.stderr:
tasks.append(
asyncio.create_task(
_read_lines_from_output_stream(process.stderr,
stderr_callback),
name="stderr",
))
timeout_in_seconds = (timeout.total_seconds() if isinstance(
timeout, datetime.timedelta) else timeout)
try:
# Gather the stream output tasks and the parent process
gather_task = asyncio.gather(*tasks, process.wait())
await asyncio.wait_for(gather_task, timeout=timeout_in_seconds)
except (asyncio.TimeoutError, asyncio.CancelledError):
with contextlib.suppress(ProcessLookupError):
if process.returncode is None:
process.terminate()
with contextlib.suppress(asyncio.CancelledError):
await gather_task
[task.cancel() for task in tasks]
await asyncio.gather(*tasks, return_exceptions=True)
raise
return cast(int, process.returncode)
async def terminate_and_wait(proc: asyncio.subprocess.Process) -> None:
try:
proc.terminate()
try:
await asyncio.wait_for(proc.wait(), timeout=2.0)
except asyncio.TimeoutError:
proc.kill()
await proc.wait()
except ProcessLookupError:
pass
async def _get_proc_output(proc: asyncio.subprocess.Process,
input: Optional[bytes],
timeout: int) -> Tuple[bytes, bytes, Optional[int]]:
try:
stdout, stderr = await asyncio.wait_for(proc.communicate(input),
timeout)
except asyncio.TimeoutError:
try:
proc.kill()
except ProcessLookupError:
pass
raise
return stdout, stderr, proc.returncode
async def _terminate_process(process: asyncio.subprocess.Process, timeout: int,
logger: logging.Logger) -> None:
returncode = process.returncode
if returncode is not None:
logger.info(f"Process has exited with {returncode}")
return
logger.info(f"Stopping process with SIGTERM, waiting {timeout} seconds")
process.terminate()
try:
returncode = await asyncio.wait_for(process.wait(), timeout=timeout)
logger.info(f"Process has exited after SIGTERM with {returncode}")
except TimeoutError:
logger.info(
f"Process hasn't exited after {timeout} seconds, SIGKILL'ing...")
process.kill()
async def drain(self, sub_process: asyncio.subprocess.Process) -> int:
# pipe, so stream is actual, not optional
out = cast(asyncio.StreamReader, sub_process.stdout)
err = cast(asyncio.StreamReader, sub_process.stderr)
ret, _, _ = await asyncio.gather(sub_process.wait(),
self.drain_out(out),
self.drain_err(err))
return ret
async def close_subprocess(
# TODO[Pylint issue 1469]: Does not recognize `asyncio.subprocess`.
subprocess: asyncio.subprocess.Process, # pylint: disable=no-member
) -> None:
"""Ensure the given subprocess is terminated."""
# We do not know what state the process is in.
# We assume the user had already exhausted
# all nice ways to terminate it.
# So just kill it.
with contextlib.suppress(ProcessLookupError):
subprocess.kill()
# Killing just sends the request / signal.
# Wait to make sure it is actually terminated.
# And automatically-created pipes and inherited fds,
# such as any given in stdin, stdout, stderr,
# are closed after termination.
await subprocess.communicate()
async def soft_kill(process: asyncio.subprocess.Process) -> None:
# First try terminating...
try:
process.terminate()
await asyncio.wait_for(process.wait(), timeout=45.0)
return
except ProcessLookupError:
# (can be thrown e.g. if the process has exited in the meantime)
return
except asyncio.TimeoutError:
pass
# ... then try killing
try:
process.kill()
await process.wait()
except ProcessLookupError:
return
async def _terminate_process(process: asyncio.subprocess.Process,
logger: logging.Logger,
timeout: int = 30) -> None:
logger.info(
f"Stopping process {process} with SIGINT, waiting {timeout} seconds")
if process.returncode is not None:
logger.info(f"Process is already terminated with {process.returncode} "
"perhaps it died prematurely")
return
process.send_signal(signal.SIGINT)
try:
await asyncio.wait_for(process.wait(), timeout=timeout)
logger.info(f"Stopped process {process} with SIGINT")
except TimeoutError:
logger.info(
f"Process {process} didn't close after {timeout} seconds, killing..."
)
finally:
if process.returncode is None:
process.kill()
async def kill_process(
proc: asyncio.subprocess.Process, wait: float,
logger: logging.Logger) -> None: # pylint: disable=no-member
if proc.returncode is None:
try:
proc.terminate()
await asyncio.sleep(wait)
if proc.returncode is None:
try:
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
except Exception:
if proc.returncode is not None:
raise
await proc.wait()
logger.info("Process killed: retcode=%d", proc.returncode)
except asyncio.CancelledError:
pass
except Exception:
if proc.returncode is None:
logger.exception("Can't kill process pid=%d", proc.pid)
else:
logger.info("Process killed: retcode=%d", proc.returncode)
async def gen_listening_ports_from_fd(
process: asyncio.subprocess.Process,
read_fd: int,
timeout: Optional[int] = None,
logger: Optional[logging.Logger] = None,
) -> Tuple[int, Optional[int]]:
if logger is None:
logger = logging.getLogger("reply-fd")
wait = asyncio.ensure_future(process.wait())
ports = asyncio.ensure_future(_read_from_fd(read_fd, logger))
done, pending = await asyncio.wait([wait, ports],
return_when=asyncio.FIRST_COMPLETED)
for fut in pending:
fut.cancel()
if ports not in done:
raise Exception(
f"Process exited with return code {process.returncode} before "
f"responding with port")
return await ports
async def kill_process(self, proc: asyncio.subprocess.Process):
proc.terminate()
try:
await asyncio.wait_for(proc.wait(), timeout=20)
except TimeoutError:
proc.kill()
