Spaces:
Running
Running
File size: 2,023 Bytes
89e57b7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """Account for every accepted Celery delivery, including retries and chains."""
import uuid
from celery import Task, current_task
from . import maintenance as gate
from .state import enabled
from .git import BackupError
HEADER = 'brainapi-backup-work'
class ContentTask(Task):
abstract = True
def apply_async(self, args=None, kwargs=None, **options):
if not enabled():
return super().apply_async(args=args, kwargs=kwargs, **options)
parent = getattr(getattr(current_task, 'request', None), 'headers', None) or {}
work_id = uuid.uuid4().hex
parent_id = getattr(getattr(current_task, 'request', None), 'backup_execution', None) or parent.get(HEADER)
gate.register_task(work_id, parent_id)
headers = dict(options.pop('headers', None) or {})
headers[HEADER] = work_id
# An ambiguous publication failure deliberately stays pending: fail closed
# rather than claim a durable barrier while a message might be in flight.
return super().apply_async(args=args, kwargs=kwargs, headers=headers, **options)
def __call__(self, *args, **kwargs):
if not enabled():
return super().__call__(*args, **kwargs)
work_id = (self.request.headers or {}).get(HEADER)
self.request.backup_execution = gate.begin_task(work_id)
token = gate._admitted.set(True)
try:
return super().__call__(*args, **kwargs)
finally:
gate._admitted.reset(token)
def after_return(self, status, retval, task_id, args, kwargs, einfo):
gate.finish_task((self.request.headers or {}).get(HEADER), getattr(self.request, 'backup_execution', None))
return super().after_return(status, retval, task_id, args, kwargs, einfo)
def on_retry(self, exc, task_id, args, kwargs, einfo):
gate.finish_task((self.request.headers or {}).get(HEADER), getattr(self.request, 'backup_execution', None))
return super().on_retry(exc, task_id, args, kwargs, einfo)
|