feat: option to fail-on-timeout

fixes #251

django_q/cluster.py

@@ -400,6 +400,29 @@ def monitor(result_queue: Queue, broker: Broker = None):
     logger.info(_(f"{name} stopped monitoring results"))
 
 
+def _check_task_timed_out(key, task: dict):
+    result = None
+    broker = get_broker()
+    cache = broker.cache
+    working_set = cache.get(key) or set()
+    if task["id"] in working_set:
+        # the previous worker has timedout and wasn't given chance to clear
+        raise Exception(f"Task Timed-out: {task}.")
+    else:
+        working_set.add(task['id'])
+        cache.set(key, working_set, timeout=Conf.RETRY * 3)
+    return result
+
+
+def _clear_task_timeout_cache(key, task):
+    broker = get_broker()
+    cache = broker.cache
+    working_set = cache.get(key) or set()
+    if task["id"] in working_set:
+        working_set.remove(task['id'])
+        cache.set(key, working_set, timeout=Conf.RETRY * 3)
+
+
 def worker(
     task_queue: Queue, result_queue: Queue, timer: Value, timeout: int = Conf.TIMEOUT
 ):
@@ -415,6 +438,8 @@ def worker(
     task_count = 0
     if timeout is None:
         timeout = -1
+
+    working_tasks_key = "DJANGO-Q-WORKING-TASKS"
     # Start reading the task queue
     for task in iter(task_queue.get, "STOP"):
         result = None
@@ -434,7 +459,11 @@ def worker(
         pre_execute.send(sender="django_q", func=f, task=task)
         # execute the payload
         timer.value = timer_value  # Busy
+
+
         try:
+            if Conf.FAIL_ON_TIMEOUT:
+                _check_task_timed_out(working_tasks_key, task)
             res = f(*task["args"], **task["kwargs"])
             result = (res, True)
         except Exception as e:
@@ -443,6 +472,9 @@ def worker(
                 error_reporter.report()
             if task.get("sync", False):
                 raise
+        if Conf.FAIL_ON_TIMEOUT:
+            _clear_task_timeout_cache(working_tasks_key, task)
+
         with timer.get_lock():
             # Process result
             task["result"] = result[0]

django_q/conf.py

@@ -123,6 +123,9 @@ class Conf:
     # Number of seconds to wait for a worker to finish.
     TIMEOUT = conf.get("timeout", None)
 
+    # Whether to fail the task when it times-out.
+    FAIL_ON_TIMEOUT = conf.get("fail_on_timeout", False)
+
     # Whether to acknowledge unsuccessful tasks.
     # This causes failed tasks to be considered delivered, thereby removing them from
     # the task queue. Defaults to False.