• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1#!/usr/bin/env python3
2#
3# Copyright 2019, The Android Open Source Project
4#
5# Licensed under the Apache License, Version 2.0 (the "License");
6# you may not use this file except in compliance with the License.
7# You may obtain a copy of the License at
8#
9#     http://www.apache.org/licenses/LICENSE-2.0
10#
11# Unless required by applicable law or agreed to in writing, software
12# distributed under the License is distributed on an "AS IS" BASIS,
13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14# See the License for the specific language governing permissions and
15# limitations under the License.
16
17"""Helper util libraries for command line operations."""
18
19import asyncio
20import sys
21import time
22from typing import Tuple, Optional, List
23
24import lib.print_utils as print_utils
25
26TIMEOUT = 50
27SIMULATE = False
28
29def run_command_nofail(cmd: List[str], **kwargs) -> None:
30  """Runs cmd list with default timeout.
31
32     Throws exception if the execution fails.
33  """
34  my_kwargs = {"timeout": TIMEOUT, "shell": False, "simulate": False}
35  my_kwargs.update(kwargs)
36  passed, out = execute_arbitrary_command(cmd, **my_kwargs)
37  if not passed:
38    raise RuntimeError(
39      "Failed to execute %s (kwargs=%s), output=%s" % (cmd, kwargs, out))
40
41def run_adb_shell_command(cmd: str) -> Tuple[bool, str]:
42  """Runs command using adb shell.
43
44  Returns:
45    A tuple of running status (True=succeeded, False=failed or timed out) and
46    std output (string contents of stdout with trailing whitespace removed).
47  """
48  return run_shell_command('adb shell "{}"'.format(cmd))
49
50def run_shell_func(script_path: str,
51                   func: str,
52                   args: List[str]) -> Tuple[bool, str]:
53  """Runs shell function with default timeout.
54
55  Returns:
56    A tuple of running status (True=succeeded, False=failed or timed out) and
57    std output (string contents of stdout with trailing whitespace removed) .
58  """
59  if args:
60    cmd = 'bash -c "source {script_path}; {func} {args}"'.format(
61      script_path=script_path,
62      func=func,
63      args=' '.join("'{}'".format(arg) for arg in args))
64  else:
65    cmd = 'bash -c "source {script_path}; {func}"'.format(
66      script_path=script_path,
67      func=func)
68
69  print_utils.debug_print(cmd)
70  return run_shell_command(cmd)
71
72def run_shell_command(cmd: str) -> Tuple[bool, str]:
73  """Runs shell command with default timeout.
74
75  Returns:
76    A tuple of running status (True=succeeded, False=failed or timed out) and
77    std output (string contents of stdout with trailing whitespace removed) .
78  """
79  return execute_arbitrary_command([cmd],
80                                   TIMEOUT,
81                                   shell=True,
82                                   simulate=SIMULATE)
83
84def execute_arbitrary_command(cmd: List[str],
85                              timeout: int,
86                              shell: bool,
87                              simulate: bool) -> Tuple[bool, str]:
88  """Run arbitrary shell command with default timeout.
89
90    Mostly copy from
91    frameworks/base/startop/scripts/app_startup/app_startup_runner.py.
92
93  Args:
94    cmd: list of cmd strings.
95    timeout: the time limit of running cmd.
96    shell: indicate if the cmd is a shell command.
97    simulate: if it's true, do not run the command and assume the running is
98        successful.
99
100  Returns:
101    A tuple of running status (True=succeeded, False=failed or timed out) and
102    std output (string contents of stdout with trailing whitespace removed) .
103  """
104  if simulate:
105    print(cmd)
106    return True, ''
107
108  print_utils.debug_print('[EXECUTE]', cmd)
109  # block until either command finishes or the timeout occurs.
110  loop = asyncio.get_event_loop()
111
112  (return_code, script_output) = loop.run_until_complete(
113    _run_command(*cmd, shell=shell, timeout=timeout))
114
115  script_output = script_output.decode()  # convert bytes to str
116
117  passed = (return_code == 0)
118  print_utils.debug_print('[$?]', return_code)
119  if not passed:
120    print('[FAILED, code:%s]' % (return_code), script_output, file=sys.stderr)
121
122  return passed, script_output.rstrip()
123
124async def _run_command(*args: List[str],
125                       shell: bool = False,
126                       timeout: Optional[int] = None) -> Tuple[int, bytes]:
127  if shell:
128    process = await asyncio.create_subprocess_shell(
129      *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT)
130  else:
131    process = await asyncio.create_subprocess_exec(
132      *args, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.STDOUT)
133
134  script_output = b''
135
136  print_utils.debug_print('[PID]', process.pid)
137
138  timeout_remaining = timeout
139  time_started = time.time()
140
141  # read line (sequence of bytes ending with b'\n') asynchronously
142  while True:
143    try:
144      line = await asyncio.wait_for(process.stdout.readline(),
145                                    timeout_remaining)
146      print_utils.debug_print('[STDOUT]', line)
147      script_output += line
148
149      if timeout_remaining:
150        time_elapsed = time.time() - time_started
151        timeout_remaining = timeout - time_elapsed
152    except asyncio.TimeoutError:
153      print_utils.debug_print('[TIMEDOUT] Process ', process.pid)
154
155      print_utils.debug_print('[TIMEDOUT] Sending SIGTERM.')
156      process.terminate()
157
158      # 5 second timeout for process to handle SIGTERM nicely.
159      try:
160        (remaining_stdout,
161         remaining_stderr) = await asyncio.wait_for(process.communicate(), 5)
162        script_output += remaining_stdout
163      except asyncio.TimeoutError:
164        print_utils.debug_print('[TIMEDOUT] Sending SIGKILL.')
165        process.kill()
166
167      # 5 second timeout to finish with SIGKILL.
168      try:
169        (remaining_stdout,
170         remaining_stderr) = await asyncio.wait_for(process.communicate(), 5)
171        script_output += remaining_stdout
172      except asyncio.TimeoutError:
173        # give up, this will leave a zombie process.
174        print_utils.debug_print('[TIMEDOUT] SIGKILL failed for process ',
175                                process.pid)
176        time.sleep(100)
177
178      return -1, script_output
179    else:
180      if not line:  # EOF
181        break
182
183  code = await process.wait()  # wait for child process to exit
184  return code, script_output
185