diff --git a/Include/Python.h b/Include/Python.h index 337119c15fe8b6f..d9765e13a3446bd 100644 --- a/Include/Python.h +++ b/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" diff --git a/Include/cpython/pystate.h b/Include/cpython/pystate.h index a9d97e47e005dff..de6dc6a127b3d8e 100644 --- a/Include/cpython/pystate.h +++ b/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); \ No newline at end of file diff --git a/Include/internal/pycore_ceval.h b/Include/internal/pycore_ceval.h index 06c4ca1619d7ce1..e9eb605fc91ccca 100644 --- a/Include/internal/pycore_ceval.h +++ b/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 diff --git a/Lib/contextlib.py b/Lib/contextlib.py index efc02bfa9243da6..6d9aa758db3e6af 100644 --- a/Lib/contextlib.py +++ b/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 \ No newline at end of file diff --git a/Lib/test/test_timeout/__init__.py b/Lib/test/test_timeout/__init__.py new file mode 100644 index 000000000000000..0cfa6b1d098ffcb --- /dev/null +++ b/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) \ No newline at end of file diff --git a/Lib/test/test_timeout/test_timeout_block.py b/Lib/test/test_timeout/test_timeout_block.py new file mode 100644 index 000000000000000..e9e1b03b778496b --- /dev/null +++ b/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") \ No newline at end of file diff --git a/Makefile.pre.in b/Makefile.pre.in index 2b34b009fd745a0..4555b2348c5e950 100644 --- a/Makefile.pre.in +++ b/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 diff --git a/Modules/_sre/sre_lib.h b/Modules/_sre/sre_lib.h index df377905bfae0d0..cbba417f866bd7a 100644 --- a/Modules/_sre/sre_lib.h +++ b/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) diff --git a/Modules/_timeoutmodule.c b/Modules/_timeoutmodule.c new file mode 100644 index 000000000000000..6e5427ba1e5a24c --- /dev/null +++ b/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); +} \ No newline at end of file diff --git a/Modules/config.c.in b/Modules/config.c.in index 704f58506048a3e..f52d56402e10e80 100644 --- a/Modules/config.c.in +++ b/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} }; diff --git a/Python/ceval.c b/Python/ceval.c index a9b31affca9890a..af400312a2bb6f0 100644 --- a/Python/ceval.c +++ b/Python/ceval.c @@ -2,6 +2,7 @@ #include "ceval.h" #include "pycore_long.h" +#include "pycore_time.h" int Py_GetRecursionLimit(void) diff --git a/Python/ceval_gil.c b/Python/ceval_gil.c index 2425bc1b39f0dcc..82bb52ddd61c52e 100644 --- a/Python/ceval_gil.c +++ b/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 diff --git a/Python/pystate.c b/Python/pystate.c index fed1df0173bacf1..daceb25f3f3d491 100644 --- a/Python/pystate.c +++ b/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; +}