gh-151989: Add with timeout feature + re c module cooperative example

Include/Python.h

@@ -114,6 +114,7 @@ __pragma(warning(disable: 4201))
 #include "cpython/cellobject.h"
 #include "iterobject.h"
 #include "cpython/initconfig.h"
+#include "cpython/pytime.h"
 #include "pystate.h"
 #include "cpython/genobject.h"
 #include "descrobject.h"
@@ -123,7 +124,6 @@ __pragma(warning(disable: 4201))
 #include "weakrefobject.h"
 #include "structseq.h"
 #include "cpython/picklebufobject.h"
-#include "cpython/pytime.h"
 #include "codecs.h"
 #include "pythread.h"
 #include "cpython/context.h"

Include/cpython/pystate.h

@@ -61,6 +61,15 @@ typedef struct _stack_chunk {
     PyObject * data[1]; /* Variable sized */
 } _PyStackChunk;
 
+/* To lower the frequency of time check
+   only when step reaches TIMEOUTCHECK_INTERVAL check is performed */
+#define TIMEOUTCHECK_INTERVAL 16
+typedef struct _timeout_block {
+    int skip_counter;
+    PyTime_t deadline;
+    struct _timeout_block *prev; /* the outter timeout block */
+} _PyTimeoutBlock;
+
 /* Minimum size of data stack chunk */
 #define _PY_DATA_STACK_CHUNK_SIZE (16*1024)
 struct _ts {
@@ -253,6 +262,8 @@ struct _ts {
         /* The interpreter guard owned by PyThreadState_EnsureFromView(), if any. */
         PyInterpreterGuard *owned_guard;
     } ensure;
+
+    _PyTimeoutBlock *timeout_block;
 };
 
 /* other API */
@@ -338,3 +349,8 @@ PyAPI_FUNC(void) _PyInterpreterState_SetEvalFrameAllowSpecialization(
     int allow_specialization);
 PyAPI_FUNC(int) _PyInterpreterState_IsSpecializationEnabled(
     PyInterpreterState *interp);
+
+
+PyAPI_FUNC(int) _PyTimeout_Push(PyThreadState *, PyTime_t);
+PyAPI_FUNC(int) _PyTimeout_Pop(PyThreadState *);
+PyAPI_FUNC(int) Py_CheckTimeOut(PyThreadState *, int);

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_EVAL_TIMEOUT_BIT (1U << 8)
 
 /* Reserve a few bits for future use */
-#define _PY_EVAL_EVENTS_BITS 8
+#define _PY_EVAL_EVENTS_BITS 16
 #define _PY_EVAL_EVENTS_MASK ((1 << _PY_EVAL_EVENTS_BITS)-1)
 
 static inline void

Lib/contextlib.py

@@ -12,12 +12,13 @@
     isgeneratorfunction as _isgeneratorfunction,
 )
 from types import GenericAlias
+import _timeout
 
 __all__ = ["asynccontextmanager", "contextmanager", "closing", "nullcontext",
            "AbstractContextManager", "AbstractAsyncContextManager",
            "AsyncExitStack", "ContextDecorator", "ExitStack",
            "redirect_stdout", "redirect_stderr", "suppress", "aclosing",
-           "chdir"]
+           "chdir", "timeout"]
 
 
 class AbstractContextManager(abc.ABC):
@@ -857,3 +858,21 @@ def __enter__(self):
 
     def __exit__(self, *excinfo):
         os.chdir(self._old_cwd.pop())
+
+
+class timeout(AbstractContextManager):
+    """Context Manager for creating a timeout block
+
+    If time is reach within block execution,
+    a `TimeoutError` will be raised
+    """
+    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_timeout/__init__.py

@@ -0,0 +1,5 @@
+import os
+from test.support import load_package_tests
+
+def load_tests(*args):
+    return load_package_tests(os.path.dirname(__file__), *args)

Lib/test/test_timeout/test_timeout_block.py

@@ -0,0 +1,21 @@
+import time
+import unittest
+
+from contextlib import timeout
+
+class TimeOutTest(unittest.TestCase):
+    def test_py(self):
+        with timeout(0.2):
+            time.sleep(0.1)
+
+        with self.assertRaises(TimeoutError):
+            with timeout(0.2):
+                while True:
+                    pass
+
+    def test_c_module_re(self):
+        import re
+        with self.assertRaises(TimeoutError):
+            with timeout(0.2):
+                # A catastrophic case for re match
+                re.match(r"(a+)+b", "a" * 30 + "c")

Makefile.pre.in

@@ -354,7 +354,8 @@ COVERAGE_REPORT_OPTIONS=--rc lcov_branch_coverage=1 --branch-coverage --title "C
 MODULE_OBJS=	\
 		Modules/config.o \
 		Modules/main.o \
-		Modules/gcmodule.o
+		Modules/gcmodule.o \
+		Modules/_timeoutmodule.o
 
 IO_H=		Modules/_io/_iomodule.h
 

Modules/_sre/sre_lib.h

@@ -550,8 +550,13 @@ typedef struct {
 
 #define _MAYBE_CHECK_SIGNALS                                       \
     do {                                                           \
-        if ((0 == (++sigcount & 0xfff)) && PyErr_CheckSignals()) { \
-            RETURN_ERROR(SRE_ERROR_INTERRUPTED);                   \
+        if ((0 == (++sigcount & 0xfff))) {                         \
+            if (PyErr_CheckSignals()) {                            \
+                RETURN_ERROR(SRE_ERROR_INTERRUPTED);               \
+            }                                                      \
+            if (Py_CheckTimeOut(PyThreadState_Get(), 0)) {         \
+                RETURN_ERROR(SRE_ERROR_INTERRUPTED);               \
+            }                                                      \
         }                                                          \
     } while (0)
 

Modules/_timeoutmodule.c

@@ -0,0 +1,130 @@
+#include "Python.h"
+#include "pycore_time.h"
+
+static PyObject *
+_timeout_enter(PyObject *self, PyObject *args)
+{
+    double seconds;
+    if (!PyArg_ParseTuple(args, "d", &seconds))
+        return NULL;
+
+    PyThreadState *tstate = PyThreadState_Get();
+    PyTime_t now;
+    PyTime_Monotonic(&now);
+
+    PyTime_t relative_deadline;
+    if (_PyTime_FromSecondsDouble(seconds, _PyTime_ROUND_TIMEOUT, &relative_deadline) < 0) {
+        PyErr_SetString(PyExc_OverflowError, "timeout value too large");
+        return NULL;
+    }
+    const PyTime_t TIMEOUT_EPS = 1000000;
+    PyTime_t deadline = now + relative_deadline + TIMEOUT_EPS;
+
+    if (_PyTimeout_Push(tstate, deadline) < 0)
+        return NULL;
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+_timeout_leave(PyObject *self, PyObject *noargs)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+    _PyTimeout_Pop(tstate);
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+_test_timeout_init(PyObject *self, PyObject *args)
+{
+    double seconds;
+    if (!PyArg_ParseTuple(args, "d", &seconds))
+        return NULL;
+
+    PyThreadState *tstate = PyThreadState_Get();
+    PyTime_t now;
+    PyTime_Monotonic(&now);
+    PyTime_t deadline = now + (PyTime_t)(seconds * 1e9);
+
+
+    (void)now;
+    (void)deadline;
+    PyTime_t fake_deadline = (PyTime_t)INT64_MAX;
+
+    if (_PyTimeout_Push(tstate, fake_deadline) < 0)
+        return NULL;
+
+    Py_RETURN_NONE;
+}
+
+static int test_skip_interval = 0;
+static PyObject *
+_test_set_timeout_skip_interval(PyObject *self, PyObject *args)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+    _PyTimeoutBlock *block = tstate->timeout_block;
+    if (block != NULL)
+        block->skip_counter = 0;
+    if (!PyArg_ParseTuple(args, "|i", &test_skip_interval))
+        return NULL;
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+_test_timeout_check(PyObject *self, PyObject *noargs)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+    int res = Py_CheckTimeOut(tstate, test_skip_interval);
+    if (res != 0) {
+        PyErr_Clear();
+    }
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+_test_benchmark_timeout_check_raw_loop(PyObject *self, PyObject *args)
+{
+    int interval, loops;
+    if (!PyArg_ParseTuple(args, "ii", &interval, &loops))
+        return NULL;
+
+    PyThreadState *tstate = PyThreadState_Get();
+    for (int i = 0; i < loops; ++i) {
+        Py_CheckTimeOut(tstate, interval);
+    }
+    Py_RETURN_NONE;
+}
+
+static PyObject *
+_test_timeout_cleanup(PyObject *self, PyObject *noargs)
+{
+    PyThreadState *tstate = PyThreadState_Get();
+
+    _PyTimeout_Pop(tstate);
+    Py_RETURN_NONE;
+}
+
+static PyMethodDef timeout_methods[] = {
+    {"enter", _timeout_enter, METH_VARARGS, "Enter a timeout block."},
+    {"leave", _timeout_leave, METH_NOARGS,  "Leave a timeout block."},
+    {"_test_timeout_init", _test_timeout_init, METH_VARARGS,  "test _PyTimeout_Push C Api"},
+    {"_test_timeout_check", _test_timeout_check, METH_NOARGS, "test Py_CheckTimeOut C Api"},
+    {"_test_timeout_cleanup", _test_timeout_cleanup, METH_NOARGS,  "test _PyTimeout_Pop CApi"},
+    {"_test_set_timeout_skip_interval", _test_set_timeout_skip_interval, METH_VARARGS,  "set test_skip_interval"},
+    {"_test_benchmark_timeout_check_raw_loop", _test_benchmark_timeout_check_raw_loop, METH_VARARGS,  "test `Py_CheckTimeOut` CApi"},
+    {NULL}
+};
+
+static struct PyModuleDef _timeoutmodule = {
+    PyModuleDef_HEAD_INIT,
+    "_timeout",
+    NULL,
+    -1,
+    timeout_methods,
+    NULL, NULL, NULL, NULL
+};
+
+PyMODINIT_FUNC
+PyInit__timeout(void)
+{
+    return PyModule_Create(&_timeoutmodule);
+}

Modules/config.c.in

@@ -22,6 +22,7 @@ extern PyObject* PyInit__tokenize(void);
 extern PyObject* PyInit__contextvars(void);
 extern PyObject* _PyWarnings_Init(void);
 extern PyObject* PyInit__string(void);
+extern PyObject* PyInit__timeout(void);
 
 struct _inittab _PyImport_Inittab[] = {
 
@@ -55,6 +56,9 @@ struct _inittab _PyImport_Inittab[] = {
     /* This lives in Objects/unicodeobject.c */
     {"_string", PyInit__string},
 
+    /* This lives in Modules/timeoutmodule.c */
+    {"_timeout", PyInit__timeout},
+
     /* Sentinel */
     {0, 0}
 };

Python/ceval.c

@@ -2,6 +2,7 @@
 
 #include "ceval.h"
 #include "pycore_long.h"
+#include "pycore_time.h"
 
 int
 Py_GetRecursionLimit(void)

Python/ceval_gil.c

@@ -1428,6 +1428,11 @@ _Py_HandlePending(PyThreadState *tstate)
         }
     }
 
+    if ((breaker & _PY_EVAL_TIMEOUT_BIT) != 0) {
+        if (Py_CheckTimeOut(tstate, TIMEOUTCHECK_INTERVAL) != 0) {
+            return -1;
+        }
+    }
 #if defined(Py_REMOTE_DEBUG) && defined(Py_SUPPORTS_REMOTE_DEBUG)
     _PyRunRemoteDebugger(tstate);
 #endif

Python/pystate.c

@@ -1637,6 +1637,7 @@ init_threadstate(_PyThreadStateImpl *_tstate,
     }
 
     tstate->_status.initialized = 1;
+    tstate->timeout_block = NULL;
 }
 
 static void
@@ -3623,3 +3624,57 @@ PyThreadState_Release(PyThreadStateToken *token)
         PyInterpreterGuard_Close(owned_guard);
     }
 }
+
+int
+_PyTimeout_Push(PyThreadState *tstate, PyTime_t deadline)
+{
+    _PyTimeoutBlock *block = PyMem_Malloc(sizeof(_PyTimeoutBlock));
+    if (block == NULL) {
+        PyErr_NoMemory();
+        return -1;
+    }
+    block->skip_counter = 0;
+    block->deadline = deadline;
+    block->prev = tstate->timeout_block;
+    tstate->timeout_block = block;
+
+    _Py_set_eval_breaker_bit(tstate, _PY_EVAL_TIMEOUT_BIT);
+    return 0;
+}
+
+int
+_PyTimeout_Pop(PyThreadState *tstate)
+{
+    _PyTimeoutBlock *block = tstate->timeout_block;
+    if (block == NULL) return 0;
+    tstate->timeout_block = block->prev;
+    PyMem_Free(block);
+    if (tstate->timeout_block == NULL) {
+        _Py_unset_eval_breaker_bit(tstate, _PY_EVAL_TIMEOUT_BIT);
+    }
+    return 0;
+}
+
+/* Check if current tstate has reached a timeout
+   when you want instant timeout check, skip_interval should be 0
+*/
+int
+Py_CheckTimeOut(PyThreadState *tstate, int skip_interval)
+{
+    _PyTimeoutBlock *block = tstate->timeout_block;
+    if (block == NULL) {
+        return 0;
+    }
+    if (skip_interval && ++block->skip_counter < skip_interval) {
+        return 0;
+    }
+    PyTime_t now;
+    PyTime_Monotonic(&now);
+    block->skip_counter = 0;
+    if (now > block->deadline) {
+        block->deadline = (PyTime_t)INT64_MAX;
+        PyErr_SetString(PyExc_TimeoutError, "Timeout");
+        return -1;
+    }
+    return 0;
+}