celery.worker.state 源代码

"""Internal worker state (global).

This includes the currently active and reserved tasks,
statistics, and revoked tasks.
"""
import os
import platform
import shelve
import sys
import weakref
import zlib
from collections import Counter

from kombu.serialization import pickle, pickle_protocol
from kombu.utils.objects import cached_property

from celery import __version__
from celery.exceptions import WorkerShutdown, WorkerTerminate
from celery.utils.collections import LimitedSet

__all__ = (
    'SOFTWARE_INFO', 'reserved_requests', 'active_requests',
    'total_count', 'revoked', 'task_reserved', 'maybe_shutdown',
    'task_accepted', 'task_ready', 'Persistent',
)

#: Worker software/platform information.
SOFTWARE_INFO = {
    'sw_ident': 'py-celery',
    'sw_ver': __version__,
    'sw_sys': platform.system(),
}

#: maximum number of revokes to keep in memory.
REVOKES_MAX = int(os.environ.get('CELERY_WORKER_REVOKES_MAX', 50000))

#: maximum number of successful tasks to keep in memory.
SUCCESSFUL_MAX = int(os.environ.get('CELERY_WORKER_SUCCESSFUL_MAX', 1000))

#: how many seconds a revoke will be active before
#: being expired when the max limit has been exceeded.
REVOKE_EXPIRES = float(os.environ.get('CELERY_WORKER_REVOKE_EXPIRES', 10800))

#: how many seconds a successful task will be cached in memory
#: before being expired when the max limit has been exceeded.
SUCCESSFUL_EXPIRES = float(os.environ.get('CELERY_WORKER_SUCCESSFUL_EXPIRES', 10800))

#: Mapping of reserved task_id->Request.
requests = {}

#: set of all reserved :class:`~celery.worker.request.Request`'s.
reserved_requests = weakref.WeakSet()

#: set of currently active :class:`~celery.worker.request.Request`'s.
active_requests = weakref.WeakSet()

#: A limited set of successful :class:`~celery.worker.request.Request`'s.
successful_requests = LimitedSet(maxlen=SUCCESSFUL_MAX,
                                 expires=SUCCESSFUL_EXPIRES)

#: count of tasks accepted by the worker, sorted by type.
total_count = Counter()

#: count of all tasks accepted by the worker
all_total_count = [0]

#: the list of currently revoked tasks.  Persistent if ``statedb`` set.
revoked = LimitedSet(maxlen=REVOKES_MAX, expires=REVOKE_EXPIRES)

#: Mapping of stamped headers flagged for revoking.
revoked_stamps = {}

should_stop = None
should_terminate = None


def reset_state():
    requests.clear()
    reserved_requests.clear()
    active_requests.clear()
    successful_requests.clear()
    total_count.clear()
    all_total_count[:] = [0]
    revoked.clear()
    revoked_stamps.clear()




def maybe_shutdown():
    """Shutdown if flags have been set."""
    if should_terminate is not None and should_terminate is not False:
        raise WorkerTerminate(should_terminate)
    elif should_stop is not None and should_stop is not False:
        raise WorkerShutdown(should_stop)



def task_reserved(request,
                  add_request=requests.__setitem__,
                  add_reserved_request=reserved_requests.add):
    """Update global state when a task has been reserved."""
    add_request(request.id, request)
    add_reserved_request(request)




def task_accepted(request,
                  _all_total_count=None,
                  add_request=requests.__setitem__,
                  add_active_request=active_requests.add,
                  add_to_total_count=total_count.update):
    """Update global state when a task has been accepted."""
    if not _all_total_count:
        _all_total_count = all_total_count
    add_request(request.id, request)
    add_active_request(request)
    add_to_total_count({request.name: 1})
    all_total_count[0] += 1



def task_ready(request,
               successful=False,
               remove_request=requests.pop,
               discard_active_request=active_requests.discard,
               discard_reserved_request=reserved_requests.discard):
    """Update global state when a task is ready."""
    if successful:
        successful_requests.add(request.id)

    remove_request(request.id, None)
    discard_active_request(request)
    discard_reserved_request(request)


C_BENCH = os.environ.get('C_BENCH') or os.environ.get('CELERY_BENCH')
C_BENCH_EVERY = int(os.environ.get('C_BENCH_EVERY') or
                    os.environ.get('CELERY_BENCH_EVERY') or 1000)
if C_BENCH:  # pragma: no cover
    import atexit
    from time import monotonic

    from billiard.process import current_process

    from celery.utils.debug import memdump, sample_mem

    all_count = 0
    bench_first = None
    bench_start = None
    bench_last = None
    bench_every = C_BENCH_EVERY
    bench_sample = []
    __reserved = task_reserved
    __ready = task_ready

    if current_process()._name == 'MainProcess':
        @atexit.register
        def on_shutdown():
            if bench_first is not None and bench_last is not None:
                print('- Time spent in benchmark: {!r}'.format(
                    bench_last - bench_first))
                print('- Avg: {}'.format(
                    sum(bench_sample) / len(bench_sample)))
                memdump()



    def task_reserved(request):
        """Called when a task is reserved by the worker."""
        global bench_start
        global bench_first
        now = None
        if bench_start is None:
            bench_start = now = monotonic()
        if bench_first is None:
            bench_first = now

        return __reserved(request)




    def task_ready(request):
        """Called when a task is completed."""
        global all_count
        global bench_start
        global bench_last
        all_count += 1
        if not all_count % bench_every:
            now = monotonic()
            diff = now - bench_start
            print('- Time spent processing {} tasks (since first '
                  'task received): ~{:.4f}s\n'.format(bench_every, diff))
            sys.stdout.flush()
            bench_start = bench_last = now
            bench_sample.append(diff)
            sample_mem()
        return __ready(request)





class Persistent:
    """Stores worker state between restarts.

    This is the persistent data stored by the worker when
    :option:`celery worker --statedb` is enabled.

    Currently only stores revoked task id's.
    """

    storage = shelve
    protocol = pickle_protocol
    compress = zlib.compress
    decompress = zlib.decompress
    _is_open = False

    def __init__(self, state, filename, clock=None):
        self.state = state
        self.filename = filename
        self.clock = clock
        self.merge()



    def open(self):
        return self.storage.open(
            self.filename, protocol=self.protocol, writeback=True,
        )




    def merge(self):
        self._merge_with(self.db)




    def sync(self):
        self._sync_with(self.db)
        self.db.sync()




    def close(self):
        if self._is_open:
            self.db.close()
            self._is_open = False




    def save(self):
        self.sync()
        self.close()


    def _merge_with(self, d):
        self._merge_revoked(d)
        self._merge_clock(d)
        return d

    def _sync_with(self, d):
        self._revoked_tasks.purge()
        d.update({
            '__proto__': 3,
            'zrevoked': self.compress(self._dumps(self._revoked_tasks)),
            'clock': self.clock.forward() if self.clock else 0,
        })
        return d

    def _merge_clock(self, d):
        if self.clock:
            d['clock'] = self.clock.adjust(d.get('clock') or 0)

    def _merge_revoked(self, d):
        try:
            self._merge_revoked_v3(d['zrevoked'])
        except KeyError:
            try:
                self._merge_revoked_v2(d.pop('revoked'))
            except KeyError:
                pass
        # purge expired items at boot
        self._revoked_tasks.purge()

    def _merge_revoked_v3(self, zrevoked):
        if zrevoked:
            self._revoked_tasks.update(pickle.loads(self.decompress(zrevoked)))

    def _merge_revoked_v2(self, saved):
        if not isinstance(saved, LimitedSet):
            # (pre 3.0.18) used to be stored as a dict
            return self._merge_revoked_v1(saved)
        self._revoked_tasks.update(saved)

    def _merge_revoked_v1(self, saved):
        add = self._revoked_tasks.add
        for item in saved:
            add(item)

    def _dumps(self, obj):
        return pickle.dumps(obj, protocol=self.protocol)

    @property
    def _revoked_tasks(self):
        return self.state.revoked

    @cached_property
    def db(self):
        self._is_open = True
        return self.open()