Cooperative timeout context manager--the deadline-scheduler approach

Include/cpython/pystate.h

@@ -61,6 +61,8 @@ typedef struct _stack_chunk {
     PyObject * data[1]; /* Variable sized */
 } _PyStackChunk;
 
+struct _timeout_block;
+
 /* Minimum size of data stack chunk */
 #define _PY_DATA_STACK_CHUNK_SIZE (16*1024)
 struct _ts {
@@ -253,6 +255,9 @@ struct _ts {
         /* The interpreter guard owned by PyThreadState_EnsureFromView(), if any. */
         PyInterpreterGuard *owned_guard;
     } ensure;
+
+    uintptr_t cancel_flags;
+    struct _timeout_block *timeout_block;
 };
 
 /* other API */

Include/internal/pycore_ceval.h

@@ -349,9 +349,10 @@ _PyEval_SpecialMethodCanSuggest(PyObject *self, int oparg);
 #define _PY_EVAL_PLEASE_STOP_BIT (1U << 5)
 #define _PY_EVAL_EXPLICIT_MERGE_BIT (1U << 6)
 #define _PY_EVAL_JIT_INVALIDATE_COLD_BIT (1U << 7)
+#define _PY_CANCEL_REQUESTED_BIT (1U << 8)
 
 /* Reserve a few bits for future use */
-#define _PY_EVAL_EVENTS_BITS 8
+#define _PY_EVAL_EVENTS_BITS 9
 #define _PY_EVAL_EVENTS_MASK ((1 << _PY_EVAL_EVENTS_BITS)-1)
 
 static inline void

Include/internal/pycore_pystate.h

@@ -171,6 +171,23 @@ extern void _PyEval_StartTheWorldAll(_PyRuntimeState *runtime);
 extern PyAPI_FUNC(void) _PyEval_StopTheWorld(PyInterpreterState *interp);
 extern PyAPI_FUNC(void) _PyEval_StartTheWorld(PyInterpreterState *interp);
 
+#define _PY_CANCEL_GENERIC (1U << 0)
+#define _PY_CANCEL_TIMEOUT (1U << 1)
+
+extern void _PyThreadState_RequestCancel(PyThreadState *tstate,
+                                         uintptr_t reason);
+extern int _PyThreadState_RequestCancelByThreadId(PyInterpreterState *interp,
+                                                  unsigned long id,
+                                                  uintptr_t reason);
+extern int _PyThreadState_CheckCancellation(PyThreadState *tstate);
+extern void _PyThreadState_ClearCancellation(PyThreadState *tstate,
+                                             uintptr_t reason);
+
+extern int _PyTimeout_Push(PyThreadState *tstate, PyTime_t timeout);
+extern int _PyTimeout_Pop(PyThreadState *tstate);
+extern int _PyTimeout_CheckNow(PyThreadState *tstate);
+extern void _PyTimeout_ClearThread(PyThreadState *tstate);
+
 
 static inline void
 _Py_EnsureFuncTstateNotNULL(const char *func, PyThreadState *tstate)

Include/internal/pycore_runtime_structs.h

@@ -8,7 +8,9 @@ extern "C" {
 #endif
 
 #include "pycore_interp_structs.h" // _PyGC_Head_UNUSED
+#include "pycore_condvar.h"       // PyCOND_T
 #include "pycore_obmalloc.h"      // struct _obmalloc_global_state
+#include "pycore_pythread.h"      // PyThread_handle_t
 
 /************ Runtime state ************/
 
@@ -87,6 +89,20 @@ struct _Py_time_runtime_state {
 #endif
 };
 
+struct _timeout_block;
+
+struct _timeout_scheduler_state {
+    PyMUTEX_T mutex;
+    PyCOND_T cond;
+    PyThread_ident_t ident;
+    PyThread_handle_t handle;
+    int initialized;
+    int running;
+    int stopping;
+    struct _timeout_block *head;
+    _PyOnceFlag once;
+};
+
 
 struct _Py_cached_objects {
     // XXX We could statically allocate the hashtable.
@@ -270,6 +286,7 @@ struct pyruntimestate {
     struct _Py_unicode_runtime_state unicode_state;
     struct _types_runtime_state types;
     struct _Py_time_runtime_state time;
+    struct _timeout_scheduler_state timeout_scheduler;
 
 #if defined(__EMSCRIPTEN__) && defined(PY_CALL_TRAMPOLINE)
     // Used in "Python/emscripten_trampoline.c" to choose between wasm-gc

Lib/contextlib.py

@@ -4,6 +4,7 @@
 import os
 import sys
 import _collections_abc
+import _timeout
 from collections import deque
 from functools import wraps
 lazy from inspect import (
@@ -17,7 +18,7 @@
            "AbstractContextManager", "AbstractAsyncContextManager",
            "AsyncExitStack", "ContextDecorator", "ExitStack",
            "redirect_stdout", "redirect_stderr", "suppress", "aclosing",
-           "chdir"]
+           "chdir", "timeout"]
 
 
 class AbstractContextManager(abc.ABC):
@@ -857,3 +858,18 @@ def __enter__(self):
 
     def __exit__(self, *excinfo):
         os.chdir(self._old_cwd.pop())
+
+
+class timeout(AbstractContextManager):
+    """Context manager that raises TimeoutError after the given delay."""
+
+    def __init__(self, seconds):
+        self._seconds = seconds
+
+    def __enter__(self):
+        _timeout.enter(self._seconds)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        _timeout.leave()
+        return False

Lib/test/test_contextlib.py

@@ -4,12 +4,19 @@
 import os
 import sys
 import tempfile
+import textwrap
 import threading
 import traceback
 import unittest
+import _timeout
+try:
+    import _interpreters
+except ModuleNotFoundError:
+    _interpreters = None
 from contextlib import *  # Tests __all__
 from test import support
-from test.support import os_helper
+from test.support import os_helper, script_helper
+from test.support import warnings_helper
 from test.support.testcase import ExceptionIsLikeMixin
 import weakref
 
@@ -1579,5 +1586,222 @@
         self.assertEqual(os.getcwd(), old_cwd)
 
 
+class TestTimeout(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        try:
+            _timeout.enter(support.SHORT_TIMEOUT)
+        except RuntimeError as exc:
+            if "failed to start timeout scheduler thread" in str(exc):
+                raise unittest.SkipTest("timeout scheduler thread unavailable")
+            raise
+        else:
+            _timeout.leave()
+
+    @unittest.skipIf(_interpreters is None, "subinterpreters required")
+    @unittest.skipIf(
+        support.is_android or support.is_apple_mobile,
+        "Subinterpreters are not supported on Android and iOS"
+    )
+    def test_import_contextlib_in_subinterpreter(self):
+        interp = _interpreters.create()
+        try:
+            excsnap = _interpreters.run_string(interp, "import contextlib")
+        finally:
+            _interpreters.destroy(interp)
+
+        self.assertIsNone(excsnap, excsnap)
+
+    def test_normal_exit(self):
+        with timeout(support.SHORT_TIMEOUT) as cm:
+            self.assertIsInstance(cm, timeout)
+
+    def test_negative_timeout(self):
+        with self.assertRaises(ValueError):
+            with timeout(-1):
+                pass
+
+    def test_direct_check_and_leave_before_enter(self):
+        code = textwrap.dedent("""
+            import _timeout
+
+            assert _timeout.check() is False
+            try:
+                _timeout.leave()
+            except RuntimeError:
+                pass
+            else:
+                raise AssertionError("inactive timeout did not fail")
+            print("ok")
+        """)
+        _, out, err = script_helper.assert_python_ok("-c", code)
+        self.assertEqual(out.splitlines(), [b"ok"])
+        self.assertEqual(err, b"")
+
+    def test_direct_leave_after_scheduler_init(self):
+        with timeout(support.SHORT_TIMEOUT):
+            pass
+
+        with self.assertRaises(RuntimeError):
+            _timeout.leave()
+
+    def test_direct_cancel_current_thread(self):
+        with self.assertRaisesRegex(RuntimeError, "thread cancelled"):
+            _timeout.cancel()
+            _timeout.check()
+
+    def test_direct_cancel_current_thread_at_eval_breaker(self):
+        with self.assertRaisesRegex(RuntimeError, "thread cancelled"):
+            _timeout.cancel()
+            pass
+
+    def test_cancel_unknown_thread(self):
+        self.assertEqual(_timeout.cancel(0), 0)
+
+    def test_cancel_thread_by_ident(self):
+        ready = threading.Event()
+        stop = threading.Event()
+        result = []
+
+        def worker():
+            ready.set()
+            try:
+                while not stop.is_set():
+                    pass
+            except RuntimeError as exc:
+                result.append(str(exc))
+
+        thread = threading.Thread(target=worker)
+        thread.start()
+        self.addCleanup(thread.join, support.SHORT_TIMEOUT)
+        self.addCleanup(stop.set)
+
+        self.assertTrue(ready.wait(support.SHORT_TIMEOUT))
+        self.assertEqual(_timeout.cancel(thread.ident), 1)
+        thread.join(support.SHORT_TIMEOUT)
+        if thread.is_alive():
+            stop.set()
+            thread.join(support.SHORT_TIMEOUT)
+
+        self.assertFalse(thread.is_alive())
+        self.assertEqual(result, ["thread cancelled"])
+
+    def test_timeout_expires_in_pure_python(self):
+        with self.assertRaises(TimeoutError):
+            with timeout(0.01):
+                while True:
+                    pass
+
+    def test_direct_check_expired_timeout(self):
+        with self.assertRaises(TimeoutError):
+            with timeout(0):
+                _timeout.check()
+
+        with self.assertRaises(RuntimeError):
+            _timeout.leave()
+
+    def test_finally_runs_after_timeout(self):
+        state = []
+
+        with self.assertRaises(TimeoutError):
+            with timeout(0.01):
+                try:
+                    while True:
+                        pass
+                finally:
+                    state.append("finally")
+
+        self.assertEqual(state, ["finally"])
+
+    def test_nested_timeout_uses_earliest_deadline(self):
+        with self.assertRaises(TimeoutError):
+            with timeout(0.01):
+                with timeout(support.SHORT_TIMEOUT):
+                    while True:
+                        pass
+
+    def test_sre_checks_timeout_cancellation(self):
+        import re
+
+        with self.assertRaises(TimeoutError):
+            with timeout(0.01):
+                re.fullmatch(r"(a+)+\Z", "a" * 100_000 + "!")
+
+    def test_thread_clear_removes_active_timeout(self):
+        code = textwrap.dedent("""
+            import _timeout
+            import threading
+
+            def worker():
+                _timeout.enter(3600)
+
+            thread = threading.Thread(target=worker)
+            thread.start()
+            thread.join()
+            print("ok")
+        """)
+        _, out, err = script_helper.assert_python_ok("-c", code)
+        self.assertEqual(out.splitlines(), [b"ok"])
+        self.assertEqual(err, b"")
+
+    @support.requires_fork()
+    @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+    def test_timeout_works_after_fork_without_active_timeout(self):
+        with timeout(support.SHORT_TIMEOUT):
+            pass
+
+        pid = os.fork()
+        if pid == 0:
+            try:
+                with timeout(0.05):
+                    while True:
+                        pass
+            except TimeoutError:
+                os._exit(0)
+            except BaseException:
+                os._exit(2)
+            else:
+                os._exit(1)
+
+        support.wait_process(pid, exitcode=0, timeout=support.SHORT_TIMEOUT)
+
+    @support.requires_fork()
+    @warnings_helper.ignore_fork_in_thread_deprecation_warnings()
+    def test_active_timeout_works_after_fork(self):
+        with timeout(0.5):
+            pid = os.fork()
+            if pid == 0:
+                try:
+                    while True:
+                        pass
+                except TimeoutError:
+                    os._exit(0)
+                except BaseException:
+                    os._exit(2)
+                else:
+                    os._exit(1)
+
+        support.wait_process(pid, exitcode=0, timeout=support.SHORT_TIMEOUT)
+
+    @unittest.skipUnless(support.Py_GIL_DISABLED, "requires free-threaded build")
+    def test_timeout_module_does_not_enable_gil(self):
+        code = textwrap.dedent("""
+            import sys
+
+            assert not sys._is_gil_enabled()
+            import _timeout
+            assert not sys._is_gil_enabled()
+            _timeout.enter(1.0)
+            _timeout.leave()
+            assert not sys._is_gil_enabled()
+            print("ok")
+        """)
+        _, out, err = script_helper.assert_python_ok(
+            "-c", code, PYTHON_GIL="0", __isolated=False)
+        self.assertEqual(out.splitlines(), [b"ok"])
+        self.assertEqual(err, b"")
+
+
 if __name__ == "__main__":
     unittest.main()