blob: f236c3c45fe0504a5cc84c5d6180e61ba9e11676 [file] [log] [blame]
# Copyright (C) 2011 Google Inc. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
# * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
# * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""Module for handling messages and concurrency for run_we_tests.py
and run_blinkpy_tests.py. This module follows the design for
multiprocessing.Pool and concurrency.futures.ProcessPoolExecutor, with the
following differences:
* Tasks are executed in stateful subprocesses via objects that implement the
Worker interface - this allows the workers to share state across tasks.
* The pool provides an asynchronous event-handling interface so the caller
may receive events as tasks are processed.
If you don't need these features, use multiprocessing.Pool or concurrency.futures
instead.
"""
import logging
import multiprocessing
import six
import sys
import traceback
from six.moves import cPickle, range, queue as Queue
from blinkpy.common.host import Host
from blinkpy.common.system import stack_utils
_log = logging.getLogger(__name__)
def get(caller, worker_factory, num_workers, host=None):
"""Returns an object that exposes a run() method that takes a list of test shards and runs them in parallel."""
return _MessagePool(caller, worker_factory, num_workers, host)
class _MessagePool(object):
def __init__(self, caller, worker_factory, num_workers, host=None):
self._caller = caller
self._worker_factory = worker_factory
self._num_workers = num_workers
self._workers = []
self._workers_stopped = set()
self._host = host
self._name = 'manager'
self._running_inline = (self._num_workers == 1)
if self._running_inline:
self._messages_to_worker = Queue.Queue()
self._messages_to_manager = Queue.Queue()
else:
self._messages_to_worker = multiprocessing.Queue()
self._messages_to_manager = multiprocessing.Queue()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, exc_traceback):
self._close()
return False
def run(self, shards):
"""Posts a list of messages to the pool and waits for them to complete."""
for message in shards:
self._messages_to_worker.put(
_Message(
self._name,
message[0],
message[1:],
from_user=True,
logs=()))
for _ in range(self._num_workers):
self._messages_to_worker.put(
_Message(
self._name,
'stop',
message_args=(),
from_user=False,
logs=()))
self.wait()
def _start_workers(self):
assert not self._workers
self._workers_stopped = set()
host = None
if self._running_inline or self._can_pickle(self._host):
host = self._host
for worker_number in range(self._num_workers):
worker = _Worker(host, self._messages_to_manager,
self._messages_to_worker, self._worker_factory,
worker_number, self._running_inline,
self if self._running_inline else None,
self._worker_log_level())
self._workers.append(worker)
worker.start()
def _worker_log_level(self):
log_level = logging.NOTSET
for handler in logging.root.handlers:
if handler.level != logging.NOTSET:
if log_level == logging.NOTSET:
log_level = handler.level
else:
log_level = min(log_level, handler.level)
return log_level
def wait(self):
try:
self._start_workers()
if self._running_inline:
self._workers[0].run()
self._loop(block=False)
else:
self._loop(block=True)
finally:
self._close()
def _close(self):
for worker in self._workers:
if worker.is_alive():
worker.terminate()
worker.join()
self._workers = []
if not self._running_inline:
# FIXME: This is a hack to get multiprocessing to not log tracebacks during shutdown :(.
multiprocessing.util._exiting = True
if self._messages_to_worker:
self._messages_to_worker.close()
self._messages_to_worker = None
if self._messages_to_manager:
self._messages_to_manager.close()
self._messages_to_manager = None
def _log_messages(self, messages):
for message in messages:
logging.root.handle(message)
def _handle_done(self, source):
self._workers_stopped.add(source)
@staticmethod
def _handle_worker_exception(source, exception_type, exception_value, _):
if exception_type == KeyboardInterrupt:
raise exception_type(exception_value)
raise WorkerException(str(exception_value))
def _can_pickle(self, host):
try:
cPickle.dumps(host)
return True
except TypeError:
return False
def _loop(self, block):
try:
while True:
if len(self._workers_stopped) == len(self._workers):
block = False
message = self._messages_to_manager.get(block)
self._log_messages(message.logs)
if message.from_user:
self._caller.handle(message.name, message.src,
*message.args)
continue
method = getattr(self, '_handle_' + message.name)
assert method, 'bad message %s' % repr(message)
method(message.src, *message.args)
except Queue.Empty:
pass
class WorkerException(BaseException):
"""Raised when we receive an unexpected/unknown exception from a worker."""
class _Message(object):
def __init__(self, src, message_name, message_args, from_user, logs):
self.src = src
self.name = message_name
self.args = message_args
self.from_user = from_user
self.logs = logs
def __repr__(self):
return '_Message(src=%s, name=%s, args=%s, from_user=%s, logs=%s)' % (
self.src, self.name, self.args, self.from_user, self.logs)
class _Worker(multiprocessing.Process):
def __init__(self, host, messages_to_manager, messages_to_worker,
worker_factory, worker_number, running_inline, manager,
log_level):
super(_Worker, self).__init__()
self.host = host
self.worker_number = worker_number
self.name = 'worker/%d' % worker_number
self.log_messages = []
self.log_level = log_level
self._running = False
self._running_inline = running_inline
self._manager = manager
self._messages_to_manager = messages_to_manager
self._messages_to_worker = messages_to_worker
self._worker = worker_factory(self)
self._logger = None
self._log_handler = None
def terminate(self):
if self._worker:
if hasattr(self._worker, 'stop'):
self._worker.stop()
self._worker = None
if self.is_alive():
super(_Worker, self).terminate()
def _close(self):
if self._log_handler and self._logger:
self._logger.removeHandler(self._log_handler)
self._log_handler = None
self._logger = None
def start(self):
if not self._running_inline:
super(_Worker, self).start()
def run(self):
if not self.host:
self.host = Host()
if not self._running_inline:
self._set_up_logging()
worker = self._worker
_log.debug('%s starting', self.name)
self._running = True
try:
if hasattr(worker, 'start'):
worker.start()
while self._running:
message = self._messages_to_worker.get()
if message.from_user:
worker.handle(message.name, message.src, *message.args)
self._yield_to_manager()
else:
assert message.name == 'stop', 'bad message %s' % repr(
message)
break
_log.debug('%s exiting', self.name)
except Queue.Empty:
assert False, '%s: ran out of messages in worker queue.' % self.name
except KeyboardInterrupt:
self._raise(sys.exc_info())
except Exception:
self._raise(sys.exc_info())
finally:
try:
if hasattr(worker, 'stop'):
worker.stop()
finally:
self._post(name='done', args=(), from_user=False)
self._close()
def stop_running(self):
self._running = False
def post(self, name, *args):
self._post(name, args, from_user=True)
self._yield_to_manager()
def _yield_to_manager(self):
if self._running_inline:
self._manager._loop(block=False)
def _post(self, name, args, from_user):
log_messages = self.log_messages
self.log_messages = []
self._messages_to_manager.put(
_Message(self.name, name, args, from_user, log_messages))
def _raise(self, exc_info):
exception_type, exception_value, exception_traceback = exc_info
if self._running_inline:
six.reraise(exception_type, exception_value, exception_traceback)
if exception_type == KeyboardInterrupt:
_log.debug('%s: interrupted, exiting', self.name)
stack_utils.log_traceback(_log.debug, exception_traceback)
else:
_log.error("%s: %s('%s') raised:",
self.name, exception_value.__class__.__name__,
str(exception_value))
stack_utils.log_traceback(_log.error, exception_traceback)
# Since tracebacks aren't picklable, send the extracted stack instead.
stack = traceback.extract_tb(exception_traceback)
self._post(
name='worker_exception',
args=(exception_type, exception_value, stack),
from_user=False)
def _set_up_logging(self):
self._logger = logging.getLogger()
# The unix multiprocessing implementation clones any log handlers into the child process,
# so we remove them to avoid duplicate logging.
for h in self._logger.handlers:
self._logger.removeHandler(h)
self._log_handler = _WorkerLogHandler(self)
self._logger.addHandler(self._log_handler)
self._logger.setLevel(self.log_level)
class _WorkerLogHandler(logging.Handler):
def __init__(self, worker):
logging.Handler.__init__(self)
self._worker = worker
self.setLevel(worker.log_level)
def emit(self, record):
self._worker.log_messages.append(record)