Split celery timeout into two separate settings (#262)

`HEALTHCHECK_CELERY_TIMEOUT` is replaced by `HEALTHCHECK_CELERY_RESULT_TIMEOUT` and `HEALTHCHECK_CELERY_QUEUE_TIMEOUT`. Where the former is the timeout for the results to be published and the letter is the queue timeout. `HEALTHCHECK_CELERY_QUEUE_TIMEOUT` is deprecated and will emit a deprecation warning. Co-authored-by: Nick Werle <[email protected]>
revsys · Nov 4, 2020 · f5244fb · f5244fb
1 parent d71f55d
commit f5244fb
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 2 deletions.
diff --git a/docs/settings.rst b/docs/settings.rst
@@ -69,3 +69,24 @@ exceeds 90% or available memory drops below 100 MB.
 
    Specify the desired memory utilization threshold, in megabytes. When available
    memory falls below the specified value, a warning will be reported.
+
+Celery Health Check
+----------------------
+Using `django.settings` you may exert more fine-grained control over the behavior of the celery health check
+
+.. list-table:: Additional Settings
+   :widths: 25 10 10 55
+   :header-rows: 1
+
+   * - Name
+     - Type
+     - Default
+     - Description
+   * - `HEALTHCHECK_CELERY_QUEUE_TIMEOUT`
+     - Number
+     - 3
+     - Specifies the maximum amount of time a task may spend in the queue before being automatically revoked with a `TaskRevokedError`.
+   * - `HEALTHCHECK_CELERY_RESULT_TIMEOUT`
+     - Number
+     - 3
+     - Specifies the maximum total time for a task to complete and return a result, including queue time.
diff --git a/health_check/contrib/celery/apps.py b/health_check/contrib/celery/apps.py
@@ -1,5 +1,8 @@
 from celery import current_app
 from django.apps import AppConfig
+from django.conf import settings
+import warnings
+
 
 from health_check.plugins import plugin_dir
 
@@ -9,6 +12,10 @@ class HealthCheckConfig(AppConfig):
 
     def ready(self):
         from .backends import CeleryHealthCheck
+        if hasattr(settings, "HEALTHCHECK_CELERY_TIMEOUT"):
+            warnings.warn("HEALTHCHECK_CELERY_TIMEOUT is depricated and may be removed in the "
+                        "future. Please use HEALTHCHECK_CELERY_RESULT_TIMEOUT and "
+                        "HEALTHCHECK_CELERY_QUEUE_TIMEOUT instead.", DeprecationWarning)
 
         for queue in current_app.amqp.queues:
             celery_class_name = 'CeleryHealthCheck' + queue.title()

diff --git a/health_check/contrib/celery/backends.py b/health_check/contrib/celery/backends.py
@@ -6,24 +6,32 @@
 )
 
 from .tasks import add
+from celery.exceptions import TaskRevokedError, TimeoutError
 
 
 class CeleryHealthCheck(BaseHealthCheckBackend):
     def check_status(self):
         timeout = getattr(settings, 'HEALTHCHECK_CELERY_TIMEOUT', 3)
+        result_timeout = getattr(settings, 'HEALTHCHECK_CELERY_RESULT_TIMEOUT', timeout)
+        queue_timeout = getattr(settings, 'HEALTHCHECK_CELERY_QUEUE_TIMEOUT', timeout)
 
         try:
             result = add.apply_async(
                 args=[4, 4],
-                expires=timeout,
+                expires=queue_timeout,
                 queue=self.queue
             )
-            result.get(timeout=timeout)
+            result.get(timeout=result_timeout)
             if result.result != 8:
                 self.add_error(ServiceReturnedUnexpectedResult("Celery returned wrong result"))
         except IOError as e:
             self.add_error(ServiceUnavailable("IOError"), e)
         except NotImplementedError as e:
             self.add_error(ServiceUnavailable("NotImplementedError: Make sure CELERY_RESULT_BACKEND is set"), e)
+        except TaskRevokedError as e:
+            self.add_error(ServiceUnavailable("TaskRevokedError: The task was revoked, likely because it spent "
+                                              "too long in the queue"), e)
+        except TimeoutError as e:
+            self.add_error(ServiceUnavailable("TimeoutError: The task took too long to return a result"), e)
         except BaseException as e:
             self.add_error(ServiceUnavailable("Unknown error"), e)