python-gevent 1.1.2-1 / usr / lib / python2.7 / dist-packages / gevent / pool.py

# Copyright (c) 2009-2011 Denis Bilenko. See LICENSE for details.
"""
Managing greenlets in a group.

The :class:`Group` class in this module abstracts a group of running
greenlets. When a greenlet dies, it's automatically removed from the
group. All running greenlets in a group can be waited on with
:meth:`Group.join`, or all running greenlets can be killed with
:meth:`Group.kill`.

The :class:`Pool` class, which is a subclass of :class:`Group`,
provides a way to limit concurrency: its :meth:`spawn <Pool.spawn>`
method blocks if the number of greenlets in the pool has already
reached the limit, until there is a free slot.
"""

from bisect import insort_right
try:
    from itertools import izip
except ImportError:
    # Python 3
    izip = zip

from gevent.hub import GreenletExit, getcurrent, kill as _kill
from gevent.greenlet import joinall, Greenlet
from gevent.timeout import Timeout
from gevent.event import Event
from gevent.lock import Semaphore, DummySemaphore

__all__ = ['Group', 'Pool']


class IMapUnordered(Greenlet):
    """
    At iterator of map results.
    """

    _zipped = False

    def __init__(self, func, iterable, spawn=None, maxsize=None, _zipped=False):
        """
        An iterator that.

        :keyword int maxsize: If given and not-None, specifies the maximum number of
            finished results that will be allowed to accumulated awaiting the reader;
            more than that number of results will cause map function greenlets to begin
            to block. This is most useful is there is a great disparity in the speed of
            the mapping code and the consumer and the results consume a great deal of resources.
            Using a bound is more computationally expensive than not using a bound.

        .. versionchanged:: 1.1b3
            Added the *maxsize* parameter.
        """
        from gevent.queue import Queue
        Greenlet.__init__(self)
        if spawn is not None:
            self.spawn = spawn
        if _zipped:
            self._zipped = _zipped
        self.func = func
        self.iterable = iterable
        self.queue = Queue()
        if maxsize:
            # Bounding the queue is not enough if we want to keep from
            # accumulating objects; the result value will be around as
            # the greenlet's result, blocked on self.queue.put(), and
            # we'll go on to spawn another greenlet, which in turn can
            # create the result. So we need a semaphore to prevent a
            # greenlet from exiting while the queue is full so that we
            # don't spawn the next greenlet (assuming that self.spawn
            # is of course bounded). (Alternatively we could have the
            # greenlet itself do the insert into the pool, but that
            # takes some rework).
            #
            # Given the use of a semaphore at this level, sizing the queue becomes
            # redundant, and that lets us avoid having to use self.link() instead
            # of self.rawlink() to avoid having blocking methods called in the
            # hub greenlet.
            self._result_semaphore = Semaphore(maxsize)
        else:
            self._result_semaphore = DummySemaphore()
        self.count = 0
        self.finished = False
        # If the queue size is unbounded, then we want to call all
        # the links (_on_finish and _on_result) directly in the hub greenlet
        # for efficiency. However, if the queue is bounded, we can't do that if
        # the queue might block (because if there's no waiter the hub can switch to,
        # the queue simply raises Full). Therefore, in that case, we use
        # the safer, somewhat-slower (because it spawns a greenlet) link() methods.
        # This means that _on_finish and _on_result can be called and interleaved in any order
        # if the call to self.queue.put() blocks..
        # Note that right now we're not bounding the queue, instead using a semaphore.
        self.rawlink(self._on_finish)

    def __iter__(self):
        return self

    def next(self):
        self._result_semaphore.release()
        value = self._inext()
        if isinstance(value, Failure):
            raise value.exc
        return value
    __next__ = next

    def _inext(self):
        return self.queue.get()

    def _ispawn(self, func, item):
        self._result_semaphore.acquire()
        self.count += 1
        g = self.spawn(func, item) if not self._zipped else self.spawn(func, *item)
        g.rawlink(self._on_result)
        return g

    def _run(self):
        try:
            func = self.func
            for item in self.iterable:
                self._ispawn(func, item)
        finally:
            self.__dict__.pop('spawn', None)
            self.__dict__.pop('func', None)
            self.__dict__.pop('iterable', None)

    def _on_result(self, greenlet):
        # This method can either be called in the hub greenlet (if the
        # queue is unbounded) or its own greenlet. If it's called in
        # its own greenlet, the calls to put() may block and switch
        # greenlets, which in turn could mutate our state. So any
        # state on this object that we need to look at, notably
        # self.count, we need to capture or mutate *before* we put.
        # (Note that right now we're not bounding the queue, but we may
        # choose to do so in the future so this implementation will be left in case.)
        self.count -= 1
        count = self.count
        finished = self.finished
        ready = self.ready()
        put_finished = False

        if ready and count <= 0 and not finished:
            finished = self.finished = True
            put_finished = True

        if greenlet.successful():
            self.queue.put(self._iqueue_value_for_success(greenlet))
        else:
            self.queue.put(self._iqueue_value_for_failure(greenlet))

        if put_finished:
            self.queue.put(self._iqueue_value_for_finished())

    def _on_finish(self, _self):
        if self.finished:
            return

        if not self.successful():
            self.finished = True
            self.queue.put(self._iqueue_value_for_self_failure())
            return

        if self.count <= 0:
            self.finished = True
            self.queue.put(self._iqueue_value_for_finished())

    def _iqueue_value_for_success(self, greenlet):
        return greenlet.value

    def _iqueue_value_for_failure(self, greenlet):
        return Failure(greenlet.exception, getattr(greenlet, '_raise_exception'))

    def _iqueue_value_for_finished(self):
        return Failure(StopIteration)

    def _iqueue_value_for_self_failure(self):
        return Failure(self.exception, self._raise_exception)


class IMap(IMapUnordered):
    # A specialization of IMapUnordered that returns items
    # in the order in which they were generated, not
    # the order in which they finish.
    # We do this by storing tuples (order, value) in the queue
    # not just value.

    def __init__(self, *args, **kwargs):
        self.waiting = []  # QQQ maybe deque will work faster there?
        self.index = 0
        self.maxindex = -1
        IMapUnordered.__init__(self, *args, **kwargs)

    def _inext(self):
        while True:
            if self.waiting and self.waiting[0][0] <= self.index:
                _, value = self.waiting.pop(0)
            else:
                index, value = self.queue.get()
                if index > self.index:
                    insort_right(self.waiting, (index, value))
                    continue
            self.index += 1
            return value

    def _ispawn(self, func, item):
        g = IMapUnordered._ispawn(self, func, item)
        self.maxindex += 1
        g.index = self.maxindex
        return g

    def _iqueue_value_for_success(self, greenlet):
        return (greenlet.index, IMapUnordered._iqueue_value_for_success(self, greenlet))

    def _iqueue_value_for_failure(self, greenlet):
        return (greenlet.index, IMapUnordered._iqueue_value_for_failure(self, greenlet))

    def _iqueue_value_for_finished(self):
        self.maxindex += 1
        return (self.maxindex, IMapUnordered._iqueue_value_for_finished(self))

    def _iqueue_value_for_self_failure(self):
        self.maxindex += 1
        return (self.maxindex, IMapUnordered._iqueue_value_for_self_failure(self))


class GroupMappingMixin(object):
    # Internal, non-public API class.
    # Provides mixin methods for implementing mapping pools. Subclasses must define:

    # - self.spawn(func, *args, **kwargs): a function that runs `func` with `args`
    # and `awargs`, potentially asynchronously. Return a value with a `get` method that
    # blocks until the results of func are available, and a `link` method.

    # - self._apply_immediately(): should the function passed to apply be called immediately,
    # synchronously?

    # - self._apply_async_use_greenlet(): Should apply_async directly call
    # Greenlet.spawn(), bypassing self.spawn? Return true when self.spawn would block

    # - self._apply_async_cb_spawn(callback, result): Run the given callback function, possiblly
    # asynchronously, possibly synchronously.

    def apply_cb(self, func, args=None, kwds=None, callback=None):
        """
        :meth:`apply` the given *func*, and, if a *callback* is given, run it with the
        results of *func* (unless an exception was raised.)

        The *callback* may be called synchronously or asynchronously. If called
        asynchronously, it will not be tracked by this group. (:class:`Group` and :class:`Pool`
        call it asynchronously in a new greenlet; :class:`~gevent.threadpool.ThreadPool` calls
        it synchronously in the current greenlet.)
        """
        result = self.apply(func, args, kwds)
        if callback is not None:
            self._apply_async_cb_spawn(callback, result)
        return result

    def apply_async(self, func, args=None, kwds=None, callback=None):
        """
        A variant of the apply() method which returns a Greenlet object.

        If *callback* is specified, then it should be a callable which
        accepts a single argument. When the result becomes ready
        callback is applied to it (unless the call failed).
        """
        if args is None:
            args = ()
        if kwds is None:
            kwds = {}
        if self._apply_async_use_greenlet():
            # cannot call spawn() directly because it will block
            return Greenlet.spawn(self.apply_cb, func, args, kwds, callback)

        greenlet = self.spawn(func, *args, **kwds)
        if callback is not None:
            greenlet.link(pass_value(callback))
        return greenlet

    def apply(self, func, args=None, kwds=None):
        """
        Rough quivalent of the :func:`apply()` builtin function blocking until
        the result is ready and returning it.

        The ``func`` will *usually*, but not *always*, be run in a way
        that allows the current greenlet to switch out (for example,
        in a new greenlet or thread, depending on implementation). But
        if the current greenlet or thread is already one that was
        spawned by this pool, the pool may choose to immediately run
        the `func` synchronously.

        Any exception ``func`` raises will be propagated to the caller of ``apply`` (that is,
        this method will raise the exception that ``func`` raised).
        """
        if args is None:
            args = ()
        if kwds is None:
            kwds = {}
        if self._apply_immediately():
            return func(*args, **kwds)
        else:
            return self.spawn(func, *args, **kwds).get()

    def map(self, func, iterable):
        """Return a list made by applying the *func* to each element of
        the iterable.

        .. seealso:: :meth:`imap`
        """
        return list(self.imap(func, iterable))

    def map_cb(self, func, iterable, callback=None):
        result = self.map(func, iterable)
        if callback is not None:
            callback(result)
        return result

    def map_async(self, func, iterable, callback=None):
        """
        A variant of the map() method which returns a Greenlet object that is executing
        the map function.

        If callback is specified then it should be a callable which accepts a
        single argument.
        """
        return Greenlet.spawn(self.map_cb, func, iterable, callback)

    def __imap(self, cls, func, *iterables, **kwargs):
        # Python 2 doesn't support the syntax that lets us mix varargs and
        # a named kwarg, so we have to unpack manually
        maxsize = kwargs.pop('maxsize', None)
        if kwargs:
            raise TypeError("Unsupported keyword arguments")
        return cls.spawn(func, izip(*iterables), spawn=self.spawn,
                         _zipped=True, maxsize=maxsize)

    def imap(self, func, *iterables, **kwargs):
        """
        imap(func, *iterables, maxsize=None) -> iterable

        An equivalent of :func:`itertools.imap`, operating in parallel.
        The *func* is applied to each element yielded from each
        iterable in *iterables* in turn, collecting the result.

        If this object has a bound on the number of active greenlets it can
        contain (such as :class:`Pool`), then at most that number of tasks will operate
        in parallel.

        :keyword int maxsize: If given and not-None, specifies the maximum number of
            finished results that will be allowed to accumulate awaiting the reader;
            more than that number of results will cause map function greenlets to begin
            to block. This is most useful is there is a great disparity in the speed of
            the mapping code and the consumer and the results consume a great deal of resources.

            .. note:: This is separate from any bound on the number of active parallel
               tasks, though they may have some interaction (for example, limiting the
               number of parallel tasks to the smallest bound).

            .. note:: Using a bound is slightly more computationally expensive than not using a bound.

            .. tip:: The :meth:`imap_unordered` method makes much better
                use of this parameter. Some additional, unspecified,
                number of objects may be required to be kept in memory
                to maintain order by this function.

        :return: An iterable object.

        .. versionchanged:: 1.1b3
            Added the *maxsize* keyword parameter.
        .. versionchanged:: 1.1a1
            Accept multiple *iterables* to iterate in parallel.
        """
        return self.__imap(IMap, func, *iterables, **kwargs)

    def imap_unordered(self, func, *iterables, **kwargs):
        """
        imap_unordered(func, *iterables, maxsize=None) -> iterable

        The same as :meth:`imap` except that the ordering of the results
        from the returned iterator should be considered in arbitrary
        order.

        This is lighter weight than :meth:`imap` and should be preferred if order
        doesn't matter.

        .. seealso:: :meth:`imap` for more details.
        """
        return self.__imap(IMapUnordered, func, *iterables, **kwargs)


class Group(GroupMappingMixin):
    """
    Maintain a group of greenlets that are still running, without
    limiting their number.

    Links to each item and removes it upon notification.

    Groups can be iterated to discover what greenlets they are tracking,
    they can be tested to see if they contain a greenlet, and they know the
    number (len) of greenlets they are tracking. If they are not tracking any
    greenlets, they are False in a boolean context.
    """

    #: The type of Greenlet object we will :meth:`spawn`. This can be changed
    #: on an instance or in a subclass.
    greenlet_class = Greenlet

    def __init__(self, *args):
        assert len(args) <= 1, args
        self.greenlets = set(*args)
        if args:
            for greenlet in args[0]:
                greenlet.rawlink(self._discard)
        # each item we kill we place in dying, to avoid killing the same greenlet twice
        self.dying = set()
        self._empty_event = Event()
        self._empty_event.set()

    def __repr__(self):
        return '<%s at 0x%x %s>' % (self.__class__.__name__, id(self), self.greenlets)

    def __len__(self):
        """
        Answer how many greenlets we are tracking. Note that if we are empty,
        we are False in a boolean context.
        """
        return len(self.greenlets)

    def __contains__(self, item):
        """
        Answer if we are tracking the given greenlet.
        """
        return item in self.greenlets

    def __iter__(self):
        """
        Iterate across all the greenlets we are tracking, in no particular order.
        """
        return iter(self.greenlets)

    def add(self, greenlet):
        """
        Begin tracking the greenlet.

        If this group is :meth:`full`, then this method may block
        until it is possible to track the greenlet.
        """
        try:
            rawlink = greenlet.rawlink
        except AttributeError:
            pass  # non-Greenlet greenlet, like MAIN
        else:
            rawlink(self._discard)
        self.greenlets.add(greenlet)
        self._empty_event.clear()

    def _discard(self, greenlet):
        self.greenlets.discard(greenlet)
        self.dying.discard(greenlet)
        if not self.greenlets:
            self._empty_event.set()

    def discard(self, greenlet):
        """
        Stop tracking the greenlet.
        """
        self._discard(greenlet)
        try:
            unlink = greenlet.unlink
        except AttributeError:
            pass  # non-Greenlet greenlet, like MAIN
        else:
            unlink(self._discard)

    def start(self, greenlet):
        """
        Start the un-started *greenlet* and add it to the collection of greenlets
        this group is monitoring.
        """
        self.add(greenlet)
        greenlet.start()

    def spawn(self, *args, **kwargs):
        """
        Begin a new greenlet with the given arguments (which are passed
        to the greenlet constructor) and add it to the collection of greenlets
        this group is monitoring.

        :return: The newly started greenlet.
        """
        greenlet = self.greenlet_class(*args, **kwargs)
        self.start(greenlet)
        return greenlet

#     def close(self):
#         """Prevents any more tasks from being submitted to the pool"""
#         self.add = RaiseException("This %s has been closed" % self.__class__.__name__)

    def join(self, timeout=None, raise_error=False):
        """
        Wait for this group to become empty *at least once*.

        If there are no greenlets in the group, returns immediately.

        .. note:: By the time the waiting code (the caller of this
           method) regains control, a greenlet may have been added to
           this group, and so this object may no longer be empty. (That
           is, ``group.join(); assert len(group) == 0`` is not
           guaranteed to hold.) This method only guarantees that the group
           reached a ``len`` of 0 at some point.

        :keyword bool raise_error: If True (*not* the default), if any
            greenlet that finished while the join was in progress raised
            an exception, that exception will be raised to the caller of
            this method. If multiple greenlets raised exceptions, which
            one gets re-raised is not determined. Only greenlets currently
            in the group when this method is called are guaranteed to
            be checked for exceptions.
        """
        if raise_error:
            greenlets = self.greenlets.copy()
            self._empty_event.wait(timeout=timeout)
            for greenlet in greenlets:
                if greenlet.exception is not None:
                    if hasattr(greenlet, '_raise_exception'):
                        greenlet._raise_exception()
                    raise greenlet.exception
        else:
            self._empty_event.wait(timeout=timeout)

    def kill(self, exception=GreenletExit, block=True, timeout=None):
        """
        Kill all greenlets being tracked by this group.
        """
        timer = Timeout._start_new_or_dummy(timeout)
        try:
            try:
                while self.greenlets:
                    for greenlet in list(self.greenlets):
                        if greenlet not in self.dying:
                            try:
                                kill = greenlet.kill
                            except AttributeError:
                                _kill(greenlet, exception)
                            else:
                                kill(exception, block=False)
                            self.dying.add(greenlet)
                    if not block:
                        break
                    joinall(self.greenlets)
            except Timeout as ex:
                if ex is not timer:
                    raise
        finally:
            timer.cancel()

    def killone(self, greenlet, exception=GreenletExit, block=True, timeout=None):
        """
        If the given *greenlet* is running and being tracked by this group,
        kill it.
        """
        if greenlet not in self.dying and greenlet in self.greenlets:
            greenlet.kill(exception, block=False)
            self.dying.add(greenlet)
            if block:
                greenlet.join(timeout)

    def full(self):
        """
        Return a value indicating whether this group can track more greenlets.

        In this implementation, because there are no limits on the number of
        tracked greenlets, this will always return a ``False`` value.
        """
        return False

    def wait_available(self, timeout=None):
        """
        Block until it is possible to :meth:`spawn` a new greenlet.

        In this implementation, because there are no limits on the number
        of tracked greenlets, this will always return immediately.
        """
        pass

    # MappingMixin methods

    def _apply_immediately(self):
        # If apply() is called from one of our own
        # worker greenlets, don't spawn a new one
        return getcurrent() in self

    def _apply_async_cb_spawn(self, callback, result):
        Greenlet.spawn(callback, result)

    def _apply_async_use_greenlet(self):
        return self.full() # cannot call self.spawn() because it will block


class Failure(object):
    __slots__ = ['exc', '_raise_exception']

    def __init__(self, exc, raise_exception=None):
        self.exc = exc
        self._raise_exception = raise_exception

    def raise_exc(self):
        if self._raise_exception:
            self._raise_exception()
        else:
            raise self.exc


class Pool(Group):

    def __init__(self, size=None, greenlet_class=None):
        """
        Create a new pool.

        A pool is like a group, but the maximum number of members
        is governed by the *size* parameter.

        :keyword int size: If given, this non-negative integer is the
            maximum count of active greenlets that will be allowed in
            this pool. A few values have special significance:

            * ``None`` (the default) places no limit on the number of
              greenlets. This is useful when you need to track, but not limit,
              greenlets, as with :class:`gevent.pywsgi.WSGIServer`. A :class:`Group`
              may be a more efficient way to achieve the same effect.
            * ``0`` creates a pool that can never have any active greenlets. Attempting
              to spawn in this pool will block forever. This is only useful
              if an application uses :meth:`wait_available` with a timeout and checks
              :meth:`free_count` before attempting to spawn.
        """
        if size is not None and size < 0:
            raise ValueError('size must not be negative: %r' % (size, ))
        Group.__init__(self)
        self.size = size
        if greenlet_class is not None:
            self.greenlet_class = greenlet_class
        if size is None:
            self._semaphore = DummySemaphore()
        else:
            self._semaphore = Semaphore(size)

    def wait_available(self, timeout=None):
        """
        Wait until it's possible to spawn a greenlet in this pool.

        :param float timeout: If given, only wait the specified number
            of seconds.

        .. warning:: If the pool was initialized with a size of 0, this
           method will block forever unless a timeout is given.

        :return: A number indicating how many new greenlets can be put into
           the pool without blocking.

        .. versionchanged:: 1.1a3
            Added the ``timeout`` parameter.
        """
        return self._semaphore.wait(timeout=timeout)

    def full(self):
        """
        Return a boolean indicating whether this pool has any room for
        members. (True if it does, False if it doesn't.)
        """
        return self.free_count() <= 0

    def free_count(self):
        """
        Return a number indicating *approximately* how many more members
        can be added to this pool.
        """
        if self.size is None:
            return 1
        return max(0, self.size - len(self))

    def add(self, greenlet):
        """
        Begin tracking the given greenlet, blocking until space is available.

        .. seealso:: :meth:`Group.add`
        """
        self._semaphore.acquire()
        try:
            Group.add(self, greenlet)
        except:
            self._semaphore.release()
            raise

    def _discard(self, greenlet):
        Group._discard(self, greenlet)
        self._semaphore.release()


class pass_value(object):
    __slots__ = ['callback']

    def __init__(self, callback):
        self.callback = callback

    def __call__(self, source):
        if source.successful():
            self.callback(source.value)

    def __hash__(self):
        return hash(self.callback)

    def __eq__(self, other):
        return self.callback == getattr(other, 'callback', other)

    def __str__(self):
        return str(self.callback)

    def __repr__(self):
        return repr(self.callback)

    def __getattr__(self, item):
        assert item != 'callback'
        return getattr(self.callback, item)