289 lines
8.4 KiB
Python
289 lines
8.4 KiB
Python
|
"""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()
|