__author__ = """Copyright Andy Whitcroft, Martin J. Bligh - 2006, 2007"""
|
|
import sys, os, signal, time, cPickle, logging
|
|
from autotest_lib.client.common_lib import error, utils
|
from autotest_lib.client.common_lib.cros import retry
|
|
|
# entry points that use subcommand must set this to their logging manager
|
# to get log redirection for subcommands
|
logging_manager_object = None
|
|
|
def parallel(tasklist, timeout=None, return_results=False):
|
"""
|
Run a set of predefined subcommands in parallel.
|
|
@param tasklist: A list of subcommand instances to execute.
|
@param timeout: Number of seconds after which the commands should timeout.
|
@param return_results: If True instead of an AutoServError being raised
|
on any error a list of the results|exceptions from the tasks is
|
returned. [default: False]
|
"""
|
run_error = False
|
for task in tasklist:
|
task.fork_start()
|
|
remaining_timeout = None
|
if timeout:
|
endtime = time.time() + timeout
|
|
results = []
|
for task in tasklist:
|
if timeout:
|
remaining_timeout = max(endtime - time.time(), 1)
|
try:
|
status = task.fork_waitfor(timeout=remaining_timeout)
|
except error.AutoservSubcommandError:
|
run_error = True
|
else:
|
if status != 0:
|
run_error = True
|
|
results.append(cPickle.load(task.result_pickle))
|
task.result_pickle.close()
|
|
if return_results:
|
return results
|
elif run_error:
|
message = 'One or more subcommands failed:\n'
|
for task, result in zip(tasklist, results):
|
message += 'task: %s returned/raised: %r\n' % (task, result)
|
raise error.AutoservError(message)
|
|
|
def parallel_simple(function, arglist, subdir_name_constructor=lambda x: str(x),
|
log=True, timeout=None, return_results=False):
|
"""
|
Each element in the arglist used to create a subcommand object,
|
where that arg is used both as a subdir name, and a single argument
|
to pass to "function".
|
|
We create a subcommand object for each element in the list,
|
then execute those subcommand objects in parallel.
|
|
NOTE: As an optimization, if len(arglist) == 1 a subcommand is not used.
|
|
@param function: A callable to run in parallel once per arg in arglist.
|
@param arglist: A list of arguments to be used one per subcommand
|
@param subdir_name_constructor: A function that returns a name for the
|
result sub-directory created per subcommand.
|
Signature is:
|
subdir_name_constructor(arg)
|
where arg is the argument passed to function.
|
@param log: If True, output will be written to output in a subdirectory
|
named after each subcommand's arg.
|
@param timeout: Number of seconds after which the commands should timeout.
|
@param return_results: If True instead of an AutoServError being raised
|
on any error a list of the results|exceptions from the function
|
called on each arg is returned. [default: False]
|
|
@returns None or a list of results/exceptions.
|
"""
|
if not arglist:
|
logging.warning('parallel_simple was called with an empty arglist, '
|
'did you forget to pass in a list of machines?')
|
|
# Bypass the multithreading if only one machine.
|
if len(arglist) == 1:
|
arg = arglist[0]
|
if return_results:
|
try:
|
result = function(arg)
|
except Exception, e:
|
return [e]
|
return [result]
|
else:
|
function(arg)
|
return
|
|
subcommands = []
|
for arg in arglist:
|
args = [arg]
|
subdir = subdir_name_constructor(arg) if log else None
|
subcommands.append(subcommand(function, args, subdir))
|
return parallel(subcommands, timeout, return_results=return_results)
|
|
|
class subcommand(object):
|
fork_hooks, join_hooks = [], []
|
|
def __init__(self, func, args, subdir = None):
|
# func(args) - the subcommand to run
|
# subdir - the subdirectory to log results in
|
if subdir:
|
self.subdir = os.path.abspath(subdir)
|
if not os.path.exists(self.subdir):
|
os.mkdir(self.subdir)
|
self.debug = os.path.join(self.subdir, 'debug')
|
if not os.path.exists(self.debug):
|
os.mkdir(self.debug)
|
else:
|
self.subdir = None
|
self.debug = None
|
|
self.func = func
|
self.args = args
|
self.pid = None
|
self.returncode = None
|
|
|
def __str__(self):
|
return str('subcommand(func=%s, args=%s, subdir=%s)' %
|
(self.func, self.args, self.subdir))
|
|
|
@classmethod
|
def register_fork_hook(cls, hook):
|
""" Register a function to be called from the child process after
|
forking. """
|
cls.fork_hooks.append(hook)
|
|
|
@classmethod
|
def register_join_hook(cls, hook):
|
""" Register a function to be called when from the child process
|
just before the child process terminates (joins to the parent). """
|
cls.join_hooks.append(hook)
|
|
|
def redirect_output(self):
|
if self.subdir and logging_manager_object:
|
tag = os.path.basename(self.subdir)
|
logging_manager_object.tee_redirect_debug_dir(self.debug, tag=tag)
|
|
|
def fork_start(self):
|
sys.stdout.flush()
|
sys.stderr.flush()
|
r, w = os.pipe()
|
self.returncode = None
|
self.pid = os.fork()
|
|
if self.pid: # I am the parent
|
os.close(w)
|
self.result_pickle = os.fdopen(r, 'r')
|
return
|
else:
|
os.close(r)
|
|
# We are the child from this point on. Never return.
|
signal.signal(signal.SIGTERM, signal.SIG_DFL) # clear handler
|
if self.subdir:
|
os.chdir(self.subdir)
|
self.redirect_output()
|
|
try:
|
for hook in self.fork_hooks:
|
hook(self)
|
result = self.func(*self.args)
|
os.write(w, cPickle.dumps(result, cPickle.HIGHEST_PROTOCOL))
|
exit_code = 0
|
except Exception, e:
|
logging.exception('function failed')
|
exit_code = 1
|
os.write(w, cPickle.dumps(e, cPickle.HIGHEST_PROTOCOL))
|
|
os.close(w)
|
|
try:
|
for hook in self.join_hooks:
|
hook(self)
|
finally:
|
sys.stdout.flush()
|
sys.stderr.flush()
|
os._exit(exit_code)
|
|
|
def _handle_exitstatus(self, sts):
|
"""
|
This is partially borrowed from subprocess.Popen.
|
"""
|
if os.WIFSIGNALED(sts):
|
self.returncode = -os.WTERMSIG(sts)
|
elif os.WIFEXITED(sts):
|
self.returncode = os.WEXITSTATUS(sts)
|
else:
|
# Should never happen
|
raise RuntimeError("Unknown child exit status!")
|
|
if self.returncode != 0:
|
print "subcommand failed pid %d" % self.pid
|
print "%s" % (self.func,)
|
print "rc=%d" % self.returncode
|
print
|
if self.debug:
|
stderr_file = os.path.join(self.debug, 'autoserv.stderr')
|
if os.path.exists(stderr_file):
|
for line in open(stderr_file).readlines():
|
print line,
|
print "\n--------------------------------------------\n"
|
raise error.AutoservSubcommandError(self.func, self.returncode)
|
|
|
def poll(self):
|
"""
|
This is borrowed from subprocess.Popen.
|
"""
|
if self.returncode is None:
|
try:
|
pid, sts = os.waitpid(self.pid, os.WNOHANG)
|
if pid == self.pid:
|
self._handle_exitstatus(sts)
|
except os.error:
|
pass
|
return self.returncode
|
|
|
def wait(self):
|
"""
|
This is borrowed from subprocess.Popen.
|
"""
|
if self.returncode is None:
|
pid, sts = os.waitpid(self.pid, 0)
|
self._handle_exitstatus(sts)
|
return self.returncode
|
|
|
def fork_waitfor(self, timeout=None):
|
if not timeout:
|
return self.wait()
|
else:
|
_, result = retry.timeout(self.wait, timeout_sec=timeout)
|
|
if result is None:
|
utils.nuke_pid(self.pid)
|
print "subcommand failed pid %d" % self.pid
|
print "%s" % (self.func,)
|
print "timeout after %ds" % timeout
|
print
|
result = self.wait()
|
|
return result
|
