| # 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-webkit-tests | |
| and test-webkitpy. 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 cPickle | |
| import logging | |
| import multiprocessing | |
| import Queue | |
| import sys | |
| import traceback | |
| from webkitpy.common.host import Host | |
| from webkitpy.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 xrange(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 xrange(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: | |
| raise 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) |