Add basic timeout helper

It's not perfect, as it doesn't work if the block is outside Python.

Author: RealOrangeOne

django_tasks/exceptions.py

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

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)