├── cq ├── tests │ ├── __init__.py │ ├── test_logs.py │ ├── test_force_chain.py │ └── test_base.py ├── management │ ├── __init__.py │ └── commands │ │ ├── __init__.py │ │ ├── cq_clear_logs.py │ │ ├── cq_retry.py │ │ ├── cq_check_lost.py │ │ ├── cq_runworker.py │ │ ├── cq_maintenance.py │ │ └── cq_load_test.py ├── migrations │ ├── __init__.py │ ├── 0004_task_retries.py │ ├── 0006_task_force_chain.py │ ├── 0005_task_last_retry.py │ ├── 0002_add_repeating.py │ ├── 0003_auto_20160911_0749.py │ ├── 0007_auto_20161010_1635.py │ └── 0001_initial.py ├── routing.py ├── router.py ├── __init__.py ├── managers.py ├── signals.py ├── decorators.py ├── serializers.py ├── views.py ├── admin.py ├── signature.py ├── apps.py ├── utils.py ├── consumers.py ├── scheduler.py ├── tasks.py ├── task.py └── models.py ├── MANIFEST.in ├── .gitignore ├── setup.cfg ├── setup.py ├── COPYRIGHT └── README.md /cq/tests/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /cq/management/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /cq/migrations/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include README.md 2 | -------------------------------------------------------------------------------- /cq/management/commands/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /cq/routing.py: -------------------------------------------------------------------------------- 1 | from .consumers import CQConsumer 2 | 3 | channel_routing = { 4 | 'cq-task': CQConsumer 5 | } 6 | -------------------------------------------------------------------------------- /cq/router.py: -------------------------------------------------------------------------------- 1 | from .views import TaskViewSet 2 | 3 | 4 | def register(router, name='cq'): 5 | router.register(r'cq/tasks', TaskViewSet, base_name='cqtask') 6 | -------------------------------------------------------------------------------- /cq/__init__.py: -------------------------------------------------------------------------------- 1 | default_app_config = 'cq.apps.CqConfig' 2 | 3 | VERSION = (0, 3, 3) 4 | 5 | __version__ = '.'.join(str(x) for x in VERSION[:(2 if VERSION[2] == 0 else 3)]) # noqa 6 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | **~ 2 | **.pyc 3 | **._* 4 | **.DS_Store 5 | **.nfs* 6 | build* 7 | .sconf_temp 8 | .sconsign.dblite 9 | dist 10 | *.egg-info 11 | db.sqlite3 12 | environ/ 13 | examples/example_site/env 14 | -------------------------------------------------------------------------------- /cq/managers.py: -------------------------------------------------------------------------------- 1 | from django.db import models 2 | 3 | 4 | class TaskManager(models.Manager): 5 | def active(self, **kwargs): 6 | return self.filter(status__in=self.model.STATUS_ACTIVE, **kwargs) 7 | -------------------------------------------------------------------------------- /cq/management/commands/cq_clear_logs.py: -------------------------------------------------------------------------------- 1 | from django.core.management.base import BaseCommand 2 | from cq.tasks import clear_logs 3 | 4 | 5 | class Command(BaseCommand): 6 | help = 'Clear all logs from REDIS.' 7 | 8 | def handle(self, *args, **options): 9 | clear_logs() 10 | -------------------------------------------------------------------------------- /cq/signals.py: -------------------------------------------------------------------------------- 1 | from django.db.models.signals import pre_save 2 | from django.dispatch import receiver 3 | 4 | from .models import RepeatingTask 5 | 6 | 7 | @receiver(pre_save, sender=RepeatingTask) 8 | def update_next_run(sender, instance, **kwargs): 9 | instance.update_next_run() 10 | -------------------------------------------------------------------------------- /setup.cfg: -------------------------------------------------------------------------------- 1 | [metadata] 2 | description-file = README.md 3 | 4 | [flake8] 5 | ignore = E501 6 | max-line-length = 100 7 | 8 | [isort] 9 | sections = FUTURE,STDLIB,THIRDPARTY,FIRSTPARTY,LOCALFOLDER 10 | default_section = THIRDPARTY 11 | known_standard_library = six 12 | known_first_party = cq 13 | multi_line_output = 3 14 | line_length = 100 15 | indent = 4 16 | -------------------------------------------------------------------------------- /cq/management/commands/cq_retry.py: -------------------------------------------------------------------------------- 1 | from django.core.management.base import BaseCommand, CommandError 2 | from cq.tasks import retry_tasks 3 | 4 | 5 | class Command(BaseCommand): 6 | help = 'Retry tasks.' 7 | 8 | # def add_arguments(self, parser): 9 | # parser.add_argument(') 10 | 11 | def handle(self, *args, **options): 12 | retry_tasks() 13 | -------------------------------------------------------------------------------- /cq/management/commands/cq_check_lost.py: -------------------------------------------------------------------------------- 1 | from django.core.management.base import BaseCommand, CommandError 2 | from cq.tasks import check_lost 3 | 4 | 5 | class Command(BaseCommand): 6 | help = 'Check for lost tasks.' 7 | 8 | # def add_arguments(self, parser): 9 | # parser.add_argument(') 10 | 11 | def handle(self, *args, **options): 12 | check_lost() 13 | -------------------------------------------------------------------------------- /cq/decorators.py: -------------------------------------------------------------------------------- 1 | from .task import TaskFunc 2 | 3 | 4 | def task(*args, **kwargs): 5 | if len(args): 6 | func_or_name = args[0] 7 | args = args[1:] 8 | else: 9 | func_or_name = None 10 | if callable(func_or_name): 11 | func = func_or_name 12 | name = None 13 | else: 14 | func = None 15 | name = func_or_name 16 | dec = TaskFunc(name, *args, **kwargs) 17 | if func: 18 | return dec(func) 19 | else: 20 | return dec 21 | -------------------------------------------------------------------------------- /cq/migrations/0004_task_retries.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.9.4 on 2016-09-19 04:09 3 | from __future__ import unicode_literals 4 | 5 | from django.db import migrations, models 6 | 7 | 8 | class Migration(migrations.Migration): 9 | 10 | dependencies = [ 11 | ('cq', '0003_auto_20160911_0749'), 12 | ] 13 | 14 | operations = [ 15 | migrations.AddField( 16 | model_name='task', 17 | name='retries', 18 | field=models.PositiveIntegerField(default=0), 19 | ), 20 | ] 21 | -------------------------------------------------------------------------------- /cq/migrations/0006_task_force_chain.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.9.4 on 2016-10-10 00:03 3 | from __future__ import unicode_literals 4 | 5 | from django.db import migrations, models 6 | 7 | 8 | class Migration(migrations.Migration): 9 | 10 | dependencies = [ 11 | ('cq', '0005_task_last_retry'), 12 | ] 13 | 14 | operations = [ 15 | migrations.AddField( 16 | model_name='task', 17 | name='force_chain', 18 | field=models.BooleanField(default=False), 19 | ), 20 | ] 21 | -------------------------------------------------------------------------------- /cq/migrations/0005_task_last_retry.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.9.4 on 2016-09-20 01:52 3 | from __future__ import unicode_literals 4 | 5 | from django.db import migrations, models 6 | 7 | 8 | class Migration(migrations.Migration): 9 | 10 | dependencies = [ 11 | ('cq', '0004_task_retries'), 12 | ] 13 | 14 | operations = [ 15 | migrations.AddField( 16 | model_name='task', 17 | name='last_retry', 18 | field=models.DateTimeField(blank=True, null=True), 19 | ), 20 | ] 21 | -------------------------------------------------------------------------------- /cq/migrations/0002_add_repeating.py: -------------------------------------------------------------------------------- 1 | from django.db import migrations 2 | from cq.tasks import maintenance 3 | from cq.models import schedule_task 4 | 5 | 6 | def add_repeating(apps, scema_editor): 7 | RepeatingTask = apps.get_model('cq.RepeatingTask') 8 | schedule_task( 9 | RepeatingTask, 10 | '* * * * *', 11 | maintenance, 12 | result_ttl=30, 13 | coalesce=False 14 | ) 15 | 16 | 17 | class Migration(migrations.Migration): 18 | dependencies = [ 19 | ('cq', '0001_initial') 20 | ] 21 | operations = [ 22 | migrations.RunPython(add_repeating, reverse_code=migrations.RunPython.noop) 23 | ] 24 | -------------------------------------------------------------------------------- /cq/serializers.py: -------------------------------------------------------------------------------- 1 | from rest_framework import serializers 2 | 3 | from .models import Task, delay 4 | 5 | 6 | class TaskSerializer(serializers.ModelSerializer): 7 | class Meta: 8 | model = Task 9 | fields = '__all__' 10 | 11 | 12 | class CreateTaskSerializer(serializers.Serializer): 13 | task = serializers.CharField(max_length=128) 14 | args = serializers.JSONField(required=False) 15 | kwargs = serializers.JSONField(required=False) 16 | 17 | def create(self, data, *args, **kwargs): 18 | return delay(data['task'], data.get('args', ()), data.get('kwargs', {})) 19 | 20 | def to_representation(self, inst): 21 | return { 22 | 'id': inst.id 23 | } 24 | -------------------------------------------------------------------------------- /cq/views.py: -------------------------------------------------------------------------------- 1 | from django.db import transaction 2 | from rest_framework import viewsets 3 | 4 | from .serializers import TaskSerializer, CreateTaskSerializer 5 | from .models import Task 6 | 7 | 8 | class TaskViewSet(viewsets.ModelViewSet): 9 | queryset = Task.objects.all() 10 | serializer_class = TaskSerializer 11 | 12 | def get_serializer(self, data=None, *args, **kwargs): 13 | if getattr(self, 'creating', False): 14 | return CreateTaskSerializer(data=data) 15 | return super().get_serializer(data, *args, **kwargs) 16 | 17 | def create(self, request, *args, **kwargs): 18 | self.creating = True 19 | with transaction.atomic(): 20 | return super().create(request, *args, **kwargs) 21 | -------------------------------------------------------------------------------- /cq/migrations/0003_auto_20160911_0749.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.10 on 2016-09-10 21:49 3 | from __future__ import unicode_literals 4 | 5 | from django.db import migrations, models 6 | 7 | 8 | class Migration(migrations.Migration): 9 | 10 | dependencies = [ 11 | ('cq', '0002_add_repeating'), 12 | ] 13 | 14 | operations = [ 15 | migrations.AlterField( 16 | model_name='task', 17 | name='status', 18 | field=models.CharField(choices=[('P', 'Pending'), ('Y', 'Retry'), ('Q', 'Queued'), ('R', 'Running'), ('F', 'Failure'), ('S', 'Success'), ('W', 'Waiting'), ('I', 'Incomplete'), ('L', 'Lost'), ('E', 'Revoked')], db_index=True, default='P', max_length=1), 19 | ), 20 | ] 21 | -------------------------------------------------------------------------------- /cq/admin.py: -------------------------------------------------------------------------------- 1 | from django import forms 2 | from django.contrib import admin 3 | 4 | from .models import RepeatingTask, Task 5 | from .task import FuncNameWidget 6 | 7 | 8 | @admin.register(Task) 9 | class TaskAdmin(admin.ModelAdmin): 10 | list_display = ('id', 'submitted', 'func_name', 'status', ) 11 | list_filter = ('status', ) 12 | 13 | 14 | class RepeatingTaskAdminForm(forms.ModelForm): 15 | class Meta: 16 | model = RepeatingTask 17 | fields = '__all__' 18 | widgets = { 19 | 'func_name': FuncNameWidget 20 | } 21 | 22 | 23 | @admin.register(RepeatingTask) 24 | class RepeatingTaskAdmin(admin.ModelAdmin): 25 | form = RepeatingTaskAdminForm 26 | list_display = ('func_name', 'args', 'kwargs', 'coalesce', 'crontab', 'last_run', 'next_run', ) 27 | -------------------------------------------------------------------------------- /cq/migrations/0007_auto_20161010_1635.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.9.4 on 2016-10-10 05:35 3 | from __future__ import unicode_literals 4 | 5 | from django.db import migrations, models 6 | import django.db.models.deletion 7 | 8 | 9 | class Migration(migrations.Migration): 10 | 11 | dependencies = [ 12 | ('cq', '0006_task_force_chain'), 13 | ] 14 | 15 | operations = [ 16 | migrations.AlterField( 17 | model_name='task', 18 | name='previous', 19 | field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, related_name='next', to='cq.Task'), 20 | ), 21 | migrations.AlterField( 22 | model_name='task', 23 | name='waiting_on', 24 | field=models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.SET_NULL, to='cq.Task'), 25 | ), 26 | ] 27 | -------------------------------------------------------------------------------- /cq/signature.py: -------------------------------------------------------------------------------- 1 | import six 2 | import inspect 3 | 4 | from .task import TaskFunc 5 | 6 | 7 | def to_func_name(func): 8 | 9 | # Convert to a module import string. 10 | if inspect.isfunction(func) or inspect.isbuiltin(func): 11 | name = '{0}.{1}'.format(func.__module__, func.__name__) 12 | elif isinstance(func, six.string_types): 13 | name = str(func) 14 | else: 15 | msg = 'Expected a callable or a string, but got: {}'.format(func) 16 | raise TypeError(msg) 17 | 18 | # Try to convert to a name before returing. Will default 19 | # to the import string. 20 | return TaskFunc.get_name(name) 21 | 22 | 23 | def to_class_name(cls): 24 | return '{0}.{1}'.format(cls.__module__, cls.__name__) 25 | 26 | 27 | def from_class_name(name): 28 | return from_func_name(name) 29 | 30 | 31 | def to_signature(func, args, kwargs): 32 | return { 33 | 'func_name': to_func_name(func), 34 | 'args': args, 35 | 'kwargs': kwargs 36 | } 37 | 38 | 39 | def from_signature(sig): 40 | func = TaskFunc.get_task(sig['func_name']).func 41 | return (func, tuple(sig.get('args', ())), sig.get('kwargs', {})) 42 | -------------------------------------------------------------------------------- /cq/apps.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from django.apps import AppConfig 4 | from django.conf import settings 5 | from django.core.cache import cache 6 | from django.db.utils import ProgrammingError 7 | from django.utils.module_loading import import_module 8 | 9 | logger = logging.getLogger('cq') 10 | 11 | 12 | def scan_tasks(*args, **kwargs): 13 | for app_name in settings.INSTALLED_APPS: 14 | try: 15 | import_module('.'.join([app_name, 'tasks'])) 16 | except ImportError: 17 | pass 18 | 19 | 20 | def requeue_tasks(*args, **kwargs): 21 | from cq.models import Task 22 | lock = 'RETRY_QUEUED_TASKS' 23 | with cache.lock(lock, timeout=10): 24 | # Find all Queued tasks and set them to Retry, since they get stuck after a reboot 25 | try: 26 | tasks = Task.objects.filter(status=Task.STATUS_QUEUED).update(status=Task.STATUS_RETRY) 27 | logger.info('Requeued {} tasks'.format(tasks)) 28 | except ProgrammingError: 29 | logger.warning("Failed requeuing tasks; database likely hasn't had migrations run.") 30 | pass 31 | 32 | 33 | class CqConfig(AppConfig): 34 | name = 'cq' 35 | 36 | def ready(self): 37 | import cq.signals 38 | scan_tasks() 39 | requeue_tasks() 40 | -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import os 2 | import re 3 | 4 | from setuptools import find_packages, setup 5 | 6 | with open('./cq/__init__.py') as f: 7 | exec(re.search(r'VERSION = .*', f.read(), re.DOTALL).group()) 8 | 9 | setup( 10 | name='django-cq', 11 | version=__version__, # noqa 12 | author='Luke Hodkinson', 13 | author_email='luke.hodkinson@uptickhq.com', 14 | maintainer='Uptick', 15 | maintainer_email='dev@uptickhq.com', 16 | url='https://github.com/uptick/django-cq', 17 | description='Distributed tasks for Django Channels.', 18 | long_description=open(os.path.join(os.path.dirname(__file__), 'README.md')).read(), 19 | long_description_content_type='text/markdown', 20 | classifiers=[ 21 | 'Development Status :: 3 - Alpha', 22 | 'Framework :: Django', 23 | 'Intended Audience :: Developers', 24 | 'License :: OSI Approved :: BSD License', 25 | 'Natural Language :: English', 26 | 'Operating System :: OS Independent', 27 | 'Programming Language :: Python', 28 | ], 29 | license='BSD', 30 | packages=find_packages(), 31 | include_package_data=True, 32 | package_data={'': ['*.txt', '*.js', '*.html', '*.*']}, 33 | install_requires=[ 34 | 'setuptools', 35 | 'six', 36 | 'croniter', 37 | 'channels>=2.1.0', 38 | 'channels_redis>=2.3.0', 39 | ], 40 | ) 41 | -------------------------------------------------------------------------------- /cq/management/commands/cq_runworker.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import time 3 | from threading import Thread 4 | 5 | from channels.management.commands.runworker import Command as BaseCommand 6 | from django_redis import get_redis_connection 7 | 8 | from django.conf import settings 9 | 10 | from ...consumers import run_task 11 | from ...utils import get_redis_key 12 | 13 | logger = logging.getLogger('cq') 14 | 15 | 16 | def launch_scheduler(*args, **kwargs): 17 | from ...scheduler import scheduler 18 | logger.info('Launching CQ scheduler.') 19 | thread = Thread(name='scheduler', target=scheduler) 20 | thread.daemon = True 21 | thread.start() 22 | 23 | 24 | class Command(BaseCommand): 25 | def handle(self, *args, **options): 26 | if getattr(settings, 'CQ_SCHEDULER', True): 27 | launch_scheduler() 28 | if getattr(settings, 'CQ_BACKEND', '').lower() == 'redis': 29 | self.handle_redis_backend() 30 | else: 31 | super().handle(*args, **options) 32 | 33 | def handle_redis_backend(self): 34 | while True: 35 | conn = get_redis_connection() 36 | while True: 37 | message = conn.brpop(get_redis_key('cq')) 38 | try: 39 | run_task(message[1].decode()) 40 | except Exception as e: 41 | logger.error(str(e)) 42 | time.sleep(1) 43 | -------------------------------------------------------------------------------- /COPYRIGHT: -------------------------------------------------------------------------------- 1 | Copyright (c) 2011, Luke Hodkinson 2 | All rights reserved. 3 | 4 | Redistribution and use in source and binary forms, with or without 5 | modification, are permitted provided that the following conditions are 6 | met: 7 | 8 | * Redistributions of source code must retain the above copyright 9 | notice, this list of conditions and the following disclaimer. 10 | * Redistributions in binary form must reproduce the above 11 | copyright notice, this list of conditions and the following 12 | disclaimer in the documentation and/or other materials provided 13 | with the distribution. 14 | * Neither the name of the author nor the names of other 15 | contributors may be used to endorse or promote products derived 16 | from this software without specific prior written permission. 17 | 18 | THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS 19 | "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT 20 | LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR 21 | A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT 22 | OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, 23 | SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT 24 | LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, 25 | DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY 26 | THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT 27 | (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE 28 | OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. 29 | -------------------------------------------------------------------------------- /cq/management/commands/cq_maintenance.py: -------------------------------------------------------------------------------- 1 | import datetime 2 | 3 | from django.core.management.base import BaseCommand 4 | from django.utils import timezone 5 | 6 | from ...models import Task 7 | from ...tasks import maintenance 8 | 9 | 10 | class Command(BaseCommand): 11 | help = 'Perform maintenance on CQ tasks.' 12 | 13 | def add_arguments(self, parser): 14 | parser.add_argument('--clear', action='store_true') 15 | parser.add_argument('--failed', action='store_true') 16 | parser.add_argument('--pending', action='store_true') 17 | parser.add_argument('--lost', action='store_true') 18 | parser.add_argument( 19 | '--prior', 20 | help='number of days ago prior to which the search occurs' 21 | ) 22 | parser.add_argument('--dry-run', action='store_true') 23 | 24 | def handle(self, *args, **options): 25 | if options['clear']: 26 | self.clear(options) 27 | else: 28 | maintenance.delay() 29 | 30 | def clear(self, options): 31 | query = {} 32 | 33 | statuses = [] 34 | if options['failed']: 35 | statuses.append(Task.STATUS_FAILURE) 36 | if options['pending']: 37 | statuses.append(Task.STATUS_PENDING) 38 | if options['lost']: 39 | statuses.append(Task.STATUS_LOST) 40 | query['status__in'] = statuses 41 | 42 | if options.get('prior', None): 43 | query['submitted__lte'] = ( 44 | timezone.now() 45 | - datetime.timedelta(days=int(options['prior'])) 46 | ) 47 | 48 | objs = Task.objects.filter(**query) 49 | self.stdout.write('Removing {} objects.'.format(len(objs))) 50 | if not options['dry_run']: 51 | objs.delete() 52 | -------------------------------------------------------------------------------- /cq/utils.py: -------------------------------------------------------------------------------- 1 | import time 2 | from contextlib import contextmanager 3 | import six 4 | import inspect 5 | import importlib 6 | import logging 7 | 8 | from django.conf import settings 9 | from redis.exceptions import RedisError 10 | from django_redis import get_redis_connection 11 | 12 | 13 | logger = logging.getLogger('cq') 14 | 15 | 16 | def get_redis_key(key): 17 | prefix = getattr(settings, 'CQ_PREFIX', '') 18 | if prefix: 19 | prefix += ':' 20 | return f'{prefix}{key}' 21 | 22 | 23 | def to_import_string(func): 24 | if inspect.isfunction(func) or inspect.isbuiltin(func): 25 | name = '{0}.{1}'.format(func.__module__, func.__name__) 26 | elif isinstance(func, six.string_types): 27 | name = str(func) 28 | elif inspect.isclass(func): 29 | return '{0}.{1}'.format(func.__module__, func.__name__) 30 | else: 31 | msg = 'Expected a callable or a string, but got: {}'.format(func) 32 | raise TypeError(msg) 33 | return name 34 | 35 | 36 | def import_attribute(name): 37 | """Return an attribute from a dotted path name (e.g. "path.to.func"). 38 | """ 39 | module_name, attribute = name.rsplit('.', 1) 40 | module = importlib.import_module(module_name) 41 | return getattr(module, attribute) 42 | 43 | 44 | @contextmanager 45 | def redis_connection(retries=3, sleep_time=0.5): 46 | while 1: 47 | try: 48 | conn = get_redis_connection() 49 | break 50 | except RedisError: 51 | if retries is None or retries == 0: 52 | raise 53 | retries -= 1 54 | time.sleep(sleep_time) 55 | try: 56 | yield conn 57 | finally: 58 | pass 59 | # This is actually not needed. The call to `get_redis_connection` 60 | # shares a single connection. 61 | # conn.release() 62 | -------------------------------------------------------------------------------- /cq/consumers.py: -------------------------------------------------------------------------------- 1 | import logging 2 | 3 | from channels.consumer import SyncConsumer 4 | from django.db import transaction 5 | 6 | from .models import Task 7 | from .task import SerialTask, TaskFunc 8 | 9 | logger = logging.getLogger('cq') 10 | 11 | 12 | class CQConsumer(SyncConsumer): 13 | def run_task(self, message): 14 | run_task(message['task_id']) 15 | 16 | 17 | def run_task(task_id): 18 | try: 19 | task = Task.objects.get(id=task_id) 20 | except Task.DoesNotExist: 21 | logger.error(f'Failed to find task with id {task_id}') 22 | return 23 | func_name = task.signature['func_name'] 24 | if task.status == Task.STATUS_REVOKED: 25 | logger.info('Not running revoked task: {}'.format(func_name)) 26 | return 27 | logger.info('{}: running task {}'.format(task_id, func_name)) 28 | task.pre_start() 29 | task_func = TaskFunc.get_task(func_name) 30 | if task_func.atomic: 31 | with transaction.atomic(): 32 | _do_run_task(task_func, task) 33 | else: 34 | _do_run_task(task_func, task) 35 | 36 | 37 | def _do_run_task(task_func, task): 38 | try: 39 | result = task.start(pre_start=False) 40 | except Exception as err: 41 | handle_failure(task_func, task, err) 42 | else: 43 | if isinstance(result, Task): 44 | task.waiting(task=result) 45 | else: 46 | if isinstance(result, SerialTask): 47 | result = result.result 48 | if task.subtasks.exists(): 49 | task.waiting(result=result) 50 | else: 51 | task.success(result) 52 | 53 | 54 | def handle_failure(task_func, task, err): 55 | """Decide whether to retry a failed task. 56 | """ 57 | if task.retries >= task_func.retries or not task_func.match_exceptions(err): 58 | task.failure(err) 59 | else: 60 | task.failure(err, retry=True) 61 | -------------------------------------------------------------------------------- /cq/scheduler.py: -------------------------------------------------------------------------------- 1 | import logging 2 | import time 3 | from datetime import timedelta 4 | from traceback import format_exc 5 | 6 | from django.core.cache import cache 7 | from django.db.utils import ProgrammingError 8 | from django.utils import timezone 9 | 10 | from .models import RepeatingTask 11 | from .utils import redis_connection, get_redis_key 12 | 13 | logger = logging.getLogger('cq') 14 | 15 | 16 | def perform_scheduling(): 17 | logger.debug('cq-scheduler: performing scheduling started') 18 | with cache.lock('cq:scheduler:lock', timeout=10): 19 | logger.debug('cq-scheduler: checking for scheduled tasks') 20 | now = timezone.now() 21 | try: 22 | rtasks = RepeatingTask.objects.filter(next_run__lte=now) 23 | logger.info('cq-scheduler: have {} repeating task(s) ready'.format(rtasks.count())) 24 | for rt in rtasks: 25 | try: 26 | rt.submit() 27 | except Exception as e: 28 | # Don't terminate if a submit fails. 29 | logger.error(format_exc()) 30 | except ProgrammingError: 31 | logger.warning('CQ scheduler not running, DB is out of date.') 32 | logger.debug('cq-scheduler: performing scheduling finished') 33 | 34 | 35 | def scheduler_internal(): 36 | logger.debug('cq-scheduler: determining winning scheduler') 37 | am_scheduler = False 38 | with redis_connection() as conn: 39 | if conn.set(get_redis_key('cq:scheduler'), 'dummy', nx=True, ex=30): 40 | # conn.expire('cq:scheduler', 30) 41 | am_scheduler = True 42 | if am_scheduler: 43 | logger.debug('cq-scheduler: winner') 44 | perform_scheduling() 45 | else: 46 | logger.debug('cq-scheduler: loser') 47 | now = timezone.now() 48 | delay = ((now + timedelta(minutes=1)).replace(second=0, microsecond=0) - now).total_seconds() 49 | logger.debug('cq-scheduler: waiting {} seconds for next schedule attempt'.format(delay)) 50 | time.sleep(delay) 51 | 52 | 53 | def scheduler(*args, **kwargs): 54 | logger.info('cq-scheduler: Scheduler thread active.') 55 | while 1: 56 | try: 57 | scheduler_internal() 58 | except Exception as e: 59 | logger.error(format_exc()) 60 | time.sleep(0.5) 61 | -------------------------------------------------------------------------------- /cq/tasks.py: -------------------------------------------------------------------------------- 1 | from datetime import timedelta 2 | 3 | from django.core.management import call_command 4 | from django.utils import timezone 5 | 6 | from .decorators import task 7 | from .models import Task 8 | from .utils import redis_connection 9 | 10 | 11 | @task 12 | def clean_up(task, *args): 13 | """ Remove stale tasks. 14 | 15 | Only remove tasks that have succeeded, are older than the TTl, have 16 | no dependencies that are still incomplete. 17 | """ 18 | now = timezone.now() 19 | to_del = Task.objects.filter( 20 | status=Task.STATUS_SUCCESS, 21 | result_expiry__lte=now 22 | ) 23 | if len(to_del): 24 | task.log('Cleaned up: {}'.format(', '.join([str(o.id) for o in to_del]))) 25 | to_del.delete() 26 | 27 | 28 | @task 29 | def clear_logs(cqt): 30 | """ Remove all logs from REDIS. 31 | """ 32 | with redis_connection() as con: 33 | for key in con.keys('cq:*:logs'): 34 | con.delete(key) 35 | 36 | 37 | @task 38 | def retry_tasks(cqtask, *args, **kwargs): 39 | retry_delay = kwargs.pop('retry_delay', 1) 40 | tasks = Task.objects.filter(status=Task.STATUS_RETRY) 41 | launched = 0 42 | for t in tasks: 43 | next_retry = (t.retries ** 2) * timedelta(minutes=retry_delay) 44 | now = timezone.now() 45 | if not t.last_retry or (now - t.last_retry) >= next_retry: 46 | cqtask.log('Retrying: {}'.format(t.id)) 47 | t.retry() 48 | launched += 1 49 | if launched >= 20: # cap at 20 50 | break 51 | 52 | 53 | @task 54 | def maintenance(task): 55 | retry_tasks(task=task) 56 | clean_up(task=task) 57 | 58 | 59 | @task 60 | def call_command_task(task, *args, **kwargs): 61 | """A wrapper to call management commands. 62 | """ 63 | return call_command(*args, **kwargs) 64 | 65 | 66 | @task 67 | def memory_details(task, method=None): 68 | if method == 'pympler': 69 | from pympler import muppy, summary 70 | all_objs = muppy.get_objects() 71 | summary.print_(summary.summarize(all_objs)) 72 | elif method == 'mem_top': 73 | from mem_top import mem_top 74 | task.log(mem_top()) 75 | else: 76 | import subprocess 77 | result = subprocess.check_output( 78 | 'ps --no-headers -eo pmem,vsize,rss,pid,cmd | sort -k 1 -nr', 79 | shell=True 80 | ) 81 | task.log('\n' + result.decode('utf8')) 82 | -------------------------------------------------------------------------------- /cq/migrations/0001_initial.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | # Generated by Django 1.10 on 2016-08-22 09:02 3 | from __future__ import unicode_literals 4 | 5 | import cq.models 6 | import django.contrib.postgres.fields.jsonb 7 | from django.db import migrations, models 8 | import django.db.models.deletion 9 | import uuid 10 | 11 | 12 | class Migration(migrations.Migration): 13 | 14 | initial = True 15 | 16 | dependencies = [ 17 | ] 18 | 19 | operations = [ 20 | migrations.CreateModel( 21 | name='RepeatingTask', 22 | fields=[ 23 | ('id', models.AutoField(auto_created=True, primary_key=True, serialize=False, verbose_name='ID')), 24 | ('crontab', models.CharField(default='* * * * *', help_text='Minute Hour Day Month Weekday', max_length=100, validators=[cq.models.validate_cron])), 25 | ('func_name', models.CharField(max_length=256, validators=[cq.models.validate_func_name])), 26 | ('args', django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=list)), 27 | ('kwargs', django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=dict)), 28 | ('last_run', models.DateTimeField(blank=True, null=True)), 29 | ('next_run', models.DateTimeField(blank=True, db_index=True, null=True)), 30 | ('coalesce', models.BooleanField(default=True)), 31 | ('result_ttl', models.PositiveIntegerField(blank=True, default=1800)), 32 | ], 33 | ), 34 | migrations.CreateModel( 35 | name='Task', 36 | fields=[ 37 | ('id', models.UUIDField(default=uuid.uuid4, editable=False, primary_key=True, serialize=False)), 38 | ('status', models.CharField(choices=[('P', 'Pending'), ('Y', 'Retry'), ('Q', 'Queued'), ('R', 'Running'), ('F', 'Failure'), ('S', 'Success'), ('W', 'Waiting'), ('I', 'Incomplete'), ('L', 'Lost')], db_index=True, default='P', max_length=1)), 39 | ('signature', django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=dict)), 40 | ('details', django.contrib.postgres.fields.jsonb.JSONField(blank=True, default=dict)), 41 | ('submitted', models.DateTimeField(auto_now_add=True)), 42 | ('started', models.DateTimeField(blank=True, null=True)), 43 | ('parent', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='subtasks', to='cq.Task')), 44 | ('previous', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='next', to='cq.Task')), 45 | ('waiting_on', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, to='cq.Task')), 46 | ('at_risk', models.CharField(choices=[('N', 'None'), ('Q', 'Queued'), ('R', 'Running')], default='N', max_length=1)), 47 | ('finished', models.DateTimeField(blank=True, null=True)), 48 | ('result_expiry', models.DateTimeField(blank=True, null=True)), 49 | ('result_ttl', models.PositiveIntegerField(blank=True, default=1800)), 50 | ], 51 | options={ 52 | 'ordering': ('-submitted',), 53 | }, 54 | ), 55 | ] 56 | -------------------------------------------------------------------------------- /cq/tests/test_logs.py: -------------------------------------------------------------------------------- 1 | from django.test import override_settings 2 | from channels.tests import ( 3 | TransactionChannelTestCase 4 | ) 5 | from django_redis import get_redis_connection 6 | 7 | from ..decorators import task 8 | from ..consumers import run_task 9 | from ..tasks import clear_logs 10 | from .test_base import SilentMixin 11 | 12 | 13 | @task 14 | def task_a(cqt, publish=True): 15 | cqt.log('First log from task_a.', publish=publish) 16 | cqt.subtask(task_b, {'publish': publish}) 17 | cqt.log('Second log from task_a.', publish=publish) 18 | return 'a' 19 | 20 | 21 | @task 22 | def task_b(cqt, publish=True): 23 | cqt.log('Log from task_b.', publish=publish) 24 | return 'b' 25 | 26 | 27 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 28 | class BasicLogTestCase(SilentMixin, TransactionChannelTestCase): 29 | def test_logs_stored_on_model(self): 30 | task = task_b.delay() 31 | while run_task(self.get_next_message('cq-tasks')): 32 | pass 33 | task.wait() 34 | logs = task.format_logs() 35 | self.assertEqual(logs, ( 36 | 'Log from task_b.' 37 | )) 38 | 39 | def test_logs_from_subtasks(self): 40 | task = task_a.delay() 41 | while 1: 42 | msg = self.get_next_message('cq-tasks') 43 | if not msg: 44 | break 45 | run_task(msg) 46 | task.wait() 47 | logs = task.format_logs() 48 | self.assertEqual(task.status, task.STATUS_SUCCESS) 49 | self.assertEqual(logs, ( 50 | 'First log from task_a.\n' 51 | 'Second log from task_a.\n' 52 | 'Log from task_b.' 53 | )) 54 | 55 | def test_no_publish(self): 56 | task = task_a.delay(publish=False) 57 | while 1: 58 | msg = self.get_next_message('cq-tasks') 59 | if not msg: 60 | break 61 | run_task(msg) 62 | task.wait() 63 | logs = task.format_logs() 64 | self.assertEqual(task.status, task.STATUS_SUCCESS) 65 | self.assertEqual(logs, ( 66 | 'First log from task_a.\n' 67 | 'Second log from task_a.\n' 68 | 'Log from task_b.' 69 | )) 70 | 71 | def test_logs_cleared_from_redis(self): 72 | task = task_a.delay() 73 | while 1: 74 | msg = self.get_next_message('cq-tasks') 75 | if not msg: 76 | break 77 | run_task(msg) 78 | task.wait() 79 | con = get_redis_connection() 80 | res = con.keys('cq:*:logs') 81 | key = task._get_log_key() 82 | res = con.keys('cq:*:logs') 83 | self.assertFalse(key.encode() in res) 84 | 85 | 86 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 87 | class ClearLogsTestCase(SilentMixin, TransactionChannelTestCase): 88 | def test_clear_logs(self): 89 | con = get_redis_connection() 90 | con.set('cq:test:logs', '') 91 | con.set('cq:another:logs', '') 92 | self.assertNotEqual(con.keys('cq:*:logs'), []) 93 | clear_logs() 94 | self.assertEqual(con.keys('cq:*:logs'), []) 95 | 96 | def test_clear_logs_when_none_exist(self): 97 | con = get_redis_connection() 98 | self.assertEqual(con.keys('cq:*:logs'), []) 99 | clear_logs() 100 | self.assertEqual(con.keys('cq:*:logs'), []) 101 | -------------------------------------------------------------------------------- /cq/management/commands/cq_load_test.py: -------------------------------------------------------------------------------- 1 | import time 2 | import sys 3 | 4 | from django.core.management.base import BaseCommand, CommandError 5 | from asgi_redis import RedisChannelLayer 6 | from cq.decorators import task 7 | from cq.models import Task, delay 8 | 9 | 10 | @task 11 | def task_j(task, a, b): 12 | return a + b 13 | 14 | 15 | @task 16 | def task_i(task): 17 | return 3 18 | 19 | 20 | @task 21 | def task_f(task): 22 | return task_i.delay().chain(task_j, (2,)) 23 | 24 | 25 | @task 26 | def task_c(task, value): 27 | if value == 'd': 28 | try: 29 | Task.objects.get_or_create(details={'failed': 'failed'}) 30 | except Task.MultipleObjectsReturned: 31 | pass 32 | raise Exception('Task D, should fail.') 33 | try: 34 | Task.objects.get_or_create(details={'failed': 'okay'}) 35 | except Task.MultipleObjectsReturned: 36 | pass 37 | return value 38 | 39 | 40 | @task 41 | def task_a(task): 42 | task.subtask(task_c, ('c',)) 43 | task.subtask(task_c, ('d',)) 44 | return task_c.delay('e') 45 | 46 | 47 | @task 48 | def task_b(task): 49 | task.subtask(task_c, ('g',)) 50 | task.subtask(task_c, ('h',)) 51 | return task_f.delay() 52 | 53 | 54 | @task 55 | def task_trunk(task): 56 | task.subtask(task_a) 57 | task.subtask(task_b) 58 | 59 | 60 | @task 61 | def task_root(task): 62 | for ii in range(50): 63 | task.subtask(task_trunk) 64 | return 'root' 65 | 66 | 67 | class Command(BaseCommand): 68 | help = 'CQ load test.' 69 | 70 | def add_arguments(self, parser): 71 | parser.add_argument('--lost', action='store_true') 72 | 73 | def handle(self, *args, **options): 74 | if options.get('lost', True): 75 | self.launch_lost() 76 | else: 77 | self.launch_one() 78 | 79 | def launch_lost(self): 80 | for ii in range(200): 81 | delay(task_root, (), {}, submit=False, status=Task.STATUS_RETRY) 82 | 83 | def launch_one(self): 84 | root = task_root.delay() 85 | root.wait() 86 | for trunk in root.subtasks.all(): 87 | self.check_trunk(trunk) 88 | 89 | # Should not have committed the dummy task. 90 | assert not Task.objects.filter(details__failed='failed').exists(), 'Should not have committed task.' 91 | 92 | # Should have committed this dummy task. 93 | assert Task.objects.filter(details__failed='okay').exists(), 'Should have committed this one.' 94 | 95 | def assertEqual(self, a, b, msg): 96 | if a != b: 97 | self.stdout.write(msg) 98 | self.stdout.write(' {} != {}'.format(a, b)) 99 | sys.exit(0) 100 | 101 | def check_trunk(self, task): 102 | task.wait() 103 | self.assertEqual(task.status, Task.STATUS_INCOMPLETE, 'trunk should be incomplete.') 104 | self.check_a(task, task.subtasks.get(signature__func_name='cq.management.commands.cq_load_test.task_a')) 105 | self.check_b(task, task.subtasks.get(signature__func_name='cq.management.commands.cq_load_test.task_b')) 106 | 107 | def check_a(self, trunk, task): 108 | task.wait() 109 | assert task.status == Task.STATUS_INCOMPLETE, 'task_a should be incomplete.' 110 | # self.assertEqual(task.result, 'e', 'task_a result should be "e"') 111 | self.check_c(task.subtasks.get(details__result='c')) 112 | self.check_d(task.subtasks.get(details__error__isnull=False)) 113 | self.check_e(task.subtasks.get(details__result='e')) 114 | 115 | def check_b(self, trunk, task): 116 | task.wait() 117 | self.assertEqual(task.status, Task.STATUS_SUCCESS, 'task_b should have succeeded.') 118 | self.assertEqual(task.result, 5, 'task_b result should be 5') 119 | self.check_f(task.subtasks.get(details__result=5)) 120 | self.check_g(task.subtasks.get(details__result='g')) 121 | self.check_h(task.subtasks.get(details__result='h')) 122 | 123 | def check_c(self, task): 124 | task.wait() 125 | assert task.result == 'c', 'task_c result should be "c"' 126 | 127 | def check_d(self, task): 128 | task.wait() 129 | assert task.result == None, 'task_d result should be "None"' 130 | self.assertEqual(task.error, 'Task D, should fail.', 'task_d should have an error') 131 | 132 | def check_e(self, task): 133 | task.wait() 134 | assert task.result == 'e', 'task_e result should be "e"' 135 | 136 | def check_f(self, task): 137 | task.wait() 138 | self.assertEqual(task.result, 5, 'task_f result should be 5') 139 | 140 | def check_g(self, task): 141 | task.wait() 142 | assert task.result == 'g', 'task_g result should be "g"' 143 | 144 | def check_h(self, task): 145 | task.wait() 146 | assert task.result == 'h', 'task_h result should be "h"' 147 | -------------------------------------------------------------------------------- /cq/tests/test_force_chain.py: -------------------------------------------------------------------------------- 1 | from django.test import override_settings 2 | from channels.tests import ( 3 | TransactionChannelTestCase 4 | ) 5 | 6 | from ..decorators import task 7 | from ..consumers import run_task 8 | from ..models import Task 9 | from .test_base import SilentMixin 10 | 11 | 12 | @task 13 | def task_a(task, force_chain=True): 14 | task.subtask(task_b).chain(task_e, force_chain=force_chain).chain(task_h, force_chain=force_chain) 15 | return 'a' 16 | 17 | 18 | @task 19 | def task_b(task): 20 | return task.subtask(task_c).chain(task_d) 21 | 22 | 23 | @task 24 | def task_c(task): 25 | return 'c' 26 | 27 | 28 | @task 29 | def task_d(task): 30 | raise Exception('d') 31 | 32 | 33 | @task 34 | def task_e(task): 35 | return task.subtask(task_f).chain(task_g) 36 | 37 | 38 | @task 39 | def task_f(task): 40 | return 'f' 41 | 42 | 43 | @task 44 | def task_g(task): 45 | return 'g' 46 | 47 | 48 | @task 49 | def task_h(task): 50 | return task.subtask(task_i).chain(task_j) 51 | 52 | 53 | @task 54 | def task_i(task): 55 | return 'i' 56 | 57 | 58 | @task 59 | def task_j(task): 60 | raise Exception('j') 61 | 62 | 63 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 64 | class ForceChainTestCase(SilentMixin, TransactionChannelTestCase): 65 | def get_task(self, func_name): 66 | for task in Task.objects.all(): 67 | if task.func_name == func_name: 68 | return task 69 | return None 70 | 71 | def test_force(self): 72 | task = task_a.delay() 73 | try: 74 | while 1: 75 | run_task(self.get_next_message('cq-tasks', require=True)) 76 | except: 77 | pass 78 | task.wait() 79 | 80 | self.assertEqual(task.status, 'I') 81 | self.assertEqual(task.subtasks.all().count(), 3) 82 | self.assertIn('cq.tests.test_force_chain.task_b', [t.func_name for t in task.subtasks.all()]) 83 | self.assertIn('cq.tests.test_force_chain.task_e', [t.func_name for t in task.subtasks.all()]) 84 | self.assertIn('cq.tests.test_force_chain.task_h', [t.func_name for t in task.subtasks.all()]) 85 | 86 | task_b = self.get_task('cq.tests.test_force_chain.task_b') 87 | task_c = self.get_task('cq.tests.test_force_chain.task_c') 88 | task_d = self.get_task('cq.tests.test_force_chain.task_d') 89 | self.assertEqual(task_b.status, 'I') 90 | self.assertEqual(task_c.status, 'S') 91 | self.assertEqual(task_d.status, 'F') 92 | 93 | task_e = self.get_task('cq.tests.test_force_chain.task_e') 94 | task_f = self.get_task('cq.tests.test_force_chain.task_f') 95 | task_g = self.get_task('cq.tests.test_force_chain.task_g') 96 | self.assertEqual(task_e.status, 'S') 97 | self.assertEqual(task_f.status, 'S') 98 | self.assertEqual(task_g.status, 'S') 99 | 100 | task_h = self.get_task('cq.tests.test_force_chain.task_h') 101 | task_i = self.get_task('cq.tests.test_force_chain.task_i') 102 | task_j = self.get_task('cq.tests.test_force_chain.task_j') 103 | self.assertEqual(task_h.status, 'I') 104 | self.assertEqual(task_i.status, 'S') 105 | self.assertEqual(task_j.status, 'F') 106 | 107 | def test_no_force(self): 108 | task = task_a.delay(force_chain=False) 109 | try: 110 | while 1: 111 | run_task(self.get_next_message('cq-tasks', require=True)) 112 | except: 113 | pass 114 | task.wait() 115 | 116 | self.assertEqual(task.status, 'I') 117 | self.assertEqual(task.subtasks.all().count(), 3) 118 | self.assertIn('cq.tests.test_force_chain.task_b', [t.func_name for t in task.subtasks.all()]) 119 | self.assertIn('cq.tests.test_force_chain.task_e', [t.func_name for t in task.subtasks.all()]) 120 | self.assertIn('cq.tests.test_force_chain.task_h', [t.func_name for t in task.subtasks.all()]) 121 | 122 | task_b = self.get_task('cq.tests.test_force_chain.task_b') 123 | task_c = self.get_task('cq.tests.test_force_chain.task_c') 124 | task_d = self.get_task('cq.tests.test_force_chain.task_d') 125 | self.assertEqual(task_b.status, 'I') 126 | self.assertEqual(task_c.status, 'S') 127 | self.assertEqual(task_d.status, 'F') 128 | 129 | task_e = self.get_task('cq.tests.test_force_chain.task_e') 130 | task_f = self.get_task('cq.tests.test_force_chain.task_f') 131 | task_g = self.get_task('cq.tests.test_force_chain.task_g') 132 | self.assertEqual(task_e.status, 'P') 133 | self.assertEqual(task_f, None) 134 | self.assertEqual(task_g, None) 135 | 136 | task_h = self.get_task('cq.tests.test_force_chain.task_h') 137 | task_i = self.get_task('cq.tests.test_force_chain.task_i') 138 | task_j = self.get_task('cq.tests.test_force_chain.task_j') 139 | self.assertEqual(task_h.status, 'P') 140 | self.assertEqual(task_i, None) 141 | self.assertEqual(task_j, None) 142 | -------------------------------------------------------------------------------- /cq/task.py: -------------------------------------------------------------------------------- 1 | import inspect 2 | import logging 3 | import uuid 4 | from functools import wraps 5 | 6 | from django import forms 7 | from django.conf import settings 8 | 9 | from .utils import import_attribute, to_import_string 10 | 11 | logger = logging.getLogger('cq') 12 | 13 | 14 | def to_func_name(func): 15 | # Try to convert to a name before returing. Will default 16 | # to the import string. 17 | return TaskFunc.get_name(to_import_string(func)) 18 | 19 | 20 | def to_signature(func, args, kwargs): 21 | return { 22 | 'func_name': to_func_name(func), 23 | 'args': args, 24 | 'kwargs': kwargs 25 | } 26 | 27 | 28 | def from_signature(sig): 29 | func = TaskFunc.get_task_func(sig['func_name']) 30 | return (func, tuple(sig.get('args', ())), sig.get('kwargs', {})) 31 | 32 | 33 | class TaskFunc(object): 34 | task_table = {} 35 | name_table = {} 36 | 37 | def __init__(self, name=None, atomic=True, **kwargs): 38 | self.name = name 39 | self.atomic = atomic 40 | self.retries = kwargs.get('retries', 0) 41 | self.retry_exceptions = kwargs.get('retry_exceptions', []) 42 | if not isinstance(self.retry_exceptions, (list, tuple)): 43 | self.retry_exceptions = [self.retry_exceptions] 44 | 45 | def __call__(self, func): 46 | @wraps(func) 47 | def _delay_args(args=(), kwargs={}, **kw): 48 | if getattr(settings, 'CQ_SERIAL', False): 49 | # Create a SerialTask here to make sure we end up 50 | # returning the task instead of the result. 51 | kwargs['task'] = SerialTask() 52 | return self.wrapper(func, args, kwargs, **kw) 53 | else: 54 | from .models import delay 55 | return delay(func, args, kwargs, **kw) 56 | 57 | @wraps(func) 58 | def _delay(*args, **kwargs): 59 | return _delay_args(args, kwargs) 60 | 61 | @wraps(func) 62 | def wrapper(*args, **kwargs): 63 | return self.wrapper(func, args, kwargs) 64 | 65 | # Add the task to the tables. 66 | self.func = wrapper 67 | func_name = to_import_string(func) 68 | self.task_table[func_name] = self 69 | logger.debug('Adding task definition: {}'.format(func_name)) 70 | if self.name and self.name != func_name: 71 | self.task_table[self.name] = self 72 | logger.debug('Adding task definition: {}'.format(self.name)) 73 | if self.name: 74 | self.name_table[self.name] = func_name 75 | 76 | wrapper.delay = _delay 77 | wrapper.delay_args = _delay_args 78 | return wrapper 79 | 80 | def wrapper(self, func, args, kwargs): 81 | task = kwargs.pop('task', None) 82 | direct = task is None 83 | serial = isinstance(task, SerialTask) 84 | if direct or serial: 85 | task = SerialTask(parent=task, previous=task) 86 | try: 87 | result = func(task, *args, **kwargs) 88 | except Exception as err: 89 | if serial: 90 | for func, args, kwargs in task._errbacks: 91 | func(*((task, err,) + tuple(args)), **kwargs) 92 | raise 93 | if direct or serial: 94 | while isinstance(result, SerialTask): 95 | result = result.result 96 | task.result = result 97 | if direct: 98 | return task.result 99 | else: 100 | return task 101 | else: 102 | return result 103 | 104 | @classmethod 105 | def get_task(cls, name): 106 | return cls.task_table[name] 107 | 108 | @classmethod 109 | def get_task_func(cls, name): 110 | try: 111 | return cls.get_task(name).func 112 | except KeyError: 113 | pass 114 | return import_attribute(name) 115 | 116 | @classmethod 117 | def get_name(cls, func_name): 118 | try: 119 | task = cls.task_table[func_name] 120 | except KeyError: 121 | return func_name 122 | return task.name or func_name 123 | 124 | def match_exceptions(self, error): 125 | def _is_exception(value): 126 | if not inspect.isclass(value): 127 | value = value.__class__ 128 | return issubclass(value, Exception) 129 | 130 | def _is_instance(value, cls): 131 | if not inspect.isclass(cls): 132 | cls = cls.__class__ 133 | return isinstance(value, cls) 134 | 135 | if not self.retry_exceptions: 136 | return True 137 | for ex in self.retry_exceptions: 138 | if callable(ex) and not _is_exception(ex) and ex(error): 139 | return True 140 | elif _is_instance(error, ex): 141 | return True 142 | return False 143 | 144 | 145 | class SerialTask(object): 146 | """ 147 | """ 148 | def __init__(self, result=None, parent=None, previous=None): 149 | self.id = uuid.uuid4() 150 | self.result = result 151 | self.parent = parent 152 | self.previous = previous 153 | self._errbacks = [] 154 | 155 | # Mimic behavior of 'logs' from the model. 156 | self.logs = [] 157 | 158 | def subtask(self, func, args=(), kwargs={}, **kw): 159 | # Note: A serial task will automatically be created. 160 | return func(*args, task=self, **kwargs) 161 | 162 | def chain(self, func, args=(), kwargs={}, **kw): 163 | # Note: A serial task will automatically be created. 164 | return func(*args, task=self, **kwargs) 165 | 166 | def errorback(self, func, args=(), kwargs={}): 167 | self._errbacks.append((func, args, kwargs)) 168 | 169 | def log(self, msg, level=logging.INFO, **kwargs): 170 | logger.log(level, msg) 171 | 172 | 173 | class FuncNameWidget(forms.TextInput): 174 | def __init__(self, *args, **kwargs): 175 | super().__init__(*args, **kwargs) 176 | self._name = 'func_name' 177 | self._list = sorted(list(TaskFunc.task_table.keys())) 178 | self.attrs.update({'list': 'list__%s' % self._name}) 179 | 180 | def render(self, name, value, attrs=None, renderer=None): 181 | text_html = super().render(name, value, attrs=attrs, renderer=renderer) 182 | data_list = '' % self._name 183 | for item in self._list: 184 | data_list += '' 186 | return (text_html + data_list) 187 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Django-CQ 2 | [![PyPI version](https://badge.fury.io/py/django-cq.svg)](https://badge.fury.io/py/django-cq) 3 | 4 | ## Description 5 | 6 | An attempt to implement a distributed task queue for use with Django channels. 7 | Modelled after RQ and Celery, complex task workflows are possible, all leveraging 8 | the Channels machinery. 9 | 10 | 11 | ## Why 12 | 13 | There are three reasons: 14 | 15 | 1. Aiming for more fault tolerant tasks. There are many occasions where information 16 | regarding how tests are progressing is needed to be stored persistently. For 17 | important tasks, this should be stored even in the case of a Redis fault, or if 18 | the worker goes down. 19 | 20 | 2. Prefer to leverage the same machinery as channels. 21 | 22 | 3. Would like to have a little extra functionality surrounding subtasks that didn't 23 | seem to be available via Celery or RQ. 24 | 25 | 26 | ## Limitations 27 | 28 | There are two limitations involved: 29 | 30 | * REDIS must be used as the Django cache. 31 | 32 | * `asgi_redis` must be used as the channel backend. 33 | 34 | There is work being done to remove these restrictions. 35 | 36 | 37 | ## Installation 38 | 39 | Use pip if you can: 40 | 41 | ```bash 42 | pip install django-cq 43 | ``` 44 | 45 | or live on the edge: 46 | 47 | ```bash 48 | pip install -e https://github.com/furious-luke/django-cq.git#egg=django-cq 49 | ``` 50 | 51 | Add the package to your settings file: 52 | 53 | ```python 54 | INSTALLED_APPS = [ 55 | 'cq', 56 | ... 57 | ] 58 | ``` 59 | 60 | And include the routing information in your channel routing list: 61 | 62 | ```python 63 | channel_routing = [ 64 | include('cq.routing.channel_routing'), 65 | ... 66 | ] 67 | ``` 68 | 69 | You'll need to migrate to include the models: 70 | 71 | ```bash 72 | ./manage.py migrate 73 | ``` 74 | 75 | You'll most likely want to create a new channel layer for your CQ 76 | tasks. The default layer has a short time-to-live on the channel messages, 77 | which causes slightly long running tasks to kill off any queued messages. 78 | Update your settings file to include the following: 79 | 80 | ```python 81 | CHANNEL_LAYERS = { 82 | 'default': { 83 | ... 84 | }, 85 | 'long': { 86 | 'BACKEND': 'asgi_redis.RedisChannelLayer', 87 | 'CONFIG': { 88 | 'hosts': [REDIS_URL], 89 | 'expiry': 1800, 90 | 'channel_capacity': { 91 | 'cq-tasks': 1000 92 | } 93 | }, 94 | 'ROUTING': 'path.to.your.channels.channel_routing', 95 | }, 96 | } 97 | 98 | CQ_CHANNEL_LAYER = 'long' 99 | ``` 100 | 101 | In order to process messages sent on the "cq-tasks" channel a worker 102 | process needs to be launched: 103 | 104 | ``` 105 | ./manage.py cq_runworker 106 | ``` 107 | 108 | 109 | ## Tasks 110 | 111 | Basic task usage is straight forward: 112 | 113 | ```python 114 | @task 115 | def send_email(cqt, addr): 116 | ... 117 | return 'OK' 118 | 119 | task = send_emails.delay('dummy@dummy.org') 120 | task.wait() 121 | print(task.result) # "OK" 122 | ``` 123 | 124 | Here, `cqt` is the task representation for the `send_email` task. This 125 | can be used to launch subtasks, chain subsequent tasks, amongst other 126 | things. 127 | 128 | Tasks may also be run in serial by just calling them: 129 | 130 | ```python 131 | result = send_email('dummy@dummy.org') 132 | print(result) # "OK" 133 | ``` 134 | 135 | 136 | ## Subtasks 137 | 138 | For more complex workflows, subtasks may be launched from within 139 | parent tasks: 140 | 141 | ```python 142 | @task 143 | def send_emails(cqt): 144 | ... 145 | for addr in email_addresses: 146 | cqt.subtask(send_email, addr) 147 | ... 148 | return 'OK' 149 | 150 | task = send_emails.delay() 151 | task.wait() 152 | print(task.result) # "OK" 153 | ``` 154 | 155 | The difference between a subtask and another task launched using `delay` from 156 | within a task is that the parent task of a subtask will not be marked as complete 157 | until all subtasks are also complete. 158 | 159 | ```python 160 | from cq.models import Task 161 | 162 | @task 163 | def parent(cqt): 164 | task_a.delay() # not a subtask 165 | cqt.subtask(task_b) # subtask 166 | 167 | parent.delay() 168 | parent.status == Task.STATUS_WAITING # True 169 | # once task_b completes 170 | parent.wait() 171 | parent.status == Task.STATUS_COMPLETE # True 172 | ``` 173 | 174 | 175 | ## Chained Tasks 176 | 177 | TODO 178 | 179 | ```python 180 | @task 181 | def calculate_something(cqt): 182 | return calc_a.delay(3).chain(add_a_to_4, (4,)) 183 | ``` 184 | 185 | 186 | ## Non-atomic Tasks 187 | 188 | By default every CQ task is atomic; no changes to the database will persist 189 | unless the task finishes without an exception. If you need to keep changes to 190 | the database, even in the event of an error, then use the `atomic` flag: 191 | 192 | ```python 193 | @task(atomic=False) 194 | def unsafe_task(cqt): 195 | pass 196 | ``` 197 | 198 | 199 | ## Logging 200 | 201 | For longer running tasks it's useful to be able to access an ongoing log 202 | of the task's progress. CQ tasks have a `log` method to send logging 203 | messages to both the standard Django log streams, and also cache them on 204 | the running task. 205 | 206 | ```python 207 | @task 208 | def long_task(cqt): 209 | cqt.log('standard old log') 210 | cqt.log('debugging log', logging.DEBUG) 211 | ``` 212 | 213 | If the current task is a subtask, the logs will go to the parent. 214 | This way there is a central task (the top-level task) which can be used 215 | to monitor the progress and status of a network of sub and chained tasks. 216 | 217 | ### Performance 218 | 219 | Due to the way logs are handled there can be issues with performance 220 | with a lot of frequent log messages. There are two ways to prevent this. 221 | 222 | Reduce the frequency of logs by setting `publish` to `False` on as many 223 | log calls as you can. This will cache the logs locally and store them 224 | on the next `publish=True` call. 225 | 226 | ```python 227 | @task 228 | def long_task(cqt): 229 | for ii in range(100): 230 | cqt.log('iteration %d' % ii, publish=False) 231 | cqt.log('done') # publish=True 232 | ``` 233 | 234 | Secondly, reducing the volume of logs may be accomplished by limiting the 235 | number of log lines that are kept. The `limit` option specifies this. The 236 | following will only keep 10 of the logged iterations: 237 | 238 | ```python 239 | @task 240 | def long_task(cqt): 241 | for ii in range(100): 242 | cqt.log('iteration %d' % ii, publish=False) 243 | cqt.log('done', limit=10) 244 | ``` 245 | 246 | ## Time-to-live 247 | 248 | TODO 249 | 250 | ## Repeating Tasks 251 | 252 | CQ comes with robust repeating tasks. There are two ways to create 253 | repeating tasks: 254 | 255 | 1. From the Django admin. 256 | 257 | 2. Using a data migration. 258 | 259 | From the admin, click into `cq` and `repeating tasks`. From there you 260 | can create a new repeating task, specifying the background task to call, 261 | and a CRON time for repetition. 262 | 263 | To create a repeating task from a migration, use the helper function 264 | `schedule_task`. 265 | 266 | ```python 267 | from django.db import migrations 268 | from cq.models import schedule_task 269 | 270 | from myapp.tasks import a_task 271 | 272 | 273 | def add_repeating(apps, scema_editor): 274 | RepeatingTask = apps.get_model('cq.RepeatingTask') 275 | schedule_task( 276 | RepeatingTask, 277 | '* * * * *', 278 | a_task 279 | ) 280 | 281 | 282 | class Migration(migrations.Migration): 283 | operations = [ 284 | migrations.RunPython(add_repeating, reverse_code=migrations.RunPython.noop) 285 | ] 286 | ``` 287 | 288 | 289 | ### Coalescing 290 | Pending or queued instances of a coalescing task will prevent other instances of the task from running. 291 | -------------------------------------------------------------------------------- /cq/tests/test_base.py: -------------------------------------------------------------------------------- 1 | import uuid 2 | from unittest.mock import patch 3 | from datetime import datetime 4 | import logging 5 | 6 | from django.test import TestCase, override_settings 7 | from django.utils import timezone 8 | try: 9 | from django.urls import reverse 10 | except ImportError: 11 | from django.core.urlresolvers import reverse 12 | from django.contrib.auth import get_user_model 13 | from channels.tests import ( 14 | TransactionChannelTestCase 15 | ) 16 | from channels.tests.base import ChannelTestCaseMixin 17 | from croniter import croniter 18 | from rest_framework.test import APITransactionTestCase 19 | 20 | from ..models import Task, RepeatingTask, delay, DuplicateSubmitError 21 | from ..decorators import task 22 | from ..consumers import run_task 23 | from ..backends import backend 24 | from ..tasks import retry_tasks 25 | 26 | 27 | User = get_user_model() 28 | 29 | 30 | class SilentMixin(object): 31 | def setUp(self): 32 | super().setUp() 33 | logging.disable(logging.CRITICAL) 34 | 35 | def tearDown(self): 36 | super().tearDown() 37 | logging.disable(logging.NOTSET) 38 | 39 | 40 | @task('a') 41 | def task_a(task, *args): 42 | if args: 43 | return args 44 | else: 45 | return 'a' 46 | 47 | 48 | @task 49 | def task_b(task): 50 | raise Exception('b') 51 | return 'b' 52 | 53 | 54 | @task 55 | def task_c(task): 56 | sub = task.subtask(task_a) 57 | return sub 58 | 59 | 60 | @task 61 | def task_d(task): 62 | task.subtask(task_a) 63 | return 'd' 64 | 65 | 66 | @task 67 | def task_e(task): 68 | sub = task.subtask(task_c) 69 | return sub 70 | 71 | 72 | @task 73 | def task_f(task): 74 | sub = task.subtask(task_b) 75 | return sub 76 | 77 | 78 | @task 79 | def task_g(task): 80 | sub = task.subtask(task_f) 81 | return sub 82 | 83 | 84 | @task 85 | def task_h(task): 86 | return task.subtask(task_i).chain(task_j, (2,)) 87 | 88 | 89 | @task 90 | def task_i(task): 91 | return 3 92 | 93 | 94 | @task('j') 95 | def task_j(task, a): 96 | return task.previous.result + a 97 | 98 | 99 | def errback(task, error): 100 | pass 101 | 102 | 103 | @task('k') 104 | def task_k(task, error=False): 105 | task.errorback(errback) 106 | if error: 107 | raise Exception 108 | 109 | 110 | @task(atomic=False) 111 | def task_l(task, uuid, error=False): 112 | Task.objects.create(id=uuid) 113 | if error: 114 | raise Exception 115 | 116 | 117 | @task 118 | def task_m(task, uuid, error=False): 119 | Task.objects.create(id=uuid) 120 | if error: 121 | raise Exception 122 | 123 | 124 | def check_retry(err): 125 | if 'hello' in str(err): 126 | return True 127 | return False 128 | 129 | 130 | class Retry(Exception): 131 | pass 132 | 133 | 134 | @task(retries=3, retry_exceptions=[Retry, check_retry]) 135 | def task_n(task, error=None): 136 | if error == 'exception': 137 | raise Retry 138 | elif error == 'func': 139 | raise Exception('hello') 140 | elif error == 'fail': 141 | raise Exception('nope') 142 | 143 | 144 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 145 | class DecoratorTestCase(SilentMixin, TransactionChannelTestCase): 146 | def test_adds_delay_function(self): 147 | self.assertTrue(hasattr(task_a, 'delay')) 148 | self.assertIsNot(task_a.delay, None) 149 | 150 | def test_task_is_still_a_function(self): 151 | self.assertEqual(task_a(), 'a') 152 | self.assertEqual(Task.objects.all().count(), 0) 153 | 154 | @patch('cq.models.Task.submit') 155 | def test_delay_creates_task(self, submit): 156 | before = timezone.now() 157 | task = task_a.delay() 158 | after = timezone.now() 159 | self.assertIsNot(task, None) 160 | self.assertGreater(len(str(task.id)), 10) 161 | self.assertGreater(task.submitted, before) 162 | self.assertLess(task.submitted, after) 163 | 164 | 165 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 166 | class TaskSuccessTestCase(SilentMixin, TransactionChannelTestCase): 167 | def test_something(self): 168 | task = task_a.delay() 169 | run_task(self.get_next_message('cq-tasks', require=True)) 170 | task.wait() 171 | self.assertEqual(task.status, task.STATUS_SUCCESS) 172 | 173 | 174 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 175 | class TaskFailureTestCase(SilentMixin, TransactionChannelTestCase): 176 | def test_something(self): 177 | task = task_b.delay() 178 | run_task(self.get_next_message('cq-tasks', require=True)) 179 | task.wait() 180 | self.assertEqual(task.status, task.STATUS_FAILURE) 181 | self.assertIsNot(task.error, None) 182 | 183 | def test_errorbacks(self): 184 | task = task_k.delay(error=True) 185 | run_task(self.get_next_message('cq-tasks', require=True)) 186 | task.wait() 187 | 188 | def test_retry_fail(self): 189 | task = task_n.delay(error='fail') 190 | run_task(self.get_next_message('cq-tasks', require=True)) 191 | task.wait() 192 | self.assertEqual(task.status, task.STATUS_FAILURE) 193 | self.assertEqual(task.retries, 0) 194 | 195 | def test_retry_success(self): 196 | task = task_n.delay() 197 | run_task(self.get_next_message('cq-tasks', require=True)) 198 | task.wait() 199 | self.assertEqual(task.status, task.STATUS_SUCCESS) 200 | self.assertEqual(task.retries, 0) 201 | 202 | def test_retry_exception(self): 203 | task = task_n.delay(error='exception') 204 | run_task(self.get_next_message('cq-tasks', require=True)) 205 | task.refresh_from_db() 206 | self.assertEqual(task.status, task.STATUS_RETRY) 207 | retry_tasks(retry_delay=0) 208 | run_task(self.get_next_message('cq-tasks', require=True)) 209 | task.refresh_from_db() 210 | self.assertEqual(task.status, task.STATUS_RETRY) 211 | self.assertEqual(task.retries, 1) 212 | retry_tasks(retry_delay=0) 213 | run_task(self.get_next_message('cq-tasks', require=True)) 214 | task.refresh_from_db() 215 | self.assertEqual(task.status, task.STATUS_RETRY) 216 | self.assertEqual(task.retries, 2) 217 | retry_tasks(retry_delay=0) 218 | run_task(self.get_next_message('cq-tasks', require=True)) 219 | task.refresh_from_db() 220 | self.assertEqual(task.status, task.STATUS_FAILURE) 221 | self.assertEqual(task.retries, 3) 222 | 223 | def test_retry_func(self): 224 | task = task_n.delay(error='func') 225 | run_task(self.get_next_message('cq-tasks', require=True)) 226 | task.refresh_from_db() 227 | self.assertEqual(task.status, task.STATUS_RETRY) 228 | retry_tasks(retry_delay=0) 229 | run_task(self.get_next_message('cq-tasks', require=True)) 230 | task.refresh_from_db() 231 | self.assertEqual(task.status, task.STATUS_RETRY) 232 | self.assertEqual(task.retries, 1) 233 | retry_tasks(retry_delay=0) 234 | run_task(self.get_next_message('cq-tasks', require=True)) 235 | task.refresh_from_db() 236 | self.assertEqual(task.status, task.STATUS_RETRY) 237 | self.assertEqual(task.retries, 2) 238 | retry_tasks(retry_delay=0) 239 | run_task(self.get_next_message('cq-tasks', require=True)) 240 | task.refresh_from_db() 241 | self.assertEqual(task.status, task.STATUS_FAILURE) 242 | self.assertEqual(task.retries, 3) 243 | 244 | 245 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 246 | class TaskRevokeTestCase(SilentMixin, TransactionChannelTestCase): 247 | def test_cancel_before_launch(self): 248 | task = delay(task_c, (), {}, submit=False) 249 | task.revoke() 250 | task.submit() 251 | run_task(self.get_next_message('cq-tasks', require=False)) 252 | run_task(self.get_next_message('cq-tasks', require=False)) 253 | task.wait() 254 | self.assertEqual(task.status, task.STATUS_REVOKED) 255 | self.assertEqual(task.subtasks.first(), None) 256 | 257 | def test_cancel_after_submit(self): 258 | task = task_c.delay() 259 | task.refresh_from_db() 260 | task.revoke() 261 | run_task(self.get_next_message('cq-tasks', require=False)) 262 | run_task(self.get_next_message('cq-tasks', require=False)) 263 | task.wait() 264 | self.assertEqual(task.status, task.STATUS_REVOKED) 265 | self.assertEqual(task.subtasks.first(), None) 266 | 267 | def test_cancel_after_run(self): 268 | task = task_c.delay() 269 | run_task(self.get_next_message('cq-tasks', require=False)) 270 | task.refresh_from_db() 271 | task.revoke() 272 | run_task(self.get_next_message('cq-tasks', require=False)) 273 | task.wait() 274 | self.assertEqual(task.status, task.STATUS_REVOKED) 275 | self.assertEqual(task.subtasks.first().status, task.STATUS_REVOKED) 276 | 277 | 278 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 279 | class AsyncSubtaskTestCase(SilentMixin, TransactionChannelTestCase): 280 | def test_returns_own_result(self): 281 | task = task_d.delay() 282 | run_task(self.get_next_message('cq-tasks', require=True)) 283 | run_task(self.get_next_message('cq-tasks', require=True)) 284 | task.wait() 285 | self.assertEqual(task.status, task.STATUS_SUCCESS) 286 | self.assertEqual(task.result, 'd') 287 | 288 | def test_returns_subtask_result(self): 289 | task = task_c.delay() 290 | run_task(self.get_next_message('cq-tasks', require=True)) 291 | run_task(self.get_next_message('cq-tasks', require=True)) 292 | task.wait() 293 | self.assertEqual(task.status, task.STATUS_SUCCESS) 294 | self.assertEqual(task.result, 'a') 295 | 296 | def test_returns_subsubtask_result(self): 297 | task = task_e.delay() 298 | run_task(self.get_next_message('cq-tasks', require=True)) 299 | run_task(self.get_next_message('cq-tasks', require=True)) 300 | run_task(self.get_next_message('cq-tasks', require=True)) 301 | task.wait() 302 | self.assertEqual(task.status, task.STATUS_SUCCESS) 303 | self.assertEqual(task.result, 'a') 304 | 305 | def test_parent_tasks_enter_waiting_state(self): 306 | task = task_e.delay() 307 | run_task(self.get_next_message('cq-tasks', require=True)) 308 | run_task(self.get_next_message('cq-tasks', require=True)) 309 | task.wait(500) 310 | task.refresh_from_db() 311 | self.assertEqual(task.status, task.STATUS_WAITING) 312 | subtask = task.subtasks.first() 313 | self.assertEqual(subtask.status, Task.STATUS_WAITING) 314 | subsubtask = subtask.subtasks.first() 315 | self.assertEqual(subsubtask.status, Task.STATUS_QUEUED) 316 | 317 | def test_returns_subtask_error(self): 318 | task = task_f.delay() 319 | run_task(self.get_next_message('cq-tasks', require=True)) 320 | run_task(self.get_next_message('cq-tasks', require=True)) 321 | task.wait() 322 | self.assertEqual(task.status, task.STATUS_INCOMPLETE) 323 | self.assertEqual(task.result, None) 324 | self.assertIsNot(task.error, None) 325 | 326 | def test_returns_subsubtask_error(self): 327 | task = task_g.delay() 328 | run_task(self.get_next_message('cq-tasks', require=True)) 329 | run_task(self.get_next_message('cq-tasks', require=True)) 330 | run_task(self.get_next_message('cq-tasks', require=True)) 331 | task.wait() 332 | self.assertEqual(task.status, task.STATUS_INCOMPLETE) 333 | self.assertEqual(task.result, None) 334 | self.assertIsNot(task.error, None) 335 | 336 | 337 | class SerialSubtaskTestCase(SilentMixin, TransactionChannelTestCase): 338 | def test_returns_own_result(self): 339 | result = task_d() 340 | self.assertEqual(result, 'd') 341 | 342 | def test_returns_subtask_result(self): 343 | result = task_c() 344 | self.assertEqual(result, 'a') 345 | 346 | def test_returns_subsubtask_result(self): 347 | result = task_e() 348 | self.assertEqual(result, 'a') 349 | 350 | 351 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 352 | class AsyncChainedTaskTestCase(SilentMixin, TransactionChannelTestCase): 353 | def test_all(self): 354 | task = task_h.delay() 355 | run_task(self.get_next_message('cq-tasks', require=True)) 356 | run_task(self.get_next_message('cq-tasks', require=True)) 357 | run_task(self.get_next_message('cq-tasks', require=True)) 358 | task.wait() 359 | self.assertEqual(task.status, task.STATUS_SUCCESS) 360 | self.assertEqual(task.result, 5) 361 | 362 | 363 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 364 | class AtomicTaskTestCase(SilentMixin, TransactionChannelTestCase): 365 | def test_non_atomic_success(self): 366 | uid = uuid.uuid4() 367 | task = task_l.delay(str(uid), error=False) 368 | run_task(self.get_next_message('cq-tasks', require=True)) 369 | task.wait() 370 | self.assertEqual(task.status, task.STATUS_SUCCESS) 371 | self.assertEqual(Task.objects.filter(id=str(uid)).exists(), True) 372 | 373 | def test_non_atomic_failure(self): 374 | uid = uuid.uuid4() 375 | task = task_l.delay(str(uid), error=True) 376 | run_task(self.get_next_message('cq-tasks', require=True)) 377 | task.wait() 378 | self.assertEqual(task.status, task.STATUS_FAILURE) 379 | self.assertEqual(Task.objects.filter(id=str(uid)).exists(), True) 380 | 381 | def test_atomic_success(self): 382 | uid = uuid.uuid4() 383 | task = task_m.delay(str(uid), error=False) 384 | run_task(self.get_next_message('cq-tasks', require=True)) 385 | task.wait() 386 | self.assertEqual(task.status, task.STATUS_SUCCESS) 387 | self.assertEqual(Task.objects.filter(id=str(uid)).exists(), True) 388 | 389 | def test_atomic_failure(self): 390 | uid = uuid.uuid4() 391 | task = task_m.delay(str(uid), error=True) 392 | run_task(self.get_next_message('cq-tasks', require=True)) 393 | task.wait() 394 | self.assertEqual(task.status, task.STATUS_FAILURE) 395 | self.assertEqual(Task.objects.filter(id=str(uid)).exists(), False) 396 | 397 | 398 | # class GetQueuedTasksTestCase(SilentMixin, TestCase): 399 | # def test_returns_empty(self): 400 | # task_ids = backend.get_queued_tasks() 401 | # self.assertEqual(task_ids, {}) 402 | 403 | # @skip('Need worker disabled.') 404 | # def test_queued(self): 405 | # chan = Channel('cq-tasks') 406 | # chan.send({'task_id': 'one'}) 407 | # chan.send({'task_id': 'two'}) 408 | # task_ids = get_queued_tasks() 409 | # self.assertEqual(task_ids, {'one', 'two'}) 410 | # cl = chan.channel_layer 411 | # while len(task_ids): 412 | # msg = cl.receive_many(['cq-tasks'], block=True) 413 | # task_ids.remove(msg[1]['task_id']) 414 | 415 | 416 | # class GetRunningTasksTestCase(SilentMixin, TestCase): 417 | # def test_empty_list(self): 418 | # task_ids = get_running_tasks() 419 | # self.assertEqual(task_ids, set()) 420 | 421 | # def test_running(self): 422 | # conn = Redis.from_url(settings.REDIS_URL) 423 | # conn.lpush('cq-current', 'one') 424 | # conn.lpush('cq-current', 'two') 425 | # task_ids = get_running_tasks() 426 | # self.assertEqual(task_ids, {'one', 'two'}) 427 | # task_ids = get_running_tasks() 428 | # self.assertEqual(task_ids, set()) 429 | 430 | 431 | class PublishCurrentTestCase(SilentMixin, TestCase): 432 | def test_publish(self): 433 | backend.clear_current() 434 | backend.set_current_task('hello') 435 | backend.publish_current(max_its=2, sleep_time=0.1) 436 | backend.set_current_task('world') 437 | backend.publish_current(max_its=3, sleep_time=0.1) 438 | task_ids = backend.get_running_tasks() 439 | self.assertEqual(task_ids, {'hello', 'world'}) 440 | task_ids = backend.get_running_tasks() 441 | self.assertEqual(task_ids, set()) 442 | 443 | 444 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 445 | class CreateRepeatingTaskTestCase(SilentMixin, TestCase): 446 | def test_create(self): 447 | rt = RepeatingTask.objects.create(func_name='cq.tests.test_base.task_a') 448 | next = croniter(rt.crontab, timezone.now()).get_next(datetime) 449 | self.assertEqual(rt.next_run, next) 450 | 451 | 452 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 453 | class RunRepeatingTaskTestCase(SilentMixin, TransactionChannelTestCase): 454 | def test_run(self): 455 | rt = RepeatingTask.objects.create(func_name='cq.tests.test_base.task_a') 456 | task = rt.submit() 457 | self.assertLess(rt.last_run, timezone.now()) 458 | self.assertGreater(rt.next_run, timezone.now()) 459 | run_task(self.get_next_message('cq-tasks', require=True)) 460 | task.wait() 461 | self.assertEqual(task.result, 'a') 462 | 463 | 464 | class ViewTestCase(SilentMixin, ChannelTestCaseMixin, APITransactionTestCase): 465 | def setUp(self): 466 | try: 467 | self.user = User.objects.create( 468 | username='a', email='a@a.org', password='a' 469 | ) 470 | except: 471 | self.user = User.objects.create( 472 | email='a@a.org', password='a' 473 | ) 474 | 475 | def test_create_and_get_task(self): 476 | 477 | # If the views aren't available, don't test. 478 | try: 479 | reverse('cqtask-list') 480 | except: 481 | return 482 | 483 | # Check task creation. 484 | data = { 485 | 'task': 'k', 486 | 'args': [False] 487 | } 488 | self.client.force_authenticate(self.user) 489 | response = self.client.post(reverse('cqtask-list'), data, format='json') 490 | self.assertEqual(response.status_code, 201) 491 | self.assertNotEqual(response.json().get('id', None), None) 492 | 493 | # Then retreival. 494 | id = response.json()['id'] 495 | response = self.client.get(reverse('cqtask-detail', kwargs={'pk': id}), data, format='json') 496 | self.assertEqual(response.status_code, 200) 497 | self.assertEqual(response.json()['status'], 'Q') 498 | run_task(self.get_next_message('cq-tasks', require=True)) 499 | response = self.client.get(reverse('cqtask-detail', kwargs={'pk': id}), data, format='json') 500 | self.assertEqual(response.status_code, 200) 501 | self.assertEqual(response.json()['status'], 'S') 502 | 503 | 504 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 505 | class SubmitExtraArgsTestCase(SilentMixin, TransactionChannelTestCase): 506 | def test_prefix_argument(self): 507 | task = task_a.delay_args(submit=False) 508 | task.submit('hello') 509 | run_task(self.get_next_message('cq-tasks', require=True)) 510 | task.wait() 511 | self.assertEqual(task.status, task.STATUS_SUCCESS) 512 | self.assertEqual(task.signature['args'][0], 'hello') 513 | self.assertEqual(task.result, ['hello']) 514 | 515 | def test_prefix_many_arguments(self): 516 | task = task_a.delay_args(submit=False) 517 | task.submit('hello', 'world') 518 | run_task(self.get_next_message('cq-tasks', require=True)) 519 | task.wait() 520 | self.assertEqual(task.status, task.STATUS_SUCCESS) 521 | self.assertEqual(task.signature['args'][0], 'hello') 522 | self.assertEqual(task.signature['args'][1], 'world') 523 | self.assertEqual(task.result, ['hello', 'world']) 524 | 525 | def test_prefix_arguments_with_existing(self): 526 | task = task_a.delay_args(args=('world',), submit=False) 527 | task.submit('hello') 528 | run_task(self.get_next_message('cq-tasks', require=True)) 529 | task.wait() 530 | self.assertEqual(task.status, task.STATUS_SUCCESS) 531 | self.assertEqual(task.signature['args'][0], 'hello') 532 | self.assertEqual(task.signature['args'][1], 'world') 533 | self.assertEqual(task.result, ['hello', 'world']) 534 | 535 | 536 | @override_settings(CQ_SERIAL=False, CQ_CHANNEL_LAYER='default') 537 | class DuplicateSubmitTestCase(SilentMixin, TransactionChannelTestCase): 538 | def test_fails_if_pending(self): 539 | task = task_a.delay(submit=False) 540 | with self.assertRaises(DuplicateSubmitError): 541 | task.submit() 542 | 543 | def test_skips_if_revoked(self): 544 | task = task_a.delay_args(submit=False) 545 | task.revoke() 546 | task.submit() 547 | with self.assertRaises(AssertionError): 548 | run_task(self.get_next_message('cq-tasks', require=True)) 549 | self.assertEqual(task.status, task.STATUS_REVOKED) 550 | -------------------------------------------------------------------------------- /cq/models.py: -------------------------------------------------------------------------------- 1 | import json 2 | import logging 3 | import re 4 | import time 5 | import uuid 6 | from datetime import datetime, timedelta 7 | from traceback import format_tb 8 | 9 | from asgiref.sync import async_to_sync 10 | from channels import DEFAULT_CHANNEL_LAYER 11 | from channels.exceptions import ChannelFull 12 | from channels.layers import get_channel_layer 13 | from croniter import croniter 14 | from django.conf import settings 15 | from django.contrib.postgres.fields import JSONField 16 | from django.core.cache import cache 17 | from django.core.exceptions import ValidationError 18 | from django.db import models, transaction 19 | from django.utils import timezone 20 | from django_redis import get_redis_connection 21 | 22 | from .managers import TaskManager 23 | from .task import TaskFunc, from_signature, to_func_name, to_signature 24 | from .utils import import_attribute, redis_connection 25 | 26 | logger = logging.getLogger('cq') 27 | 28 | 29 | class CQError(Exception): 30 | pass 31 | 32 | 33 | class DuplicateSubmitError(CQError): 34 | pass 35 | 36 | 37 | class Task(models.Model): 38 | """A persistent representation of a background task. 39 | """ 40 | STATUS_PENDING = 'P' 41 | STATUS_RETRY = 'Y' 42 | STATUS_QUEUED = 'Q' 43 | STATUS_RUNNING = 'R' 44 | STATUS_FAILURE = 'F' 45 | STATUS_SUCCESS = 'S' 46 | STATUS_WAITING = 'W' 47 | STATUS_INCOMPLETE = 'I' 48 | STATUS_LOST = 'L' 49 | STATUS_REVOKED = 'E' 50 | STATUS_CHOICES = ( 51 | (STATUS_PENDING, 'Pending'), 52 | (STATUS_RETRY, 'Retry'), 53 | (STATUS_QUEUED, 'Queued'), 54 | (STATUS_RUNNING, 'Running'), 55 | (STATUS_FAILURE, 'Failure'), 56 | (STATUS_SUCCESS, 'Success'), 57 | (STATUS_WAITING, 'Waiting'), 58 | (STATUS_INCOMPLETE, 'Incomplete'), 59 | (STATUS_LOST, 'Lost'), 60 | (STATUS_REVOKED, 'Revoked') 61 | ) 62 | STATUS_DONE = {STATUS_FAILURE, STATUS_SUCCESS, STATUS_INCOMPLETE, 63 | STATUS_LOST, STATUS_REVOKED} 64 | STATUS_ERROR = {STATUS_FAILURE, STATUS_LOST, STATUS_INCOMPLETE, 65 | STATUS_REVOKED} 66 | STATUS_ACTIVE = {STATUS_PENDING, STATUS_QUEUED, STATUS_RUNNING, 67 | STATUS_WAITING} 68 | 69 | AT_RISK_NONE = 'N' 70 | AT_RISK_QUEUED = 'Q' 71 | AT_RISK_RUNNING = 'R' 72 | AT_RISK_CHOICES = ( 73 | (AT_RISK_NONE, 'None'), 74 | (AT_RISK_QUEUED, 'Queued'), 75 | (AT_RISK_RUNNING, 'Running'), 76 | ) 77 | 78 | id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) 79 | status = models.CharField(max_length=1, choices=STATUS_CHOICES, default=STATUS_PENDING, db_index=True) 80 | signature = JSONField(default=dict, blank=True) 81 | details = JSONField(default=dict, blank=True) 82 | parent = models.ForeignKey('self', blank=True, null=True, related_name='subtasks', on_delete=models.CASCADE) 83 | previous = models.ForeignKey('self', related_name='next', blank=True, null=True, on_delete=models.SET_NULL) 84 | waiting_on = models.ForeignKey('self', blank=True, null=True, on_delete=models.SET_NULL) 85 | submitted = models.DateTimeField(auto_now_add=True) 86 | started = models.DateTimeField(null=True, blank=True) 87 | finished = models.DateTimeField(null=True, blank=True) 88 | result_ttl = models.PositiveIntegerField(default=1800, blank=True) 89 | result_expiry = models.DateTimeField(null=True, blank=True) 90 | at_risk = models.CharField(max_length=1, choices=AT_RISK_CHOICES, default=AT_RISK_NONE) 91 | retries = models.PositiveIntegerField(default=0) 92 | last_retry = models.DateTimeField(null=True, blank=True) 93 | force_chain = models.BooleanField(default=False) 94 | 95 | objects = TaskManager() 96 | 97 | class Meta: 98 | ordering = ('-submitted',) 99 | 100 | def __str__(self): 101 | return '{} - {}'.format(self.id, self.func_name) 102 | 103 | def retry(self): 104 | self.status = self.STATUS_PENDING 105 | self.started = None 106 | self.finished = None 107 | self.details = {} 108 | self.at_risk = self.AT_RISK_NONE 109 | self.retries += 1 110 | self.last_retry = timezone.now() 111 | self.save( 112 | update_fields=( 113 | 'status', 'started', 'finished', 114 | 'details', 'at_risk', 'retries', 115 | 'last_retry' 116 | ) 117 | ) 118 | self.submit() 119 | 120 | def submit(self, *pre_args): 121 | """ To be run from server. 122 | """ 123 | with cache.lock(str(self.id), timeout=2): 124 | 125 | # Need to reload just in case it's been modified elsewhere. 126 | self.refresh_from_db() 127 | 128 | # If we've been moved to revoke, don't run. If it's anything 129 | # other than pending, error. 130 | if self.status == self.STATUS_REVOKED: 131 | return 132 | elif self.status != self.STATUS_PENDING: 133 | raise DuplicateSubmitError( 134 | 'Task %s cannot be submitted multiple times.' % self.id 135 | ) 136 | self.status = self.STATUS_QUEUED 137 | 138 | # Prepend arguments. 139 | if len(pre_args) > 0: 140 | func, args, kwargs = from_signature(self.signature) 141 | args = pre_args + tuple(args) 142 | self.signature = to_signature(func, args, kwargs) 143 | 144 | # The database sometimes has not finished writing a commit 145 | # before the worker begins executing. In these cases we need 146 | # to wait for the commit. 147 | with transaction.atomic(): 148 | self.save(update_fields=('status', 'signature')) 149 | transaction.on_commit(lambda: self.send()) 150 | 151 | def send(self): 152 | if getattr(settings, 'CQ_BACKEND', '').lower() == 'redis': 153 | layer_name = 'cq' 154 | conn = get_redis_connection() 155 | conn.lpush(layer_name, str(self.id)) 156 | else: 157 | layer_name = getattr(settings, 'CQ_CHANNEL_LAYER', DEFAULT_CHANNEL_LAYER) 158 | layer = get_channel_layer(layer_name) 159 | logger.info('Sending CQ message on "{}" layer.'.format(layer_name)) 160 | try: 161 | async_to_sync(layer.send)('cq-task', {'type': 'run_task', 'task_id': str(self.id)}) 162 | logger.info('Message sent.') 163 | except ChannelFull: 164 | logger.error('CQ: Channel layer full.') 165 | self.status = self.STATUS_RETRY 166 | self.save(update_fields=('status',)) 167 | 168 | def wait(self, timeout=None): 169 | """Wait for task to finish. To be called from server. 170 | """ 171 | start = timezone.now() 172 | end = start 173 | if timeout is not None: 174 | start += timedelta(milliseconds=timeout) 175 | delta = timedelta(milliseconds=500) 176 | self.refresh_from_db() 177 | while (self.status not in self.STATUS_DONE and 178 | (timeout is None or start < end)): 179 | time.sleep(0.5) 180 | self.refresh_from_db() 181 | start += delta 182 | 183 | def pre_start(self): 184 | logger.debug('{}: task prestart'.format(self.id)) 185 | self.status = self.STATUS_RUNNING 186 | self.started = timezone.now() 187 | self.save(update_fields=('status', 'started')) 188 | 189 | # Ensure our logs are fresh. 190 | self._task_logs = [] 191 | with redis_connection() as con: 192 | con.delete(self._get_log_key()) 193 | logger.debug('{}: done'.format(self.id)) 194 | 195 | def start(self, result=None, pre_start=True): 196 | """To be run from workers. 197 | """ 198 | logger.debug('{}: task start'.format(self.id)) 199 | if pre_start: 200 | self.pre_start() 201 | func, args, kwargs = from_signature(self.signature) 202 | if result is not None: 203 | args = (result,) + tuple(args) 204 | task_func = TaskFunc.get_task(self.signature['func_name']) 205 | if task_func.atomic: 206 | with transaction.atomic(): 207 | return func(*args, task=self, **kwargs) 208 | else: 209 | return func(*args, task=self, **kwargs) 210 | logger.debug('{}: done'.format(self.id)) 211 | 212 | def revoke(self): 213 | with cache.lock(str(self.id), timeout=2): 214 | if self.status not in self.STATUS_DONE: 215 | self.status = self.STATUS_REVOKED 216 | self.save(update_fields=('status',)) 217 | for child in self.subtasks.all(): 218 | child.revoke() 219 | for next in self.next.all(): 220 | next.revoke() 221 | 222 | def subtask(self, func, args=(), kwargs={}, **kw): 223 | """Launch a subtask. 224 | 225 | Subtasks are run at the same time as the current task. The current 226 | task will not be considered complete until the subtask finishes. 227 | """ 228 | return delay(func, args, kwargs, parent=self, **kw) 229 | 230 | def chain(self, func, args=(), kwargs={}, **kw): 231 | """Chain a task. 232 | 233 | Chained tasks are run after completion of the current task, and are 234 | passed the result of the current task. 235 | """ 236 | return chain(func, args, kwargs, previous=self, **kw) 237 | 238 | def errorback(self, func, args=(), kwargs={}): 239 | self.details.setdefault('errbacks', []).append( 240 | to_signature(func, args, kwargs) 241 | ) 242 | self.save(update_fields=('details',)) 243 | 244 | def waiting(self, task=None, result=None): 245 | logger.info('{}: Waiting task: {}'.format(self.id, self.func_name)) 246 | self.status = self.STATUS_WAITING 247 | self.waiting_on = task 248 | if task is not None and task.parent != self: 249 | assert task.parent is None 250 | task.parent = self 251 | task.save(update_fields=('parent',)) 252 | if result is not None: 253 | logger.debug('Setting task result: {} = {}'.format( 254 | self.func_name, result 255 | )) 256 | self.details['result'] = result 257 | 258 | # Must publish all remaining local logs to REDIS. 259 | self._publish_logs() 260 | 261 | with transaction.atomic(): 262 | self.save(update_fields=('status', 'waiting_on', 'details')) 263 | transaction.on_commit(lambda: self.post_waiting()) 264 | 265 | def success(self, result=None): 266 | """To be run from workers. 267 | """ 268 | logger.info('{}: Task succeeded: {}'.format(self.id, self.func_name)) 269 | self.status = self.STATUS_SUCCESS 270 | if result is not None: 271 | logger.debug('Setting task result: {} = {}'.format( 272 | self.func_name, result 273 | )) 274 | self.details['result'] = result 275 | self.finished = timezone.now() 276 | self.result_expiry = self.finished + timedelta(seconds=self.result_ttl) 277 | self._store_logs() 278 | with transaction.atomic(): 279 | self.save(update_fields=('status', 'details', 'finished', 'result_expiry')) 280 | transaction.on_commit(lambda: self.post_success(self.result)) 281 | 282 | def post_success(self, result): 283 | if self.parent: 284 | self.parent.child_succeeded(self, result) 285 | self.launch_next() 286 | 287 | def post_waiting(self): 288 | self.launch_subtasks() 289 | 290 | def launch_subtasks(self): 291 | 292 | # Launch subtasks, but don't fire off any chained subtasks. Chained 293 | # tasks get registered as subtasks also in order to pass logs. 294 | for next in self.subtasks.all(): 295 | if next.previous is None: 296 | next.submit() 297 | 298 | def launch_next(self, force_chain=False): 299 | for next in self.next.all(): 300 | if not force_chain or next.force_chain: 301 | next.submit() 302 | 303 | def _store_logs(self): 304 | if self.parent: 305 | return 306 | key = self._get_log_key() 307 | with redis_connection() as con: 308 | try: 309 | logs = con.lrange(key, 0, -1) 310 | logs = [json.loads(l.decode()) for l in logs] 311 | except: 312 | logs = [] 313 | try: 314 | logs.extend(self._task_logs) 315 | except AttributeError: 316 | pass 317 | if 'logs' not in self.details: 318 | self.details['logs'] = [] 319 | self.details['logs'].extend(logs) 320 | with redis_connection() as con: 321 | con.delete(key) 322 | 323 | def child_succeeded(self, task, result): 324 | logger.info('{}: Task child succeeded: {}'.format(self.id, self.func_name)) 325 | if task == self.waiting_on and self.status not in self.STATUS_ERROR: 326 | logger.info('Setting task result: {} = {}'.format( 327 | self.func_name, result 328 | )) 329 | self.details['result'] = result 330 | self.save(update_fields=('details',)) 331 | if all([s.status == self.STATUS_SUCCESS for s in self.subtasks.all()]): 332 | logger.debug('All children succeeded: {}'.format(self.func_name)) 333 | self.success() 334 | 335 | def failure(self, err, retry=False): 336 | """To be run from workers. 337 | """ 338 | 339 | # Set the error details. 340 | self.details['error'] = str(err) 341 | self.details['exception'] = err.__class__.__name__ 342 | try: 343 | self.details['traceback'] = ''.join(format_tb(err.__traceback__)) 344 | except: 345 | pass 346 | 347 | # Set the status and start formatting the output message. 348 | if self.status == self.STATUS_WAITING or self.status == self.STATUS_INCOMPLETE: 349 | msg = 'Task incomplete: {}'.format(self.func_name) 350 | self.status = self.STATUS_INCOMPLETE 351 | else: 352 | msg = 'Task failed: {}'.format(self.func_name) 353 | self.status = self.STATUS_FAILURE 354 | 355 | # Finish the message. 356 | msg += '\nError: {}'.format(self.details['error']) 357 | if 'traceback' in self.details: 358 | msg += '\nTraceback:\n{}'.format(self.details['traceback']) 359 | logger.error('{}: {}'.format(self.id, msg)) 360 | 361 | if retry: 362 | self.status = self.STATUS_RETRY 363 | self.finished = timezone.now() 364 | self._store_logs() 365 | self.save(update_fields=('status', 'details', 'finished')) 366 | 367 | if not retry: 368 | if self.parent: 369 | self.parent.failure(err) 370 | for eb in self.details.get('errbacks', []): 371 | func, args, kwargs = from_signature(eb) 372 | func(*((self, err,) + tuple(args)), **kwargs) 373 | 374 | # Check if we want to force the subsequent chained 375 | # items to run. 376 | self.launch_next(force_chain=True) 377 | 378 | def log(self, msg, level=logging.INFO, origin=None, publish=True, 379 | limit=40): 380 | """Log to the task, and to the system logger. 381 | 382 | Will push the logged message to the topmost task. 383 | """ 384 | if self.parent: 385 | self.parent.log(msg, level, origin or self, publish=publish) 386 | else: 387 | logger.log(level, '{}: {}'.format(self.id, msg)) 388 | data = { 389 | 'message': msg, 390 | 'timestamp': str(timezone.now()) 391 | } 392 | if origin: 393 | data['origin'] = str(origin.id) 394 | try: 395 | self._task_logs.append(data) 396 | except AttributeError: 397 | self._task_logs = [data] 398 | 399 | # Don't try to set too much in the cache, it can cause 400 | # problems. Instead, cap it at the past `limit` logs. Also, use 401 | # `publish` to control when publishing happens. 402 | if publish: 403 | self._publish_logs(limit) 404 | 405 | def _publish_logs(self, limit=40): 406 | key = self._get_log_key() 407 | logs = [json.dumps(l) for l in self._task_logs[-limit:]] 408 | if logs: 409 | with redis_connection() as con: 410 | con.rpush(key, *logs) 411 | con.ltrim(key, -limit, -1) 412 | self._task_logs = [] 413 | 414 | @property 415 | def result(self): 416 | return self.details.get('result', None) 417 | 418 | @property 419 | def error(self): 420 | return self.details.get('error', None) 421 | 422 | @property 423 | def logs(self): 424 | logs = self.details.get('logs', None) 425 | if logs is None: 426 | with redis_connection() as con: 427 | key = self._get_log_key() 428 | try: 429 | logs = con.lrange(key, 0, -1) 430 | except: 431 | logs = [] 432 | logs = [json.loads(l.decode()) for l in logs] 433 | return logs 434 | 435 | @property 436 | def func_name(self): 437 | return self.signature.get('func_name', None) 438 | 439 | def format_logs(self): 440 | return '\n'.join([l['message'] for l in self.logs]) 441 | 442 | def _get_log_key(self): 443 | return 'cq:{}:logs'.format(self.id) 444 | 445 | 446 | def validate_cron(value): 447 | if value.strip() != value: 448 | raise ValidationError('Leading nor trailing spaces are allowed') 449 | columns = value.split() 450 | if columns != value.split(' '): 451 | raise ValidationError('Use only a single space as a column separator') 452 | if len(columns) != 5: 453 | raise ValidationError('Entry has to consist of exactly 5 columns') 454 | pattern = r'^(\*|\d+(-\d+)?(,\d+(-\d+)?)*)(/\d+)?$' 455 | p = re.compile(pattern) 456 | for i, c in enumerate(columns): 457 | if not p.match(c): 458 | raise ValidationError("Incorrect value {} in column {}".format( 459 | c, i + 1 460 | )) 461 | 462 | 463 | def validate_func_name(value): 464 | """Try to import a function before accepting it. 465 | """ 466 | try: 467 | import_attribute(value) 468 | except: 469 | raise ValidationError('Unable to import task.') 470 | 471 | 472 | class RepeatingTask(models.Model): 473 | """Basic repeating tasks. 474 | 475 | Uses CRON style strings to set repeating tasks. 476 | """ 477 | crontab = models.CharField(max_length=100, default='* * * * *', 478 | validators=[validate_cron], 479 | help_text='Minute Hour Day Month Weekday') 480 | func_name = models.CharField(max_length=256, validators=[validate_func_name]) 481 | args = JSONField(default=list, blank=True) 482 | kwargs = JSONField(default=dict, blank=True) 483 | result_ttl = models.PositiveIntegerField(default=1800, blank=True) 484 | last_run = models.DateTimeField(blank=True, null=True) 485 | next_run = models.DateTimeField(blank=True, null=True, db_index=True) 486 | coalesce = models.BooleanField(default=True) 487 | 488 | def __str__(self): 489 | arguments = ', '.join(map(lambda x: repr(x), self.args)) 490 | if self.kwargs: 491 | arguments += ', '.join(map( 492 | lambda row: '{}={}'.format(row[0], repr(row[1])), self.kwargs.items() 493 | )) 494 | task = '{}({})'.format(self.func_name, arguments) 495 | if self.last_run: 496 | task += ' ({})'.format(self.last_run) 497 | return task 498 | 499 | def submit(self): 500 | if self.coalesce and Task.objects.active(signature__func_name=self.func_name): 501 | logger.info('Coalescing task: {}'.format(self.func_name)) 502 | return None 503 | logger.info('Launching scheduled task: {}'.format(self)) 504 | with transaction.atomic(): 505 | task = delay(self.func_name, tuple(self.args), self.kwargs, 506 | submit=False, result_ttl=self.result_ttl) 507 | self.last_run = timezone.now() 508 | self.update_next_run() 509 | self.save(update_fields=('last_run', 'next_run')) 510 | task.submit() 511 | return task 512 | 513 | def update_next_run(self): 514 | self.next_run = croniter( 515 | self.crontab, timezone.localtime(timezone.now()) 516 | ).get_next(datetime) 517 | 518 | @classmethod 519 | def schedule(cls, crontab, func, args=(), kwargs={}): 520 | return schedule_task(cls, crontab, func, args, kwargs) 521 | 522 | 523 | def schedule_task(cls, crontab, func, args=(), kwargs={}, **_kwargs): 524 | """Create a repeating task. 525 | """ 526 | # This is mostly for creating scheduled tasks in migrations. The 527 | # signals don't run in migrations, so we need to explicitly set 528 | # the `next_run` value. 529 | next = croniter(crontab, timezone.localtime(timezone.now())).get_next(datetime) 530 | return cls.objects.create( 531 | crontab=crontab, 532 | func_name=to_func_name(func), 533 | args=args, 534 | kwargs=kwargs, 535 | next_run=next, 536 | **_kwargs 537 | ) 538 | 539 | 540 | def chain(func, args, kwargs, parent=None, previous=None, submit=True, 541 | **kw): 542 | """Run a task after an existing task. 543 | 544 | The result is passed as the first argument to the chained task. 545 | If no parent is specified, automatically use the parent of the 546 | predecessor. Note that I'm not sure this is the correct behavior, 547 | but is useful for making sure logs to where they should. 548 | """ 549 | sig = to_signature(func, args, kwargs) 550 | if parent is None and previous: 551 | parent = previous.parent 552 | task = Task.objects.create(signature=sig, parent=parent, previous=previous, 553 | **kw) 554 | 555 | # Need to check immediately if the parent task has completed and 556 | # launch the subtask if so. 557 | if parent is not None and submit: 558 | with cache.lock(str(parent.id), timeout=2): 559 | parent.refresh_from_db() 560 | if parent.status == Task.STATUS_SUCCESS: 561 | task.submit() 562 | 563 | # If we have no parent, and we want to submit then do so now. This 564 | # happens for a straight-up delay. 565 | elif parent is None and submit: 566 | task.submit() 567 | 568 | return task 569 | 570 | 571 | def delay(func, args, kwargs, parent=None, submit=True, **kw): 572 | task = chain(func, args, kwargs, parent, submit=submit, **kw) 573 | # if submit: 574 | # task.submit() 575 | return task 576 | --------------------------------------------------------------------------------