2018-02-05 14:58:02 +00:00
|
|
|
import time
|
|
|
|
|
|
2020-04-24 14:36:21 +01:00
|
|
|
from gds_metrics.metrics import Histogram
|
2017-08-31 12:52:06 +01:00
|
|
|
from celery import Celery, Task
|
2018-07-12 15:09:38 +01:00
|
|
|
from celery.signals import worker_process_shutdown
|
2020-04-24 11:21:41 +01:00
|
|
|
from flask import g, request
|
2019-10-23 16:14:59 +01:00
|
|
|
from flask.ctx import has_request_context
|
2018-07-12 15:09:38 +01:00
|
|
|
|
|
|
|
|
|
|
|
|
|
@worker_process_shutdown.connect
|
2020-04-24 11:21:41 +01:00
|
|
|
def log_on_worker_shutdown(sender, signal, pid, exitcode, **kwargs):
|
|
|
|
|
# imported here to avoid circular imports
|
|
|
|
|
from app import notify_celery
|
|
|
|
|
|
|
|
|
|
# if the worker has already restarted at least once, then we no longer have app context and current_app won't work
|
|
|
|
|
# to create a new one. Instead we have to create a new app context from the original flask app and use that instead.
|
|
|
|
|
with notify_celery._app.app_context():
|
|
|
|
|
# if the worker has restarted
|
|
|
|
|
notify_celery._app.logger.info('worker shutdown: PID: {} Exitcode: {}'.format(pid, exitcode))
|
2017-08-31 12:52:06 +01:00
|
|
|
|
|
|
|
|
|
2018-02-12 15:29:03 +00:00
|
|
|
def make_task(app):
|
2020-04-24 15:15:41 +01:00
|
|
|
SQS_APPLY_ASYNC_DURATION_SECONDS = Histogram(
|
|
|
|
|
'sqs_apply_async_duration_seconds',
|
2020-04-24 14:36:21 +01:00
|
|
|
'Time taken to put task on queue',
|
|
|
|
|
['task_name']
|
|
|
|
|
)
|
|
|
|
|
|
2018-02-12 15:29:03 +00:00
|
|
|
class NotifyTask(Task):
|
|
|
|
|
abstract = True
|
|
|
|
|
start = None
|
2018-02-05 14:58:02 +00:00
|
|
|
|
2018-02-12 15:29:03 +00:00
|
|
|
def on_success(self, retval, task_id, args, kwargs):
|
|
|
|
|
elapsed_time = time.time() - self.start
|
|
|
|
|
app.logger.info(
|
|
|
|
|
"{task_name} took {time}".format(
|
|
|
|
|
task_name=self.name, time="{0:.4f}".format(elapsed_time)
|
|
|
|
|
)
|
2018-02-05 14:58:02 +00:00
|
|
|
)
|
2017-08-31 12:52:06 +01:00
|
|
|
|
2018-02-12 15:29:03 +00:00
|
|
|
def on_failure(self, exc, task_id, args, kwargs, einfo):
|
|
|
|
|
# ensure task will log exceptions to correct handlers
|
2018-02-22 15:05:37 +00:00
|
|
|
app.logger.exception('Celery task: {} failed'.format(self.name))
|
2018-02-12 15:29:03 +00:00
|
|
|
super().on_failure(exc, task_id, args, kwargs, einfo)
|
|
|
|
|
|
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
|
|
|
# ensure task has flask context to access config, logger, etc
|
|
|
|
|
with app.app_context():
|
|
|
|
|
self.start = time.time()
|
2019-10-23 16:14:59 +01:00
|
|
|
# Remove 'request_id' from the kwargs (so the task doesn't get an unexpected kwarg), then add it to g
|
|
|
|
|
# so that it gets logged
|
|
|
|
|
g.request_id = kwargs.pop('request_id', None)
|
2018-02-12 15:29:03 +00:00
|
|
|
return super().__call__(*args, **kwargs)
|
2017-08-31 12:52:06 +01:00
|
|
|
|
2019-10-23 16:14:59 +01:00
|
|
|
def apply_async(self, args=None, kwargs=None, task_id=None, producer=None,
|
|
|
|
|
link=None, link_error=None, **options):
|
|
|
|
|
kwargs = kwargs or {}
|
|
|
|
|
|
|
|
|
|
if has_request_context() and hasattr(request, 'request_id'):
|
|
|
|
|
kwargs['request_id'] = request.request_id
|
|
|
|
|
|
2020-04-24 15:15:41 +01:00
|
|
|
with SQS_APPLY_ASYNC_DURATION_SECONDS.labels(self.name).time():
|
2020-04-24 14:36:21 +01:00
|
|
|
return super().apply_async(args, kwargs, task_id, producer, link, link_error, **options)
|
2019-10-23 16:14:59 +01:00
|
|
|
|
2018-02-12 15:29:03 +00:00
|
|
|
return NotifyTask
|
2017-06-01 14:32:19 +01:00
|
|
|
|
2016-02-09 13:31:45 +00:00
|
|
|
|
|
|
|
|
class NotifyCelery(Celery):
|
2017-06-09 16:20:02 +01:00
|
|
|
|
2017-07-19 13:50:29 +01:00
|
|
|
def init_app(self, app):
|
2017-08-31 12:52:06 +01:00
|
|
|
super().__init__(
|
|
|
|
|
app.import_name,
|
|
|
|
|
broker=app.config['BROKER_URL'],
|
2018-02-12 15:29:03 +00:00
|
|
|
task_cls=make_task(app),
|
2017-08-31 12:52:06 +01:00
|
|
|
)
|
2016-02-09 13:31:45 +00:00
|
|
|
|
2017-08-31 12:52:06 +01:00
|
|
|
self.conf.update(app.config)
|
2020-04-24 11:21:41 +01:00
|
|
|
self._app = app
|