feat: Support task backoff

Author: khvn26

settings/dev.py

@@ -58,6 +58,7 @@
 ENABLE_CLEAN_UP_OLD_TASKS = True
 ENABLE_TASK_PROCESSOR_HEALTH_CHECK = True
 RECURRING_TASK_RUN_RETENTION_DAYS = 15
+TASK_BACKOFF_DEFAULT_SECONDS = 5
 TASK_DELETE_BATCH_SIZE = 2000
 TASK_DELETE_INCLUDE_FAILED_TASKS = False
 TASK_DELETE_RETENTION_DAYS = 15

src/task_processor/exceptions.py

@@ -1,3 +1,6 @@
+from datetime import datetime
+
+
 class TaskProcessingError(Exception):
     pass
 
@@ -6,5 +9,21 @@ class InvalidArgumentsError(TaskProcessingError):
     pass
 
 
+class TaskBackoffError(TaskProcessingError):
+    """
+    Raise this exception inside a task to indicate that it should be retried after a delay.
+    This is typically used when a task fails due to a temporary issue, such as
+    a network error or a service being unavailable.
+    """
+
+    def __init__(
+        self,
+        delay_until: datetime | None = None,
+    ) -> None:
+        super().__init__()
+        delay_until = delay_until
+        self.delay_until = delay_until
+
+
 class TaskQueueFullError(Exception):
     pass

src/task_processor/models.py

@@ -7,9 +7,9 @@
 from django.db import models
 from django.utils import timezone
 
-from task_processor.exceptions import TaskProcessingError, TaskQueueFullError
+from task_processor.exceptions import TaskQueueFullError
 from task_processor.managers import RecurringTaskManager, TaskManager
-from task_processor.task_registry import registered_tasks
+from task_processor.task_registry import get_task, registered_tasks
 from task_processor.types import TaskCallable
 
 _django_json_encoder_default = DjangoJSONEncoder().default
@@ -75,15 +75,8 @@ def run(self) -> None:
 
     @property
     def callable(self) -> TaskCallable[typing.Any]:
-        try:
-            task = registered_tasks[self.task_identifier]
-            return task.task_function
-        except KeyError as e:
-            raise TaskProcessingError(
-                "No task registered with identifier '%s'. Ensure your task is "
-                "decorated with @register_task_handler.",
-                self.task_identifier,
-            ) from e
+        task = get_task(self.task_identifier)
+        return task.task_handler.unwrapped
 
 
 class Task(AbstractBaseTask):

src/task_processor/processor.py

@@ -5,10 +5,12 @@
 from contextlib import ExitStack
 from datetime import timedelta
 
+from dateutil.relativedelta import relativedelta
 from django.conf import settings
 from django.utils import timezone
 
 from task_processor import metrics
+from task_processor.exceptions import TaskBackoffError
 from task_processor.managers import RecurringTaskManager, TaskManager
 from task_processor.models import (
     AbstractBaseTask,
@@ -120,6 +122,7 @@ def _run_task(
     ctx.enter_context(timer)
 
     task_identifier = task.task_identifier
+    registered_task = get_task(task_identifier)
 
     logger.debug(
         f"Running task {task_identifier} id={task.pk} args={task.args} kwargs={task.kwargs}"
@@ -157,9 +160,24 @@ def _run_task(
             exc_info=True,
         )
 
+        if isinstance(e, TaskBackoffError):
+            delay_until = e.delay_until or timezone.now() + relativedelta(
+                seconds=settings.TASK_BACKOFF_DEFAULT_SECONDS,
+            )
+            registered_task.task_handler.delay(
+                delay_until=delay_until,
+                args=task.args,
+                kwargs=task.kwargs,
+            )
+            logger.info(
+                "Backoff requested. Task '%s' set to retry at %s",
+                task_identifier,
+                delay_until,
+            )
+
     labels = {
         "task_identifier": task_identifier,
-        "task_type": get_task(task_identifier).task_type.value.lower(),
+        "task_type": registered_task.task_type.value.lower(),
         "result": result.lower(),
     }
 

src/task_processor/task_registry.py

@@ -3,8 +3,12 @@
 import typing
 from dataclasses import dataclass
 
+from task_processor.exceptions import TaskProcessingError
 from task_processor.types import TaskCallable
 
+if typing.TYPE_CHECKING:
+    from task_processor.decorators import TaskHandler  # noqa: F401
+
 logger = logging.getLogger(__name__)
 
 
@@ -16,7 +20,7 @@ class TaskType(enum.Enum):
 @dataclass
 class RegisteredTask:
     task_identifier: str
-    task_function: TaskCallable[typing.Any]
+    task_handler: "TaskHandler"
     task_type: TaskType = TaskType.STANDARD
     task_kwargs: dict[str, typing.Any] | None = None
 
@@ -43,18 +47,25 @@ def initialise() -> None:
 def get_task(task_identifier: str) -> RegisteredTask:
     global registered_tasks
 
-    return registered_tasks[task_identifier]
+    try:
+        return registered_tasks[task_identifier]
+    except KeyError:
+        raise TaskProcessingError(
+            "No task registered with identifier '%s'. Ensure your task is "
+            "decorated with @register_task_handler.",
+            task_identifier,
+        )
 
 
 def register_task(
     task_identifier: str,
-    callable_: TaskCallable[typing.Any],
+    task_handler: "TaskHandler",
 ) -> None:
     global registered_tasks
 
     registered_task = RegisteredTask(
         task_identifier=task_identifier,
-        task_function=callable_,
+        task_handler=task_handler,
     )
     registered_tasks[task_identifier] = registered_task