1 files changed, 166 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..c3d96059c91c
--- /dev/null
+++ b/startop/scripts/lib/cmd_utils.py
@@ -0,0 +1,166 @@
+#!/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_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) .
+  """
+  cmd = 'bash -c "source {script_path}; {func} {args}"'.format(
+    script_path=script_path,
+    func=func,
+    args=' '.join("'{}'".format(arg) for arg in args))
+  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