def exclusive_lock(task_identifier):
""" Obtain an exclusively lock, using the task_identifier as a unique ID.
This helps prevents the case of multiple workers executing the same task at
the same time, which can cause unexpected side effects.
"""
# See: https://celery.readthedocs.io/en/latest/tutorials/task-cookbook.html
# Plan to timeout a few seconds before the limit
# (After `LOCK_EXPIRE` seconds have passed, the cache will self-clean)
timeout_at = monotonic() + LOCK_EXPIRE - 3
# Try to add the value to the cache.
# Returns False if already cached (another worker added it already)
# Returns True otherwise (this worker is the first to add the key)
got_lock = cache.add(task_identifier, 'true', LOCK_EXPIRE)
# Yield our ability to obtain a lock, but always cleanup
try:
yield got_lock
finally:
# If `got_lock` was False, we don't own the lock - never clean up
# If we're close to the timeout, just let the cache self-clean
if got_lock and monotonic() < timeout_at:
cache.delete(task_identifier)
def test_purge(self):
# purge now enforces rules
# cant purge(1) now. but .purge(now=...) still works
s = LimitedSet(maxlen=10)
[s.add(i) for i in range(10)]
s.maxlen = 2
s.purge()
assert len(s) == 2
# expired
s = LimitedSet(maxlen=10, expires=1)
[s.add(i) for i in range(10)]
s.maxlen = 2
s.purge(now=monotonic() + 100)
assert len(s) == 0
# not expired
s = LimitedSet(maxlen=None, expires=1)
[s.add(i) for i in range(10)]
s.maxlen = 2
s.purge(now=lambda: monotonic() - 100)
assert len(s) == 2
# expired -> minsize
s = LimitedSet(maxlen=10, minlen=10, expires=1)
[s.add(i) for i in range(20)]
s.minlen = 3
s.purge(now=monotonic() + 3)
assert s.minlen == len(s)
assert len(s._heap) <= s.maxlen * (
100. + s.max_heap_percent_overload) / 100
def pstatus(p):
runtime = monotonic() - p.runtime
elapsed = monotonic() - p.elapsed
return F_PROGRESS.format(
p,
runtime=humanize_seconds(runtime, now=runtime),
elapsed=humanize_seconds(elapsed, now=elapsed),
)
def pstatus(p, status=None):
runtime = format(monotonic() - p.runtime, '.4f')
elapsed = format(monotonic() - p.elapsed, '.4f')
return F_PROGRESS.format(
p,
status=status or '',
runtime=humanize_seconds(runtime, now=runtime),
elapsed=humanize_seconds(elapsed, now=elapsed),
)
def test_on_accepted_acks_early(self):
job = self.xRequest()
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
assert job.acknowledged
prev, module._does_debug = module._does_debug, False
try:
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
finally:
module._does_debug = prev
def test_on_accepted_acks_early(self):
job = TaskRequest(
self.mytask.name, uuid(), [1], {'f': 'x'}, app=self.app,
)
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
self.assertTrue(job.acknowledged)
prev, module._does_debug = module._does_debug, False
try:
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
finally:
module._does_debug = prev
def handle_worker(self, hostname_worker):
(hostname, worker) = hostname_worker
last_write, obj = self._last_worker_write[hostname]
if not last_write or \
monotonic() - last_write > self.worker_update_freq:
obj = self.WorkerState.objects.update_or_create(
hostname=hostname,
defaults={'last_heartbeat': self.get_heartbeat(worker)},
)
self._last_worker_write[hostname] = (monotonic(), obj)
return obj
def task_lock(lock_id, oid):
timeout_at = monotonic() + LOCK_EXPIRE - 3
# cache.add fails if the key already exists
status = cache.add(lock_id, oid, LOCK_EXPIRE)
try:
yield status
finally:
# memcache delete is very slow, but we have to use it to take
# advantage of using add() for atomic locking
if monotonic() < timeout_at and status:
# don't release the lock if we exceeded the timeout
# to lessen the chance of releasing an expired lock
# owned by someone else
# also don't release the lock if we didn't acquire it
cache.delete(lock_id)
def _retry_on_error(self, fun, *args, **kwargs):
ts = monotonic() + self._retry_timeout
while 1:
try:
return fun(*args, **kwargs)
except (pycassa.InvalidRequestException,
pycassa.TimedOutException,
pycassa.UnavailableException,
pycassa.AllServersUnavailable,
socket.error,
socket.timeout,
Thrift.TException) as exc:
if monotonic() > ts:
raise
logger.warning('Cassandra error: %r. Retrying...', exc)
time.sleep(self._retry_wait)
def test_shrink_raises_exception(self):
worker = Mock(name='worker')
x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
x.scale_up(3)
x._last_action = monotonic() - 10000
x.pool.shrink_raises_exception = True
x._shrink(1)
def test_should_sync(self):
self.assertTrue(self.s.should_sync())
self.s._last_sync = monotonic()
self.s._tasks_since_sync = 0
self.assertFalse(self.s.should_sync())
self.s._last_sync -= self.s.sync_every + 1
self.assertTrue(self.s.should_sync())
def drain_events_until(self, p, timeout=None, on_interval=None, wait=None):
wait = wait or self.result_consumer.drain_events
time_start = monotonic()
while 1:
# Total time spent may exceed a single call to wait()
if timeout and monotonic() - time_start >= timeout:
raise socket.timeout()
try:
yield self.wait_for(p, wait, timeout=1)
except socket.timeout:
pass
if on_interval:
on_interval()
if p.ready: # got event on the wanted channel.
break
def test_periodic_task_disabled_and_enabled(self):
# Get the entry for m2
e1 = self.s.schedule[self.m2.name]
# Increment the entry (but make sure it doesn't sync)
self.s._last_sync = monotonic()
self.s.schedule[e1.name] = self.s.reserve(e1)
assert self.s.flushed == 1
# Fetch the raw object from db, change the args
# and save the changes.
m2 = PeriodicTask.objects.get(pk=self.m2.pk)
m2.enabled = False
m2.save()
# get_schedule should now see the schedule has changed.
# and remove entry for m2
assert self.m2.name not in self.s.schedule
assert self.s.flushed == 2
m2.enabled = True
m2.save()
# get_schedule should now see the schedule has changed.
# and add entry for m2
assert self.m2.name in self.s.schedule
assert self.s.flushed == 3
def test_on_accepted_acks_late(self):
job = TaskRequest(
self.mytask.name, uuid(), [1], {'f': 'x'}, app=self.app,
)
self.mytask.acks_late = True
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
self.assertFalse(job.acknowledged)
def test_shrink_raises_ValueError(self, debug):
worker = Mock(name='worker')
x = autoscale.Autoscaler(self.pool, 10, 3, worker=worker)
x.scale_up(3)
x._last_scale_up = monotonic() - 10000
x.pool.shrink_raises_ValueError = True
x.scale_down(1)
self.assertTrue(debug.call_count)
def task_reserved(request): # noqa
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 _enter(self, eta, priority, entry):
secs = max(eta - monotonic(), 0)
g = self._Greenlet.spawn_later(secs, entry)
self._queue.add(g)
g.link(self._entry_exit)
g.entry = entry
g.eta = eta
g.priority = priority
g.cancelled = False
return g
def runtest(self, fun, n=50, index=0, repeats=1):
n = getattr(fun, '__iterations__', None) or n
print('{0}: [[[{1}({2})]]]'.format(repeats, fun.__name__, n))
with blockdetection(self.block_timeout):
with self.fbi.investigation():
runtime = elapsed = monotonic()
i = 0
failed = False
self.progress = Progress(
fun, i, n, index, repeats, elapsed, runtime, 0,
)
_marker.delay(pstatus(self.progress))
try:
for i in range(n):
runtime = monotonic()
self.progress = Progress(
fun, i + 1, n, index, repeats, runtime, elapsed, 0,
)
try:
fun()
except StopSuite:
raise
except Exception as exc:
print('-> {0!r}'.format(exc))
import traceback
print(traceback.format_exc())
print(pstatus(self.progress))
else:
print(pstatus(self.progress))
except Exception:
failed = True
self.speaker.beep()
raise
finally:
print('{0} {1} iterations in {2}'.format(
'failed after' if failed else 'completed',
i + 1, humanize_seconds(monotonic() - elapsed),
))
if not failed:
self.progress = Progress(
fun, i + 1, n, index, repeats, runtime, elapsed, 1,
)
def task_reserved(request): # noqa
"""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 test_on_accepted_terminates(self):
signum = signal.SIGTERM
pool = Mock()
job = self.xRequest()
with self.assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.terminate(pool, signal='TERM')
assert not pool.terminate_job.call_count
job.on_accepted(pid=314, time_accepted=monotonic())
pool.terminate_job.assert_called_with(314, signum)
def test_terminate__task_started(self):
pool = Mock()
signum = signal.SIGTERM
job = self.get_request(self.mytask.s(1, f='x'))
with self.assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.time_start = monotonic()
job.worker_pid = 313
job.terminate(pool, signal='TERM')
pool.terminate_job.assert_called_with(job.worker_pid, signum)
def test_terminate__task_started(self):
pool = Mock()
signum = signal.SIGKILL
job = Request({"task": self.mytask.name, "id": uuid(), "args": [1], "kwrgs": {"f": "x"}}, app=self.app)
with assert_signal_called(
task_revoked, sender=job.task, request=job, terminated=True, expired=False, signum=signum
):
job.time_start = monotonic()
job.worker_pid = 313
job.terminate(pool, signal="KILL")
pool.terminate_job.assert_called_with(job.worker_pid, signum)
def add(self, item, now=None):
# type: (Any, float) -> None
"""Add a new item, or reset the expiry time of an existing item."""
now = now or monotonic()
if item in self._data:
self.discard(item)
entry = (now, item)
self._data[item] = entry
heappush(self._heap, entry)
if self.maxlen and len(self._data) >= self.maxlen:
self.purge()
def on_accepted(self, pid, time_accepted):
"""Handler called when task is accepted by worker pool."""
self.worker_pid = pid
# Convert monotonic time_accepted to absolute time
self.time_start = time() - (monotonic() - time_accepted)
task_accepted(self)
if not self.task.acks_late:
self.acknowledge()
self.send_event('task-started')
if _does_debug:
debug('Task accepted: %s[%s] pid:%r', self.name, self.id, pid)
if self._terminate_on_ack is not None:
self.terminate(*self._terminate_on_ack)
def test_on_accepted_terminates(self):
signum = signal.SIGKILL
pool = Mock()
job = TaskRequest(
self.mytask.name, uuid(), [1], {'f': 'x'}, app=self.app,
)
with assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.terminate(pool, signal='KILL')
self.assertFalse(pool.terminate_job.call_count)
job.on_accepted(pid=314, time_accepted=monotonic())
pool.terminate_job.assert_called_with(314, signum)
def test_terminate__task_started(self):
pool = Mock()
signum = signal.SIGKILL
job = TaskRequest(
self.mytask.name, uuid(), [1], {'f': 'x'}, app=self.app,
)
with assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.time_start = monotonic()
job.worker_pid = 313
job.terminate(pool, signal='KILL')
pool.terminate_job.assert_called_with(job.worker_pid, signum)
def runtest(self, fun, n=50, index=0, repeats=1):
n = getattr(fun, '__iterations__', None) or n
header = '[[[{0}({1})]]]'.format(fun.__name__, n)
if repeats > 1:
header = '{0} #{1}'.format(header, repeats)
self.print(header)
with blockdetection(self.block_timeout):
with self.fbi.investigation():
runtime = elapsed = monotonic()
i = 0
failed = False
self.progress = Progress(
fun, i, n, index, repeats, elapsed, runtime, 0,
)
_marker.delay(pstatus(self.progress))
try:
for i in range(n):
runtime = monotonic()
self.progress = Progress(
fun, i + 1, n, index, repeats, runtime, elapsed, 0,
)
self.execute_test(fun)
except Exception:
failed = True
self.speaker.beep()
raise
finally:
if n > 1 or failed:
self.print('{0} {1} iterations in {2}'.format(
'failed after' if failed else 'completed',
i + 1, humanize_seconds(monotonic() - elapsed),
), file=self.stderr if failed else self.stdout)
if not failed:
self.progress = Progress(
fun, i + 1, n, index, repeats, runtime, elapsed, 1,
)
def task_ready(request): # noqa
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 {0} tasks (since first '
'task received): ~{1:.4f}s\n'.format(bench_every, diff))
sys.stdout.flush()
bench_start = bench_last = now
bench_sample.append(diff)
sample_mem()
return __ready(request)
def test_terminate__pool_ref(self):
pool = Mock()
signum = signal.SIGTERM
job = self.get_request(self.mytask.s(1, f='x'))
job._apply_result = Mock(name='_apply_result')
with self.assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.time_start = monotonic()
job.worker_pid = 314
job.terminate(pool, signal='TERM')
job._apply_result().terminate.assert_called_with(signum)
job._apply_result = Mock(name='_apply_result2')
job._apply_result.return_value = None
job.terminate(pool, signal='TERM')
def test_terminate__task_started(self):
pool = Mock()
signum = signal.SIGTERM
job = Request({
'task': self.mytask.name,
'id': uuid(),
'args': [1],
'kwrgs': {'f': 'x'},
}, app=self.app)
with assert_signal_called(
task_revoked, sender=job.task, request=job,
terminated=True, expired=False, signum=signum):
job.time_start = monotonic()
job.worker_pid = 313
job.terminate(pool, signal='TERM')
pool.terminate_job.assert_called_with(job.worker_pid, signum)
def scale_down(self, n):
if n and self._last_action and (monotonic() - self._last_action >
self.keepalive):
self._last_action = monotonic()
return self._shrink(n)
def test_on_accepted_acks_late(self):
job = self.xRequest()
self.mytask.acks_late = True
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
self.assertFalse(job.acknowledged)
def should_sync(self):
return ((not self._last_sync or
(monotonic() - self._last_sync) > self.sync_every)
or (self.sync_every_tasks
and self._tasks_since_sync >= self.sync_every_tasks))
def scale_up(self, n):
self._last_scale_up = monotonic()
return self._grow(n)
def trace_task(uuid, args, kwargs, request=None):
# R - is the possibly prepared return value.
# I - is the Info object.
# T - runtime
# Rstr - textual representation of return value
# retval - is the always unmodified return value.
# state - is the resulting task state.
# This function is very long because we've unrolled all the calls
# for performance reasons, and because the function is so long
# we want the main variables (I, and R) to stand out visually from the
# the rest of the variables, so breaking PEP8 is worth it ;)
R = I = T = Rstr = retval = state = None
task_request = None
time_start = monotonic()
try:
try:
kwargs.items
except AttributeError:
raise InvalidTaskError(
'Task keyword arguments is not a mapping')
push_task(task)
task_request = Context(request or {}, args=args,
called_directly=False, kwargs=kwargs)
root_id = task_request.root_id or uuid
task_priority = task_request.delivery_info.get('priority') if \
inherit_parent_priority else None
push_request(task_request)
try:
# -*- PRE -*-
if prerun_receivers:
send_prerun(sender=task, task_id=uuid, task=task,
args=args, kwargs=kwargs)
loader_task_init(uuid, task)
if track_started:
store_result(
uuid, {'pid': pid, 'hostname': hostname}, STARTED,
request=task_request,
)
# -*- TRACE -*-
try:
R = retval = fun(*args, **kwargs)
state = SUCCESS
except Reject as exc:
I, R = Info(REJECTED, exc), ExceptionInfo(internal=True)
state, retval = I.state, I.retval
I.handle_reject(task, task_request)
except Ignore as exc:
I, R = Info(IGNORED, exc), ExceptionInfo(internal=True)
state, retval = I.state, I.retval
I.handle_ignore(task, task_request)
except Retry as exc:
I, R, state, retval = on_error(
task_request, exc, uuid, RETRY, call_errbacks=False)
except Exception as exc:
I, R, state, retval = on_error(task_request, exc, uuid)
except BaseException:
raise
else:
try:
# callback tasks must be applied before the result is
# stored, so that result.children is populated.
# groups are called inline and will store trail
# separately, so need to call them separately
# so that the trail's not added multiple times :(
# (Issue #1936)
callbacks = task.request.callbacks
if callbacks:
if len(task.request.callbacks) > 1:
sigs, groups = [], []
for sig in callbacks:
sig = signature(sig, app=app)
if isinstance(sig, group):
groups.append(sig)
else:
sigs.append(sig)
for group_ in groups:
group_.apply_async(
(retval,),
parent_id=uuid, root_id=root_id,
priority=task_priority
)
if sigs:
group(sigs, app=app).apply_async(
(retval,),
parent_id=uuid, root_id=root_id,
priority=task_priority
)
else:
signature(callbacks[0], app=app).apply_async(
(retval,), parent_id=uuid, root_id=root_id,
priority=task_priority
)
# execute first task in chain
chain = task_request.chain
if chain:
_chsig = signature(chain.pop(), app=app)
_chsig.apply_async(
(retval,), chain=chain,
parent_id=uuid, root_id=root_id,
priority=task_priority
)
mark_as_done(
uuid, retval, task_request, publish_result,
)
except EncodeError as exc:
I, R, state, retval = on_error(task_request, exc, uuid)
else:
Rstr = saferepr(R, resultrepr_maxsize)
T = monotonic() - time_start
if task_on_success:
task_on_success(retval, uuid, args, kwargs)
if success_receivers:
send_success(sender=task, result=retval)
if _does_info:
info(LOG_SUCCESS, {
'id': uuid,
'name': get_task_name(task_request, name),
'return_value': Rstr,
'runtime': T,
})
# -* POST *-
if state not in IGNORE_STATES:
if task_after_return:
task_after_return(
state, retval, uuid, args, kwargs, None,
)
finally:
try:
if postrun_receivers:
send_postrun(sender=task, task_id=uuid, task=task,
args=args, kwargs=kwargs,
retval=retval, state=state)
finally:
pop_task()
pop_request()
if not eager:
try:
backend_cleanup()
loader_cleanup()
except (KeyboardInterrupt, SystemExit, MemoryError):
raise
except Exception as exc:
logger.error('Process cleanup failed: %r', exc,
exc_info=True)
except MemoryError:
raise
except Exception as exc:
if eager:
raise
R = report_internal_error(task, exc)
if task_request is not None:
I, _, _, _ = on_error(task_request, exc, uuid)
return trace_ok_t(R, I, T, Rstr)
def scale_down(self, n):
if self._last_scale_up and (monotonic() - self._last_scale_up >
self.keepalive):
return self._shrink(n)
def runtest(self, fun, n=50, index=0, repeats=1):
n = getattr(fun, '__iterations__', None) or n
print('{0}: [[[{1}({2})]]]'.format(repeats, fun.__name__, n))
with blockdetection(self.block_timeout):
with self.fbi.investigation():
runtime = elapsed = monotonic()
i = 0
failed = False
self.progress = Progress(
fun,
i,
n,
index,
repeats,
elapsed,
runtime,
0,
)
_marker.delay(pstatus(self.progress))
try:
for i in range(n):
runtime = monotonic()
self.progress = Progress(
fun,
i + 1,
n,
index,
repeats,
runtime,
elapsed,
0,
)
try:
fun()
except StopSuite:
raise
except Exception as exc:
print('-> {0!r}'.format(exc))
import traceback
print(traceback.format_exc())
print(pstatus(self.progress))
else:
print(pstatus(self.progress))
except Exception:
failed = True
self.speaker.beep()
raise
finally:
print('{0} {1} iterations in {2}s'.format(
'failed after' if failed else 'completed',
i + 1,
humanize_seconds(monotonic() - elapsed),
))
if not failed:
self.progress = Progress(
fun,
i + 1,
n,
index,
repeats,
runtime,
elapsed,
1,
)
def __init__(self, gap=5.0):
self.gap = gap
self.last_noise = monotonic() - self.gap * 2
def pstatus(p):
return F_PROGRESS.format(
p,
runtime=humanize_seconds(monotonic() - p.runtime, now='0 seconds'),
elapsed=humanize_seconds(monotonic() - p.elapsed, now='0 seconds'),
)
def runtest(self, fun, n=50, index=0, repeats=1):
n = getattr(fun, '__iterations__', None) or n
header = '[[[{0}({1})]]]'.format(fun.__name__, n)
if repeats > 1:
header = '{0} #{1}'.format(header, repeats)
self.print(header)
with blockdetection(self.block_timeout):
with self.fbi.investigation():
runtime = elapsed = monotonic()
i = 0
failed = False
self.progress = Progress(
fun,
i,
n,
index,
repeats,
elapsed,
runtime,
0,
)
_marker.delay(pstatus(self.progress))
try:
for i in range(n):
runtime = monotonic()
self.progress = Progress(
fun,
i + 1,
n,
index,
repeats,
runtime,
elapsed,
0,
)
self.execute_test(fun)
except Exception:
failed = True
self.speaker.beep()
raise
finally:
if n > 1 or failed:
self.print('{0} {1} iterations in {2}'.format(
'failed after' if failed else 'completed',
i + 1,
humanize_seconds(monotonic() - elapsed),
),
file=self.stderr if failed else self.stdout)
if not failed:
self.progress = Progress(
fun,
i + 1,
n,
index,
repeats,
runtime,
elapsed,
1,
)
def test_on_accepted_time_start(self):
job = self.xRequest()
job.on_accepted(pid=os.getpid(), time_accepted=monotonic())
assert time() - job.time_start < 1
def test_should_sync(self):
self.assertTrue(self.s.should_sync())
self.s._last_sync = monotonic()
self.assertFalse(self.s.should_sync())
self.s._last_sync -= self.s.sync_every
self.assertTrue(self.s.should_sync())
def test_should_sync(self):
assert self.s.should_sync()
self.s._last_sync = monotonic()
assert not self.s.should_sync()
self.s._last_sync -= self.s.sync_every
assert self.s.should_sync()
def __init__(self, gap=5.0, file=None):
self.gap = gap
self.file = sys.stdout if file is None else file
self.last_noise = monotonic() - self.gap * 2
def beep(self):
now = monotonic()
if now - self.last_noise >= self.gap:
self.emit()
self.last_noise = now
