/circuits/core/workers.py

 1# Module:   workers
 2# Date:     6th February 2011
 3# Author:   James Mills, prologic at shortcircuit dot net dot au
 4
 5"""Workers
 6
 7Worker components used to perform "work" in independent threads or
 8processes. Worker(s) are typically used by a Pool (circuits.core.pools)
 9to create a pool of workers. Worker(s) are not registered with a Manager
10or another Component - instead they are managed by the Pool. If a Worker
11is used independently it should not be registered as it causes its
12main event handler ``_on_task`` to execute in the other thread blocking it.
13"""
14
15from threading import current_thread
16from weakref import WeakKeyDictionary
17from multiprocessing import cpu_count
18from multiprocessing.pool import ThreadPool
19from multiprocessing import Pool as ProcessPool
20
21from .events import Event
22from .handlers import handler
23from .components import BaseComponent
24
25
26DEFAULT_WORKERS = 10
27
28
29class task(Event):
30    """task Event
31
32    This Event is used to initiate a new task to be performed by a Worker
33    or a Pool of Worker(s).
34
35    :param f: The function to be executed.
36    :type  f: function
37
38    :param args: Arguments to pass to the function
39    :type  args: tuple
40
41    :param kwargs: Keyword Arguments to pass to the function
42    :type  kwargs: dict
43    """
44
45    success = True
46    failure = True
47
48    def __init__(self, f, *args, **kwargs):
49        "x.__init__(...) initializes x; see x.__class__.__doc__ for signature"
50
51        super(task, self).__init__(f, *args, **kwargs)
52
53
54class Worker(BaseComponent):
55    """A thread/process Worker Component
56
57    This Component creates a Worker (either a thread or process) which
58    when given a ``Task``, will execute the given function in the task
59    in the background in its thread/process.
60
61    :param process: True to start this Worker as a process (Thread otherwise)
62    :type process: bool
63    """
64
65    channel = "worker"
66
67    def init(self, process=False, workers=None, channel=channel):
68        if not hasattr(current_thread(), "_children"):
69            current_thread()._children = WeakKeyDictionary()
70
71        self.workers = workers or (cpu_count() if process else DEFAULT_WORKERS)
72        Pool = ProcessPool if process else ThreadPool
73        self.pool = Pool(self.workers)
74
75    @handler("stopped", "unregistered", channel="*")
76    def _on_stopped(self, event, *args):
77        if event.name == "unregistered" and args[0] is not self:
78            return
79
80        self.pool.close()
81        self.pool.join()
82
83    @handler("task")
84    def _on_task(self, f, *args, **kwargs):
85        result = self.pool.apply_async(f, args, kwargs)
86        while not result.ready():
87            yield
88        yield result.get()