| # |
| # Copyright (C) 2016 The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| # |
| """Defines WorkQueue for delegating asynchronous work to subprocesses.""" |
| from __future__ import absolute_import |
| |
| import collections |
| import logging |
| import multiprocessing |
| import os |
| import signal |
| import sys |
| import traceback |
| |
| |
| def logger(): |
| return logging.getLogger(__name__) |
| |
| |
| def worker_sigterm_handler(_signum, _frame): |
| """Raises SystemExit so atexit/finally handlers can be executed.""" |
| sys.exit() |
| |
| |
| class TaskError(Exception): |
| """An error for an exception raised in a worker process. |
| |
| Exceptions raised in the worker will not be printed by default, and will |
| also not halt execution. We catch these exceptions in the worker process |
| and pass them through the queue. Results are checked, and if the result is |
| a TaskError the TaskError is raised in the caller's process. The message |
| for the TaskError is the stack trace of the original exception, and will be |
| printed if the TaskError is not caught. |
| """ |
| def __init__(self, trace): |
| super(TaskError, self).__init__(trace) |
| |
| |
| def create_windows_process_group(): |
| import ndk.win32 |
| job = ndk.win32.CreateJobObject() |
| |
| limit_info = ndk.win32.JOBOBJECT_EXTENDED_LIMIT_INFORMATION( |
| BasicLimitInformation=ndk.win32.JOBOBJECT_BASIC_LIMIT_INFORMATION( |
| LimitFlags=ndk.win32.JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE)) |
| |
| ndk.win32.SetInformationJobObject( |
| job, ndk.win32.JobObjectExtendedLimitInformation, limit_info) |
| ndk.win32.AssignProcessToJobObject(job, ndk.win32.GetCurrentProcess()) |
| |
| |
| def assign_self_to_new_process_group(): |
| if sys.platform == 'win32': |
| return create_windows_process_group() |
| else: |
| return os.setpgrp() |
| |
| |
| def kill_process_group(group): |
| if sys.platform == 'win32': |
| import ndk.win32 |
| ndk.win32.CloseHandle(group) |
| else: |
| os.kill(0, signal.SIGTERM) |
| |
| |
| class Worker(object): |
| IDLE_STATUS = 'IDLE' |
| EXCEPTION_STATUS = 'EXCEPTION' |
| |
| def __init__(self, data, task_queue, result_queue, manager): |
| """Creates a Worker object. |
| |
| Args: |
| task_queue: A multiprocessing.Queue of Tasks to retrieve work from. |
| result_queue: A multiprocessing.Queue to push results to. |
| """ |
| self.data = data |
| self.task_queue = task_queue |
| self.result_queue = result_queue |
| # For multiprocess.Manager.Value, the type is actually ignored. |
| # https://stackoverflow.com/a/21290961/632035 |
| self._status = manager.Value('', self.IDLE_STATUS) |
| self._status_lock = manager.Lock() |
| self.process = multiprocessing.Process(target=self.main) |
| |
| @property |
| def status(self): |
| with self._status_lock: |
| return self._status.value |
| |
| @status.setter |
| def status(self, value): |
| with self._status_lock: |
| self._status.value = value |
| |
| def put_result(self, result, status): |
| with self._status_lock: |
| self._status.value = status |
| self.result_queue.put(result) |
| |
| @property |
| def pid(self): |
| return self.process.pid |
| |
| def is_alive(self): |
| return self.process.is_alive() |
| |
| def start(self): |
| self.process.start() |
| |
| def terminate(self): |
| self.process.terminate() |
| |
| def join(self, timeout=None): |
| self.process.join(timeout) |
| |
| def main(self): |
| """Main loop for worker processes.""" |
| group = assign_self_to_new_process_group() |
| signal.signal(signal.SIGTERM, worker_sigterm_handler) |
| try: |
| while True: |
| logger().debug('worker %d waiting for work', os.getpid()) |
| task = self.task_queue.get() |
| logger().debug('worker %d running task', os.getpid()) |
| result = task.run(self) |
| logger().debug('worker %d putting result', os.getpid()) |
| self.put_result(result, self.IDLE_STATUS) |
| except SystemExit: |
| pass |
| except: # pylint: disable=bare-except |
| logger().debug('worker %d raised exception', os.getpid()) |
| trace = ''.join(traceback.format_exception(*sys.exc_info())) |
| self.put_result(TaskError(trace), self.EXCEPTION_STATUS) |
| finally: |
| # multiprocessing.Process.terminate() doesn't kill our descendents. |
| signal.signal(signal.SIGTERM, signal.SIG_IGN) |
| logger().debug('worker %d killing process group', os.getpid()) |
| kill_process_group(group) |
| signal.signal(signal.SIGTERM, signal.SIG_DFL) |
| logger().debug('worker %d exiting', os.getpid()) |
| |
| |
| class Task(object): |
| """A task to be executed by a worker process.""" |
| def __init__(self, func, args, kwargs): |
| """Creates a task. |
| |
| Args: |
| func: An invocable object to be executed by a worker process. |
| args: Arguments to be passed to the task. |
| kwargs: Keyword arguments to be passed to the task. |
| """ |
| self.func = func |
| self.args = args |
| self.kwargs = kwargs |
| |
| def run(self, worker_data): |
| """Invokes the task.""" |
| return self.func(worker_data, *self.args, **self.kwargs) |
| |
| |
| class ProcessPoolWorkQueue(object): |
| """A pool of processes for executing work asynchronously.""" |
| |
| join_timeout = 8 # Timeout for join before trying SIGKILL. |
| |
| def __init__(self, num_workers=multiprocessing.cpu_count(), |
| task_queue=None, result_queue=None, worker_data=None): |
| """Creates a WorkQueue. |
| |
| Worker threads are spawned immediately and remain live until both |
| terminate() and join() are called. |
| |
| Args: |
| num_workers: Number of worker processes to spawn. |
| task_queue: multiprocessing.Queue for tasks. Allows multiple work |
| queues to share a single task queue. If None, the work queue |
| creates its own. |
| result_queue: multiprocessing.Queue for results. Allows multiple |
| work queues to share a single result queue. If None, the work |
| queue creates its own. |
| worker_data: Data to be passed to every task run by this work |
| queue. |
| """ |
| self.manager = multiprocessing.Manager() |
| |
| self.task_queue = task_queue |
| self.owns_task_queue = False |
| if task_queue is None: |
| self.task_queue = self.manager.Queue() |
| self.owns_task_queue = True |
| |
| self.result_queue = result_queue |
| self.owns_result_queue = False |
| if result_queue is None: |
| self.result_queue = self.manager.Queue() |
| self.owns_result_queue = True |
| |
| self.worker_data = worker_data |
| |
| self.workers = [] |
| # multiprocessing.JoinableQueue's join isn't able to implement |
| # finished() because it doesn't come in a non-blocking flavor. |
| self.num_tasks = 0 |
| self._spawn_workers(num_workers) |
| |
| def add_task(self, func, *args, **kwargs): |
| """Queues up a new task for execution. |
| |
| Tasks are executed in order of insertion as worker processes become |
| available. |
| |
| Args: |
| func: An invocable object to be executed by a worker process. |
| args: Arguments to be passed to the task. |
| kwargs: Keyword arguments to be passed to the task. |
| """ |
| self.task_queue.put(Task(func, args, kwargs)) |
| self.num_tasks += 1 |
| |
| def get_result(self): |
| """Gets a result from the queue, blocking until one is available.""" |
| result = self.result_queue.get() |
| if type(result) == TaskError: |
| raise result |
| self.num_tasks -= 1 |
| return result |
| |
| def terminate(self): |
| """Terminates all worker processes.""" |
| for worker in self.workers: |
| logger().debug('terminating %d', worker.pid) |
| worker.terminate() |
| |
| def join(self): |
| """Waits for all worker processes to exit.""" |
| for worker in self.workers: |
| logger().debug('joining %d', worker.pid) |
| worker.join(self.join_timeout) |
| if worker.is_alive(): |
| logger().error( |
| 'worker %d will not die; sending SIGKILL', worker.pid) |
| os.killpg(worker.pid, signal.SIGKILL) |
| worker.join() |
| self.workers = [] |
| |
| def finished(self): |
| """Returns True if all tasks have completed execution.""" |
| return self.num_tasks == 0 |
| |
| def _spawn_workers(self, num_workers): |
| """Spawns the worker processes. |
| |
| Args: |
| num_workers: Number of worker proceeses to spawn. |
| """ |
| for _ in range(num_workers): |
| worker = Worker( |
| self.worker_data, self.task_queue, self.result_queue, |
| self.manager) |
| worker.start() |
| self.workers.append(worker) |
| |
| |
| class DummyWorker(object): |
| def __init__(self, data): |
| self.data = data |
| |
| |
| class DummyWorkQueue(object): |
| """A fake WorkQueue that does not parallelize. |
| |
| Useful for debugging when trying to determine if an issue is being caused |
| by multiprocess specific behavior. |
| """ |
| # pylint: disable=unused-argument |
| def __init__(self, num_workers=None, task_queue=None, result_queue=None, |
| worker_data=None): |
| """Creates a SerialWorkQueue.""" |
| self.task_queue = collections.deque() |
| self.worker_data = worker_data |
| # pylint: enable=unused-argument |
| |
| def add_task(self, func, *args, **kwargs): |
| """Queues up a new task for execution. |
| |
| Tasks are executed when get_result is called. |
| |
| Args: |
| func: An invocable object to be executed by a worker process. |
| args: Arguments to be passed to the task. |
| kwargs: Keyword arguments to be passed to the task. |
| """ |
| self.task_queue.append(Task(func, args, kwargs)) |
| |
| def get_result(self): |
| """Executes a task and returns the result.""" |
| task = self.task_queue.popleft() |
| try: |
| return task.run(DummyWorker(self.worker_data)) |
| except: |
| trace = ''.join(traceback.format_exception(*sys.exc_info())) |
| raise TaskError(trace) |
| |
| def terminate(self): |
| """Does nothing.""" |
| pass |
| |
| def join(self): |
| """Does nothing.""" |
| pass |
| |
| @property |
| def num_tasks(self): |
| return len(self.task_queue) |
| |
| @property |
| def workers(self): |
| return [] |
| |
| def finished(self): |
| """Returns True if all tasks have completed execution.""" |
| return self.num_tasks == 0 |
| |
| |
| WorkQueue = ProcessPoolWorkQueue |
