delete failed task results eventually

Author: aaxelb

project/settings.py

@@ -341,8 +341,17 @@ def split(string, delim):
     },
 }
 
-CELERY_RESULT_EXPIRES = 60 * 60 * 24 * 3  # 4 days
 CELERY_RESULT_BACKEND = 'share.celery:CeleryDatabaseBackend'
+CELERY_RESULT_EXPIRES = int(os.environ.get(
+    'CELERY_RESULT_EXPIRES',
+    60 * 60 * 24 * 3,  # 3 days
+))
+# only successful tasks get the default expiration (above)
+# -- failed tasks kept longer (see `share.celery`)
+FAILED_CELERY_RESULT_EXPIRES = int(os.environ.get(
+    'FAILED_CELERY_RESULT_EXPIRES',
+    60 * 60 * 24 * 11,  # 11 days
+))
 
 # Don't reject tasks that were present on a worker when it was killed
 CELERY_TASK_REJECT_ON_WORKER_LOST = False

share/celery.py

@@ -1,13 +1,15 @@
+import datetime
 import functools
 import logging
 
-
 from celery import states
 from celery.app.task import Context
 from celery.backends.base import BaseDictBackend
 from celery.utils.time import maybe_timedelta
 
+from django.conf import settings
 from django.db import transaction
+from django.db.models import Q
 from django.utils import timezone
 
 import sentry_sdk
@@ -90,7 +92,10 @@ def _store_result(self, task_id, result, status, traceback=None, request=None, *
 
     @die_on_unhandled
     def cleanup(self, expires=None):
-        TaskResultCleaner(expires or self.expires).clean()
+        TaskResultCleaner(
+            success_ttl=(expires or self.expires),
+            nonsuccess_ttl=settings.FAILED_CELERY_RESULT_EXPIRES,
+        ).clean()
 
     @die_on_unhandled
     def _get_task_meta_for(self, task_id):
@@ -111,20 +116,19 @@ class TaskResultCleaner:
 
     TaskModel = CeleryTaskResult
 
-    TASK_TTLS = {
-    }
-
-    NO_ARCHIVE = {
-    }
-
-    def __init__(self, expires, bucket=None, delete=True, chunk_size=5000):
-        self.bucket = bucket
+    def __init__(self, success_ttl, nonsuccess_ttl=None, delete=True, chunk_size=5000):
         self.chunk_size = chunk_size
         self.delete = delete
-        self.expires = expires
+        self.success_ttl = success_ttl
+        self.nonsuccess_ttl = nonsuccess_ttl or success_ttl
 
-    def get_ttl(self, task_name):
-        return timezone.now() - maybe_timedelta(self.TASK_TTLS.get(task_name, self.expires))
+    @property
+    def success_cutoff(self) -> datetime.datetime:
+        return timezone.now() - maybe_timedelta(self.success_ttl)
+
+    @property
+    def nonsuccess_cutoff(self) -> datetime.datetime:
+        return timezone.now() - maybe_timedelta(self.nonsuccess_ttl)
 
     def get_task_names(self):
         qs = self.TaskModel.objects.values('task_name').annotate(name=GroupBy('task_name'))
@@ -137,12 +141,15 @@ def get_task_names(self):
 
     def clean(self):
         for name in self.get_task_names():
-            logger.debug('Looking for succeeded %s tasks modified before %s', name, self.get_ttl(name))
-
-            queryset = self.TaskModel.objects.filter(
-                task_name=name,
-                status=states.SUCCESS,
-                date_modified__lt=self.get_ttl(name)
+            success_q = Q(status=states.SUCCESS, date_modified__lt=self.success_cutoff)
+            nonsuccess_q = (
+                ~Q(status=states.SUCCESS)
+                & Q(date_modified__lt=self.nonsuccess_cutoff)
+            )
+            queryset = (
+                self.TaskModel.objects
+                .filter(task_name=name)
+                .filter(success_q | nonsuccess_q)
             )
 
             if not queryset.exists():

tests/share/test_celery.py

@@ -1,49 +1,72 @@
-import pytest
-import datetime
-
+import contextlib
+from datetime import timedelta
 from unittest import mock
 
+import pytest
 from django.utils import timezone
 
 from share.celery import TaskResultCleaner, CeleryTaskResult
-
 from tests import factories
 
 
-@pytest.mark.usefixtures('nested_django_db')
-class TestResultArchiver:
+@contextlib.contextmanager
+def long_now(new_now=None):
+    _now = new_now or timezone.now()
+    with mock.patch.object(timezone, 'now', return_value=_now):
+        yield _now
 
-    @pytest.fixture(scope='class', autouse=True)
-    def task_result_data(self, class_scoped_django_db):
-        return factories.CeleryTaskResultFactory.create_batch(100)
+
+@pytest.mark.django_db
+class TestResultCleaner:
 
     def test_delete_false(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520), delete=False)
+        factories.CeleryTaskResultFactory.create_batch(10)
+        trc = TaskResultCleaner(timedelta(weeks=520), delete=False)
         assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 0
-        assert CeleryTaskResult.objects.count() != 0
+        assert CeleryTaskResult.objects.count() == 10
 
     def test_delete_queryset(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 100
+        factories.CeleryTaskResultFactory.create_batch(10)
+        trc = TaskResultCleaner(timedelta(weeks=520))
+        assert trc.delete_queryset(CeleryTaskResult.objects.all()) == 10
         assert CeleryTaskResult.objects.count() == 0
 
-    def test_get_ttl_default(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        assert ((timezone.now() - datetime.timedelta(weeks=520)) - trc.get_ttl('non-existant-task')) < datetime.timedelta(seconds=2)
-
-    def test_get_ttl(self):
-        trc = TaskResultCleaner(datetime.timedelta(weeks=520))
-        trc.TASK_TTLS['existant-task'] = datetime.timedelta(days=1)
-        assert ((timezone.now() - datetime.timedelta(days=1)) - trc.get_ttl('existant-task')) < datetime.timedelta(seconds=2)
-
-    def test_clean(self):
-        trc = TaskResultCleaner(0, bucket=mock.Mock())
-        factories.CeleryTaskResultFactory.create_batch(100, status='SUCCESS')
-        trc.clean()
-        assert CeleryTaskResult.objects.count() <= 100  # There's an autouse fixture that makes 100
-
-    def test_clean_chunksize(self):
-        trc = TaskResultCleaner(0, bucket=mock.Mock(), chunk_size=1)
-        factories.CeleryTaskResultFactory.create_batch(100, status='SUCCESS')
-        trc.clean()
-        assert CeleryTaskResult.objects.count() <= 100  # There's an autouse fixture that makes 100
+    def test_success_cutoff(self, settings):
+        with long_now() as _now:
+            trc = TaskResultCleaner(timedelta(days=3).total_seconds())
+            _expected = _now - timedelta(days=3)
+            assert trc.success_cutoff == _expected
+
+    def test_nonsuccess_cutoff(self, settings):
+        with long_now() as _now:
+            trc = TaskResultCleaner(
+                success_ttl=timedelta(days=3),
+                nonsuccess_ttl=timedelta(days=5),
+            )
+            assert trc.success_cutoff == _now - timedelta(days=3)
+            assert trc.nonsuccess_cutoff == _now - timedelta(days=5)
+
+    @pytest.mark.parametrize('batch_size', [1, 1111])
+    def test_clean(self, batch_size):
+        with long_now() as _now:
+            with long_now(_now - timedelta(days=7)):
+                # all should be deleted:
+                factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+                factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            with long_now(_now - timedelta(days=4)):
+                # successes should be deleted:
+                factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+                factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            # none should be deleted:
+            factories.CeleryTaskResultFactory.create_batch(10, status='SUCCESS')
+            factories.CeleryTaskResultFactory.create_batch(7, status='FAILED')
+            # end setup
+            assert CeleryTaskResult.objects.count() == 51
+            trc = TaskResultCleaner(
+                success_ttl=timedelta(days=3),
+                nonsuccess_ttl=timedelta(days=5),
+                chunk_size=batch_size,
+            )
+            trc.clean()
+            assert CeleryTaskResult.objects.filter(status='SUCCESS').count() == 10
+            assert CeleryTaskResult.objects.exclude(status='SUCCESS').count() == 14