Implement timeout

django_tasks/backends/database/backend.py

@@ -46,6 +46,7 @@ def _task_to_db_task(
             queue_name=task.queue_name,
             run_after=task.run_after,
             backend_name=self.alias,
+            timeout=task.timeout,
         )
 
     def enqueue(

django_tasks/backends/database/migrations/0015_dbtaskresult_timeout.py

@@ -0,0 +1,17 @@
+# Generated by Django 5.1.6 on 2025-02-16 16:53
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("django_tasks_database", "0014_remove_dbtaskresult_exception_data"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="dbtaskresult",
+            name="timeout",
+            field=models.PositiveBigIntegerField(default=600, verbose_name="timeout"),
+        ),
+    ]

django_tasks/backends/database/models.py

@@ -15,6 +15,7 @@
 from django_tasks.task import (
     DEFAULT_PRIORITY,
     DEFAULT_QUEUE_NAME,
+    DEFAULT_TIMEOUT,
     MAX_PRIORITY,
     MIN_PRIORITY,
     ResultStatus,
@@ -99,6 +100,7 @@ class DBTaskResult(GenericBase[P, T], models.Model):
     backend_name = models.TextField(_("backend name"))
 
     run_after = models.DateTimeField(_("run after"), null=True)
+    timeout = models.PositiveBigIntegerField(_("timeout"), default=DEFAULT_TIMEOUT)
 
     return_value = models.JSONField(_("return value"), default=None, null=True)
 
@@ -141,6 +143,7 @@ def task(self) -> Task[P, T]:
             queue_name=self.queue_name,
             run_after=self.run_after,
             backend=self.backend_name,
+            timeout=self.timeout,
         )
 
     @property

django_tasks/exceptions.py

@@ -13,3 +13,9 @@ class InvalidTaskBackendError(ImproperlyConfigured):
 
 class ResultDoesNotExist(ObjectDoesNotExist):
     pass
+
+
+class TimeoutException(BaseException):
+    """
+    Something timed out.
+    """

django_tasks/task.py

@@ -32,6 +32,7 @@
 MIN_PRIORITY = -100
 MAX_PRIORITY = 100
 DEFAULT_PRIORITY = 0
+DEFAULT_TIMEOUT = 600  # 10 minutes
 
 TASK_REFRESH_ATTRS = {
     "_exception_class",
@@ -78,6 +79,9 @@ class Task(Generic[P, T]):
     immediately, or whatever the backend decides
     """
 
+    timeout: int = DEFAULT_TIMEOUT
+    """The maximum duration the task can take to execute before being aborted"""
+
     def __post_init__(self) -> None:
         self.get_backend().validate_task(self)
 
@@ -95,6 +99,7 @@ def using(
         queue_name: Optional[str] = None,
         run_after: Optional[Union[datetime, timedelta]] = None,
         backend: Optional[str] = None,
+        timeout: Optional[int] = None,
     ) -> Self:
         """
         Create a new task with modified defaults
@@ -110,6 +115,8 @@ def using(
             changes["run_after"] = run_after
         if backend is not None:
             changes["backend"] = backend
+        if timeout is not None:
+            changes["timeout"] = timeout
 
         return replace(self, **changes)
 
@@ -188,6 +195,7 @@ def task(
     queue_name: str = DEFAULT_QUEUE_NAME,
     backend: str = DEFAULT_TASK_BACKEND_ALIAS,
     enqueue_on_commit: Optional[bool] = None,
+    timeout: int = DEFAULT_TIMEOUT,
 ) -> Callable[[Callable[P, T]], Task[P, T]]: ...
 
 
@@ -199,6 +207,7 @@ def task(
     queue_name: str = DEFAULT_QUEUE_NAME,
     backend: str = DEFAULT_TASK_BACKEND_ALIAS,
     enqueue_on_commit: Optional[bool] = None,
+    timeout: int = DEFAULT_TIMEOUT,
 ) -> Union[Task[P, T], Callable[[Callable[P, T]], Task[P, T]]]:
     """
     A decorator used to create a task.
@@ -212,6 +221,7 @@ def wrapper(f: Callable[P, T]) -> Task[P, T]:
             queue_name=queue_name,
             backend=backend,
             enqueue_on_commit=enqueue_on_commit,
+            timeout=timeout,
         )
 
     if function:
@@ -223,7 +233,7 @@ def wrapper(f: Callable[P, T]) -> Task[P, T]:
 @dataclass(frozen=True)
 class TaskResult(Generic[T]):
     task: Task
-    """The task for which this is a result"""
+    """The task for which this is a result, as it was run"""
 
     id: str
     """A unique identifier for the task result"""

django_tasks/utils.py

@@ -1,14 +1,19 @@
+import ctypes
 import inspect
 import json
 import random
 import time
+from contextlib import contextmanager
 from functools import wraps
+from threading import Timer, current_thread
 from traceback import format_exception
-from typing import Any, Callable, TypeVar
+from typing import Any, Callable, Iterator, TypeVar
 
 from django.utils.crypto import RANDOM_STRING_CHARS
 from typing_extensions import ParamSpec
 
+from .exceptions import TimeoutException
+
 T = TypeVar("T")
 P = ParamSpec("P")
 
@@ -74,3 +79,40 @@ def get_random_id() -> str:
     it's not cryptographically secure.
     """
     return "".join(random.choices(RANDOM_STRING_CHARS, k=32))
+
+
+def _do_timeout(tid: int) -> None:
+    """
+    Raise `TimeoutException` in the given thread.
+
+    Here be dragons.
+    """
+    # Since we're in Python, no GIL lock is needed
+    ret = ctypes.pythonapi.PyThreadState_SetAsyncExc(
+        ctypes.c_ulong(tid), ctypes.py_object(TimeoutException)
+    )
+
+    if ret == 0:
+        raise RuntimeError("Timeout failed - thread not found")
+    elif ret != 1:
+        ret = ctypes.pythonapi.PyThreadState_SetAsyncExc(ctypes.c_ulong(tid), None)
+        raise RuntimeError("Timeout failed")
+
+
+@contextmanager
+def timeout(timeout: float) -> Iterator[Timer]:
+    """
+    Run the wrapped code for at most `timeout` seconds before aborting.
+
+    This works by starting a timer thread, and using "magic" raises an exception
+    in the main process after the timer expires.
+
+    Raises `TimeoutException` when the timeout occurs.
+    """
+    timeout_timer = Timer(timeout, _do_timeout, args=[current_thread().ident])
+    try:
+        timeout_timer.start()
+
+        yield timeout_timer
+    finally:
+        timeout_timer.cancel()

tests/tests/test_utils.py

@@ -1,10 +1,12 @@
 import datetime
 import subprocess
+import time
 from unittest.mock import Mock
 
 from django.test import SimpleTestCase
 
 from django_tasks import utils
+from django_tasks.exceptions import TimeoutException
 from tests import tasks as test_tasks
 
 
@@ -116,3 +118,10 @@ def test_complex_exception(self) -> None:
             self.assertIn("KeyError: datetime.datetime", traceback)
         else:
             self.fail("KeyError not raised")
+
+
+class TimeoutTestCase(SimpleTestCase):
+    def test_sleep_timeout(self) -> None:
+        with self.assertRaises(TimeoutException):
+            with utils.timeout(0.25):
+                time.sleep(0.5)