blob: 2639183ad5a41e65a8dc6902a15fb99b1efc39cb [file] [log] [blame]
Kuang-che Wu6e4beca2018-06-27 17:45:02 +08001# -*- coding: utf-8 -*-
Kuang-che Wu88875db2017-07-20 10:47:53 +08002# Copyright 2017 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5"""Utility functions and classes."""
6
7from __future__ import print_function
Kuang-che Wu11713052019-05-30 16:21:54 +08008import difflib
Kuang-che Wu88875db2017-07-20 10:47:53 +08009import logging
Kuang-che Wu88875db2017-07-20 10:47:53 +080010import subprocess
11import threading
12import time
13
14import psutil
Kuang-che Wuae6824b2019-08-27 22:20:01 +080015from six.moves import queue
Kuang-che Wu88875db2017-07-20 10:47:53 +080016
Kuang-che Wu44278142019-03-04 11:33:57 +080017from bisect_kit import errors
18
Kuang-che Wu88875db2017-07-20 10:47:53 +080019logger = logging.getLogger(__name__)
20
21
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +080022class TimeoutExpired(Exception):
23 """Timeout expired.
24
25 This may be raised by blocking calls like Popen.wait(), check_call(),
26 check_output(), etc.
27 """
28
29
Kuang-che Wu88875db2017-07-20 10:47:53 +080030class Popen(object):
31 """Wrapper of subprocess.Popen. Support output logging.
32
33 Attributes:
34 duration: Wall time of program execution in seconds.
35 returncode: The child return code.
36 """
37
38 def __init__(self,
39 args,
40 stdout_callback=None,
41 stderr_callback=None,
Kuang-che Wubcafc552019-08-15 15:27:02 +080042 log_stdout=True,
Kuang-che Wu88875db2017-07-20 10:47:53 +080043 **kwargs):
44 """Initializes Popen.
45
46 Args:
47 args: Command line arguments.
48 stdout_callback: Callback function for stdout. Called once per line.
49 stderr_callback: Callback function for stderr. Called once per line.
Kuang-che Wubcafc552019-08-15 15:27:02 +080050 log_stdout: Whether write the stdout output of the child process to log.
Kuang-che Wu88875db2017-07-20 10:47:53 +080051 **kwargs: Additional arguments passing to subprocess.Popen.
52 """
53 if 'stdout' in kwargs:
54 raise ValueError('stdout argument not allowed, it will be overridden.')
55 if 'stderr' in kwargs:
56 raise ValueError('stderr argument not allowed, it will be overridden.')
57
58 self.stdout_callback = stdout_callback
59 self.stderr_callback = stderr_callback
Kuang-che Wubcafc552019-08-15 15:27:02 +080060 self.log_stdout = log_stdout
Kuang-che Wu88875db2017-07-20 10:47:53 +080061 self.stdout_lines = []
62 self.stderr_lines = []
63 self.duration = -1
64 self.start = time.time()
Kuang-che Wuae6824b2019-08-27 22:20:01 +080065 self.queue = queue.Queue(65536)
Kuang-che Wu88875db2017-07-20 10:47:53 +080066 if isinstance(args, str):
67 logger.debug('cwd=%s, run %r', kwargs.get('cwd'), args)
68 else:
Kuang-che Wu62dda1e2018-03-27 20:37:52 +080069 logger.debug('cwd=%s, run %r', kwargs.get('cwd'),
70 subprocess.list2cmdline(args))
Kuang-che Wu88875db2017-07-20 10:47:53 +080071 self.p = subprocess.Popen(
72 args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs)
73
74 self.stdout_thread = threading.Thread(
75 target=self._reader_thread, args=('stdout', self.p.stdout))
76 self.stdout_thread.setDaemon(True)
77 self.stdout_thread.start()
78
79 self.stderr_thread = threading.Thread(
80 target=self._reader_thread, args=('stderr', self.p.stderr))
81 self.stderr_thread.setDaemon(True)
82 self.stderr_thread.start()
83
84 @property
85 def returncode(self):
86 return self.p.returncode
87
88 def _reader_thread(self, where, child_file):
89 """Reader thread to help reading stdout and stderr.
90
91 Args:
92 where: 'stdout' or 'stderr'.
93 child_file: file object which producing output.
94 """
95 for line in iter(child_file.readline, ''):
96 self.queue.put((where, line))
97 self.queue.put((where, ''))
98
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +080099 def wait(self, timeout=None):
Kuang-che Wu88875db2017-07-20 10:47:53 +0800100 """Waits child process.
101
102 Returns:
103 return code.
104 """
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800105 t0 = time.time()
Kuang-che Wu88875db2017-07-20 10:47:53 +0800106 ended = 0
107 while ended < 2:
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800108 if timeout is not None:
109 try:
110 remaining_time = timeout - (time.time() - t0)
111 if remaining_time > 0:
112 where, line = self.queue.get(block=True, timeout=remaining_time)
113 else:
114 # We follow queue.get's behavior to raise queue.Empty, so it's
115 # always queue.Empty when time is up, no matter remaining_time is
116 # negative or positive.
117 raise queue.Empty
118 except queue.Empty:
119 logger.debug('child process time out (%.1f seconds), kill it',
120 timeout)
121 self.p.kill()
122 raise TimeoutExpired
123 else:
124 where, line = self.queue.get(block=True)
Kuang-che Wu88875db2017-07-20 10:47:53 +0800125 # line includes '\n', will be '' if EOF.
126 if not line:
127 ended += 1
128 continue
Kuang-che Wubcafc552019-08-15 15:27:02 +0800129 if self.log_stdout or where == 'stderr':
Kuang-che Wu88875db2017-07-20 10:47:53 +0800130 logger.debug('[%s] %s', where, line.rstrip('\n'))
131 if self.stdout_callback and where == 'stdout':
132 self.stdout_callback(line)
133 if self.stderr_callback and where == 'stderr':
134 self.stderr_callback(line)
135 self.p.wait()
136 self.duration = time.time() - self.start
137 logger.debug('returncode %d', self.returncode)
138 return self.returncode
139
140 def terminate(self):
141 """Terminates child and descendant processes."""
142 # Need to ignore failures because sometimes they are expected.
143 # For example, the owner of child process is different to current and
144 # unable to be killed by current process. 'cros_sdk' is one of such case.
Kuang-che Wu455c7342017-11-28 13:21:32 +0800145 for proc in psutil.Process(self.p.pid).children(recursive=True):
Kuang-che Wu88875db2017-07-20 10:47:53 +0800146 try:
147 proc.terminate()
Kuang-che Wu68db08a2018-03-30 11:50:34 +0800148 except psutil.AccessDenied:
Kuang-che Wu88875db2017-07-20 10:47:53 +0800149 logger.warning('Unable to terminate pid=%d; ignore', proc.pid)
150 try:
151 self.p.terminate()
152 except OSError:
153 logger.warning('Unable to terminate pid=%d; ignore', self.p.pid)
154 time.sleep(0.1)
155 try:
156 self.p.kill()
157 except OSError:
158 logger.warning('Unable to kill pid=%d; ignore', self.p.pid)
159
160
161def call(*args, **kwargs):
162 """Run command.
163
164 Modeled after subprocess.call.
165
166 Returns:
167 Exit code of sub-process.
168 """
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800169 timeout = kwargs.get('timeout')
170 # TODO(kcwu): let current function capture this optional parameter after
171 # migrated to python3
172 if 'timeout' in kwargs:
173 del kwargs['timeout']
Kuang-che Wu88875db2017-07-20 10:47:53 +0800174 p = Popen(args, **kwargs)
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800175 return p.wait(timeout=timeout)
Kuang-che Wu88875db2017-07-20 10:47:53 +0800176
177
178def check_output(*args, **kwargs):
179 """Runs command and return output.
180
181 Modeled after subprocess.check_output.
182
183 Returns:
184 stdout string of execution.
185
186 Raises:
187 subprocess.CalledProcessError if the exit code is non-zero.
188 """
189 stdout_lines = []
190
191 def collect_stdout(line):
192 stdout_lines.append(line)
193
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800194 timeout = kwargs.get('timeout')
195 # TODO(kcwu): let current function capture this optional parameter after
196 # migrated to python3
197 if 'timeout' in kwargs:
198 del kwargs['timeout']
Kuang-che Wu88875db2017-07-20 10:47:53 +0800199 p = Popen(args, stdout_callback=collect_stdout, **kwargs)
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800200 p.wait(timeout=timeout)
Kuang-che Wu88875db2017-07-20 10:47:53 +0800201 stdout = ''.join(stdout_lines)
202 if p.returncode != 0:
203 raise subprocess.CalledProcessError(p.returncode, args, stdout)
204
205 return stdout
206
207
208def check_call(*args, **kwargs):
209 """Runs command and ensures it succeeded.
210
211 Modeled after subprocess.check_call.
212
213 Raises:
214 subprocess.CalledProcessError if the exit code is non-zero.
215 """
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800216 timeout = kwargs.get('timeout')
217 # TODO(kcwu): let current function capture this optional parameter after
218 # migrated to python3
219 if 'timeout' in kwargs:
220 del kwargs['timeout']
Kuang-che Wu88875db2017-07-20 10:47:53 +0800221 p = Popen(args, **kwargs)
Kuang-che Wuc26dcdf2019-11-01 16:30:06 +0800222 p.wait(timeout=timeout)
Kuang-che Wu88875db2017-07-20 10:47:53 +0800223 if p.returncode != 0:
224 raise subprocess.CalledProcessError(p.returncode, args)
225
226
Kuang-che Wu44278142019-03-04 11:33:57 +0800227def ssh_cmd(host, *args, **kwargs):
228 """Runs remote command using ssh.
229
230 Args:
231 host: remote host address
232 args: command and args running on the remote host
233 kwargs:
234 connect_timeout: connection timeout in seconds (int)
235
236 Raises:
237 subprocess.CalledProcessError if the exit code is non-zero.
238 """
239 cmd = ['ssh']
240 if kwargs.get('connect_timeout'):
241 cmd += ['-oConnectTimeout=%d' % kwargs['connect_timeout']]
242 cmd.append(host)
243 cmd += list(args)
244 try:
245 return check_output(*cmd)
246 except subprocess.CalledProcessError as e:
247 # ssh's own error code is 255.
248 if e.returncode == 255:
249 raise errors.SshConnectionError('ssh connection to %r failed' % host)
250 raise
251
252
Kuang-che Wu88875db2017-07-20 10:47:53 +0800253def version_key_func(v):
254 """Splits version string into components.
255
256 Split version number by '.', and convert to `int` if possible. After this
257 conversion, version numbers can be compared ordering directly. Usually this is
258 used with sort function together.
259
260 Example,
261 >>> version_key_func('1.a.3')
262 [1, 'a', 3]
263
264 Args:
265 v: version string
266
267 Returns:
268 list of int or string
269 """
270 return [int(x) if x.isdigit() else x for x in v.split('.')]
271
272
273def is_version_lesseq(a, b):
274 """Compares whether version `a` is less or equal to version `b`.
275
276 Note this only compares the numeric values component-wise. That is, '1.1' is
277 less than '2.0', but '1.1' may or may not be older than '2.0' according to
278 chromium version semantic.
279
280 Args:
281 a: version string
282 b: version string
283
284 Returns:
285 bool: True if a <= b
286 """
287 return version_key_func(a) <= version_key_func(b)
288
289
290def is_direct_relative_version(a, b):
291 r"""Determines two versions are direct-relative.
292
293 "Direct-relative" means "one is ancestor of the other".
294
295 This follows chromium and chromiumos version semantic.
296 https://www.chromium.org/developers/version-numbers
297
298 That is, [Major+1].[Minor] is a descendant of [Major+1].1, which is branched
299 from [Major+1].0, which is a child of [Major].0. Thus, [Major+1].[Minor] is
300 not direct-relative to any [Major].[Minor>0].
301
302 For example, in this chart, 3.3 is not direct-relative to 2.2.
303
304 -> 2.0 ------------------> 3.0 -------------
305 \ \
306 -> 2.1 -> 2.2 .... -> 3.1 -> 3.2 -> 3.3 ....
307
308 Args:
309 a: version string
310 b: version string
311
312 Returns:
313 bool: True if `a` and `b` are direct-relative.
314 """
315 a = version_key_func(a)
316 b = version_key_func(b)
317 assert len(a) == len(b)
318 if a > b:
319 a, b = b, a
320
321 branched = False
322 for x, y in zip(a, b):
323 if branched:
324 if x != 0:
325 return False
326 elif x != y:
327 branched = True
328
329 return True
Kuang-che Wu11713052019-05-30 16:21:54 +0800330
331
332def show_similar_candidates(key, value, candidates):
333 logger.error('incorrect %s: %r; possible candidates:', key, value)
334 if not candidates:
335 logger.error('(no candidates at all)')
336 return
337 similar_candidates = difflib.get_close_matches(value, candidates)
338 if not similar_candidates:
339 logger.error('(no similar candidates)')
340 return
341 for candidate in similar_candidates:
342 logger.error(' %s', candidate)