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 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 _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 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 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 _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(self, proc: asyncio.subprocess.Process):
proc.terminate()
try:
await asyncio.wait_for(proc.wait(), timeout=20)
except TimeoutError:
proc.kill()