diff options
Diffstat (limited to 'startop/scripts/lib/cmd_utils.py')
| -rw-r--r-- | startop/scripts/lib/cmd_utils.py | 184 |
1 files changed, 184 insertions, 0 deletions
diff --git a/startop/scripts/lib/cmd_utils.py b/startop/scripts/lib/cmd_utils.py new file mode 100644 index 000000000000..6071f145fe1d --- /dev/null +++ b/startop/scripts/lib/cmd_utils.py @@ -0,0 +1,184 @@ +#!/usr/bin/env python3 +# +# Copyright 2019, The Android Open Source Project +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Helper util libraries for command line operations.""" + +import asyncio +import sys +import time +from typing import Tuple, Optional, List + +import lib.print_utils as print_utils + +TIMEOUT = 50 +SIMULATE = False + +def run_command_nofail(cmd: List[str], **kwargs) -> None: + """Runs cmd list with default timeout. + + Throws exception if the execution fails. + """ + my_kwargs = {"timeout": TIMEOUT, "shell": False, "simulate": False} + my_kwargs.update(kwargs) + passed, out = execute_arbitrary_command(cmd, **my_kwargs) + if not passed: + raise RuntimeError( + "Failed to execute %s (kwargs=%s), output=%s" % (cmd, kwargs, out)) + +def run_adb_shell_command(cmd: str) -> Tuple[bool, str]: + """Runs command using adb shell. + + Returns: + A tuple of running status (True=succeeded, False=failed or timed out) and + std output (string contents of stdout with trailing whitespace removed). + """ + return run_shell_command('adb shell "{}"'.format(cmd)) + +def run_shell_func(script_path: str, + func: str, + args: List[str]) -> Tuple[bool, str]: + """Runs shell function with default timeout. + + Returns: + A tuple of running status (True=succeeded, False=failed or timed out) and + std output (string contents of stdout with trailing whitespace removed) . + """ + if args: + cmd = 'bash -c "source {script_path}; {func} {args}"'.format( + script_path=script_path, + func=func, + args=' '.join("'{}'".format(arg) for arg in args)) + else: + cmd = 'bash -c "source {script_path}; {func}"'.format( + script_path=script_path, + func=func) + + print_utils.debug_print(cmd) + return run_shell_command(cmd) + +def run_shell_command(cmd: str) -> Tuple[bool, str]: + """Runs shell command with default timeout. + + Returns: + A tuple of running status (True=succeeded, False=failed or timed out) and + std output (string contents of stdout with trailing whitespace removed) . + """ + return execute_arbitrary_command([cmd], + TIMEOUT, + shell=True, + simulate=SIMULATE) + +def execute_arbitrary_command(cmd: List[str], + timeout: int, + shell: bool, + simulate: bool) -> Tuple[bool, str]: + """Run arbitrary shell command with default timeout. + + Mostly copy from + frameworks/base/startop/scripts/app_startup/app_startup_runner.py. + + Args: + cmd: list of cmd strings. + timeout: the time limit of running cmd. + shell: indicate if the cmd is a shell command. + simulate: if it's true, do not run the command and assume the running is + successful. + + Returns: + A tuple of running status (True=succeeded, False=failed or timed out) and + std output (string contents of stdout with trailing whitespace removed) . + """ + if simulate: + print(cmd) + return True, '' + + print_utils.debug_print('[EXECUTE]', cmd) + # block until either command finishes or the timeout occurs. + loop = asyncio.get_event_loop() + + (return_code, script_output) = loop.run_until_complete( + _run_command(*cmd, shell=shell, timeout=timeout)) + + script_output = script_output.decode() # convert bytes to str + + passed = (return_code == 0) + print_utils.debug_print('[$?]', return_code) + if not passed: + print('[FAILED, code:%s]' % (return_code), script_output, file=sys.stderr) + + return passed, script_output.rstrip() + +async def _run_command(*args: List[str], + shell: bool = False, + timeout: Optional[int] = None) -> Tuple[int, bytes]: + if shell: + process = await asyncio.create_subprocess_shell( + *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT) + else: + process = await asyncio.create_subprocess_exec( + *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT) + + script_output = b'' + + print_utils.debug_print('[PID]', process.pid) + + timeout_remaining = timeout + time_started = time.time() + + # read line (sequence of bytes ending with b'\n') asynchronously + while True: + try: + line = await asyncio.wait_for(process.stdout.readline(), + timeout_remaining) + print_utils.debug_print('[STDOUT]', line) + script_output += line + + if timeout_remaining: + time_elapsed = time.time() - time_started + timeout_remaining = timeout - time_elapsed + except asyncio.TimeoutError: + print_utils.debug_print('[TIMEDOUT] Process ', process.pid) + + print_utils.debug_print('[TIMEDOUT] Sending SIGTERM.') + process.terminate() + + # 5 second timeout for process to handle SIGTERM nicely. + try: + (remaining_stdout, + remaining_stderr) = await asyncio.wait_for(process.communicate(), 5) + script_output += remaining_stdout + except asyncio.TimeoutError: + print_utils.debug_print('[TIMEDOUT] Sending SIGKILL.') + process.kill() + + # 5 second timeout to finish with SIGKILL. + try: + (remaining_stdout, + remaining_stderr) = await asyncio.wait_for(process.communicate(), 5) + script_output += remaining_stdout + except asyncio.TimeoutError: + # give up, this will leave a zombie process. + print_utils.debug_print('[TIMEDOUT] SIGKILL failed for process ', + process.pid) + time.sleep(100) + + return -1, script_output + else: + if not line: # EOF + break + + code = await process.wait() # wait for child process to exit + return code, script_output |