fix: better handle errors in celery result backend

- use celery's own "retry on result-backend error" logic
    - remove `die_on_unhandled` hack
    - new env var: `CELERY_RESULT_BACKEND_MAX_RETRIES` (default 173)
    - retry on error caused by `get_or_create` in overlapping
      transactions (with less-strict transaction isolation levels)
- port fixes from celery's `BaseKeyValueStoreBackend`:
    - avoid clobbering successes with "worker lost" or other errors
    - avoid error trying to get non-existent task results
- use celery's `BaseBackend` instead of `BaseDictBackend`
  (equivalent for back-compat; let's use the better name)

Author: aaxelb

project/settings.py

@@ -336,6 +336,8 @@ def split(string, delim):
 }
 
 CELERY_RESULT_BACKEND = 'share.celery:CeleryDatabaseBackend'
+CELERY_RESULT_BACKEND_ALWAYS_RETRY = True
+CELERY_RESULT_BACKEND_MAX_RETRIES = int(os.environ.get('CELERY_RESULT_BACKEND_MAX_RETRIES', 173))
 CELERY_RESULT_EXPIRES = int(os.environ.get(
     'CELERY_RESULT_EXPIRES',
     60 * 60 * 24 * 3,  # 3 days

share/celery.py

@@ -1,19 +1,17 @@
 import datetime
-import functools
 import logging
 
 from celery import states
 from celery.app.task import Context
-from celery.backends.base import BaseDictBackend
+from celery.backends.base import BaseBackend
 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.db.utils import IntegrityError
 from django.utils import timezone
 
-import sentry_sdk
-
 from share.models import CeleryTaskResult
 from share.models.sql import GroupBy
 from trove.util.django import pk_chunked
@@ -22,25 +20,8 @@
 logger = logging.getLogger(__name__)
 
 
-def die_on_unhandled(func):
-    @functools.wraps(func)
-    def wrapped(*args, **kwargs):
-        err = None
-        try:
-            return func(*args, **kwargs)
-        except Exception as e:
-            err = e
-            try:
-                logger.exception('Celery internal method %s failed', func)
-                sentry_sdk.capture_exception()
-            finally:
-                if err:
-                    raise SystemExit(57)  # Something a bit less generic than 1 or -1
-    return wrapped
-
-
 # Based on https://github.com/celery/django-celery-results/commit/f88c677d66ba1eaf1b7cb1f3b8c910012990984f
-class CeleryDatabaseBackend(BaseDictBackend):
+class CeleryDatabaseBackend(BaseBackend):
     """
 
     Implemented from scratch rather than subclassed due to:
@@ -53,8 +34,28 @@ class CeleryDatabaseBackend(BaseDictBackend):
     """
     TaskModel = CeleryTaskResult
 
-    @die_on_unhandled
+    # override BaseBackend
+    def exception_safe_to_retry(self, exc):
+        # retry error from multiple simultaneous `get_or_create` calls --
+        # re-raises an IntegrityError while handling TaskModel.DoesNotExist
+        # https://github.com/django/django/blob/9e7cc2b628fe8fd3895986af9b7fc9525034c1b0/django/db/models/query.py#L959
+        return (
+            isinstance(exc, IntegrityError)
+            and isinstance(exc.__context__, self.TaskModel.DoesNotExist)
+        )
+
+    # implement for BaseBackend
     def _store_result(self, task_id, result, status, traceback=None, request=None, **kwargs):
+        _already_successful = (
+            self.TaskModel.objects
+            .filter(task_id=task_id, status=states.SUCCESS)
+            .exists()
+        )
+        if _already_successful:
+            # avoid clobbering prior successful result, which could be caused by network partition or lost worker, ostensibly:
+            # https://github.com/celery/celery/blob/92514ac88afc4ccdff31f3a1018b04499607ca1e/celery/backends/base.py#L967-L972
+            return
+
         fields = {
             'result': result,
             'traceback': traceback,
@@ -88,20 +89,21 @@ def _store_result(self, task_id, result, status, traceback=None, request=None, *
                     setattr(obj, key, value)
             obj.save()
 
-        return obj
-
-    @die_on_unhandled
+    # override BaseBackend
     def cleanup(self, expires=None):
         TaskResultCleaner(
             success_ttl=(expires or self.expires),
             nonsuccess_ttl=settings.FAILED_CELERY_RESULT_EXPIRES,
         ).clean()
 
-    @die_on_unhandled
+    # implement for BaseBackend
     def _get_task_meta_for(self, task_id):
-        return self.TaskModel.objects.get(task_id=task_id).as_dict()
+        try:
+            return self.TaskModel.objects.get(task_id=task_id).as_dict()
+        except self.TaskModel.DoesNotExist:
+            return {'status': states.PENDING, 'result': None}
 
-    @die_on_unhandled
+    # implement for BaseBackend
     def _forget(self, task_id):
         try:
             self.TaskModel.objects.get(task_id=task_id).delete()