# -*- coding: utf-8 -*- # import os import datetime import sys import time from django.conf import settings from django.utils import timezone from django.core.cache import cache from django.db import transaction from celery import subtask from celery.signals import worker_ready, worker_shutdown, task_prerun, \ task_postrun, after_task_publish from django_celery_beat.models import PeriodicTask from common.utils import get_logger, TeeObj, get_object_or_none from common.const import celery_task_pre_key from .utils import get_after_app_ready_tasks, get_after_app_shutdown_clean_tasks from ..models import CeleryTask logger = get_logger(__file__) @worker_ready.connect def on_app_ready(sender=None, headers=None, body=None, **kwargs): if cache.get("CELERY_APP_READY", 0) == 1: return cache.set("CELERY_APP_READY", 1, 10) tasks = get_after_app_ready_tasks() logger.debug("Start need start task: [{}]".format( ", ".join(tasks)) ) for task in tasks: subtask(task).delay() @worker_shutdown.connect def after_app_shutdown(sender=None, headers=None, body=None, **kwargs): if cache.get("CELERY_APP_SHUTDOWN", 0) == 1: return cache.set("CELERY_APP_SHUTDOWN", 1, 10) tasks = get_after_app_shutdown_clean_tasks() logger.debug("App shutdown signal recv") logger.debug("Clean need cleaned period tasks: [{}]".format( ', '.join(tasks)) ) PeriodicTask.objects.filter(name__in=tasks).delete() @after_task_publish.connect def after_task_publish_signal_handler(sender, headers=None, **kwargs): CeleryTask.objects.create( id=headers["id"], status=CeleryTask.WAITING, name=headers["task"] ) cache.set(headers["id"], True, 3600) @task_prerun.connect def pre_run_task_signal_handler(sender, task_id=None, task=None, **kwargs): time.sleep(0.1) for i in range(5): if cache.get(task_id, False): break else: time.sleep(0.1) continue t = get_object_or_none(CeleryTask, id=task_id) if t is None: logger.warn("Not get the task: {}".format(task_id)) return now = datetime.datetime.now().strftime("%Y-%m-%d") log_path = os.path.join(now, task_id + '.log') full_path = os.path.join(CeleryTask.LOG_DIR, log_path) if not os.path.exists(os.path.dirname(full_path)): os.makedirs(os.path.dirname(full_path)) with transaction.atomic(): t.date_start = timezone.now() t.status = CeleryTask.RUNNING t.log_path = log_path t.save() f = open(full_path, 'w') tee = TeeObj(f) sys.stdout = tee task.log_f = tee @task_postrun.connect def post_run_task_signal_handler(sender, task_id=None, task=None, **kwargs): t = get_object_or_none(CeleryTask, id=task_id) if t is None: logger.warn("Not get the task: {}".format(task_id)) return with transaction.atomic(): t.status = CeleryTask.FINISHED t.date_finished = timezone.now() t.save() task.log_f.flush() sys.stdout = task.log_f.origin_stdout task.log_f.close()