diff --git a/CMakeLists.txt b/CMakeLists.txt index a8dc8156..68afcf7f 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -48,6 +48,10 @@ if(PK_ENABLE_OS) add_definitions(-DPK_ENABLE_OS=1) endif() +if(PK_ENABLE_WATCHDOG) + add_definitions(-DPK_ENABLE_WATCHDOG=1) +endif() + if(PK_BUILD_MODULE_LZ4) add_subdirectory(3rd/lz4) include_directories(3rd/lz4) diff --git a/CMakeOptions.txt b/CMakeOptions.txt index d9cefc70..4e21ab44 100644 --- a/CMakeOptions.txt +++ b/CMakeOptions.txt @@ -8,6 +8,7 @@ endif() # system features option(PK_ENABLE_OS "" OFF) option(PK_ENABLE_DETERMINISM "" FALSE) +option(PK_ENABLE_WATCHDOG "" OFF) # modules option(PK_BUILD_MODULE_LZ4 "" OFF) diff --git a/include/pocketpy/config.h b/include/pocketpy/config.h index 18ca5d68..12fcbb10 100644 --- a/include/pocketpy/config.h +++ b/include/pocketpy/config.h @@ -12,6 +12,10 @@ #define PK_ENABLE_OS 1 #endif +#ifndef PK_ENABLE_WATCHDOG // can be overridden by cmake +#define PK_ENABLE_WATCHDOG 0 +#endif + // GC min threshold #ifndef PK_GC_MIN_THRESHOLD // can be overridden by cmake #define PK_GC_MIN_THRESHOLD 32768 diff --git a/include/pocketpy/interpreter/vm.h b/include/pocketpy/interpreter/vm.h index 8a13dcb7..2dabd1d4 100644 --- a/include/pocketpy/interpreter/vm.h +++ b/include/pocketpy/interpreter/vm.h @@ -23,6 +23,11 @@ typedef struct TraceInfo { py_TraceFunc func; } TraceInfo; +typedef struct WatchdogInfo { + py_i64 timeout; + py_i64 last_reset_time; +} WatchdogInfo; + typedef struct VM { py_Frame* top_frame; @@ -48,6 +53,7 @@ typedef struct VM { py_StackRef curr_class; py_StackRef curr_decl_based_function; TraceInfo trace_info; + WatchdogInfo watchdog_info; py_TValue vectorcall_buffer[PK_MAX_CO_VARNAMES]; InternedNames names; diff --git a/include/pocketpy/pocketpy.h b/include/pocketpy/pocketpy.h index dff6af1d..a61f023b 100644 --- a/include/pocketpy/pocketpy.h +++ b/include/pocketpy/pocketpy.h @@ -111,6 +111,16 @@ PK_API void py_sys_settrace(py_TraceFunc func); /// Setup the callbacks for the current VM. PK_API py_Callbacks* py_callbacks(); +/// Begin the watchdog with a timeout in milliseconds. +/// `PK_ENABLE_WATCHDOG` must be defined to `1` to use this feature. +/// You need to call `py_watchdog_reset()` periodically to keep the watchdog alive. +/// If the timeout is reached, `TimeoutError` will be raised. +PK_API void py_watchdog_begin(py_i64 timeout); +/// Reset the watchdog. +PK_API void py_watchdog_reset(); +/// End the watchdog. +PK_API void py_watchdog_end(); + /// Get the current source location of the frame. PK_API const char* py_Frame_sourceloc(py_Frame* frame, int* lineno); /// Python equivalent to `globals()` with respect to the given frame. @@ -553,6 +563,7 @@ PK_API void py_clearexc(py_StackRef p0); #define NameError(n) py_exception(tp_NameError, "name '%n' is not defined", (n)) #define TypeError(...) py_exception(tp_TypeError, __VA_ARGS__) #define RuntimeError(...) py_exception(tp_RuntimeError, __VA_ARGS__) +#define TimeoutError(...) py_exception(tp_TimeoutError, __VA_ARGS__) #define OSError(...) py_exception(tp_OSError, __VA_ARGS__) #define ValueError(...) py_exception(tp_ValueError, __VA_ARGS__) #define IndexError(...) py_exception(tp_IndexError, __VA_ARGS__) @@ -757,6 +768,7 @@ enum py_PredefinedType { tp_IndexError, tp_ValueError, tp_RuntimeError, + tp_TimeoutError, tp_ZeroDivisionError, tp_NameError, tp_UnboundLocalError, diff --git a/include/typings/pkpy.pyi b/include/typings/pkpy.pyi index 793897d3..69b9c691 100644 --- a/include/typings/pkpy.pyi +++ b/include/typings/pkpy.pyi @@ -25,6 +25,19 @@ def currentvm() -> int: """Return the current VM index.""" +def watchdog_begin(timeout: int): + """ + Begin the watchdog with a timeout in milliseconds. + `PK_ENABLE_WATCHDOG` must be defined to `1` to use this feature. + You need to call `watchdog_reset()` periodically to keep the watchdog alive. + If the timeout is reached, `TimeoutError` will be raised. + """ +def watchdog_reset(): + """Reset the watchdog.""" +def watchdog_end(): + """End the watchdog.""" + + class ComputeThread: def __init__(self, vm_index: Literal[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]): ... diff --git a/src/interpreter/ceval.c b/src/interpreter/ceval.c index 90ebfd49..3f889bae 100644 --- a/src/interpreter/ceval.c +++ b/src/interpreter/ceval.c @@ -7,6 +7,7 @@ #include "pocketpy/pocketpy.h" #include "pocketpy/objects/error.h" #include +#include static bool stack_format_object(VM* self, c11_sv spec); @@ -107,6 +108,18 @@ FrameResult VM__run_top_frame(VM* self) { } } +#if PK_ENABLE_WATCHDOG + if(self->watchdog_info.timeout > 0){ + py_i64 now = clock() / (CLOCKS_PER_SEC / 1000); + py_i64 delta = now - self->watchdog_info.last_reset_time; + if(delta > self->watchdog_info.timeout) { + self->watchdog_info.last_reset_time = now; + TimeoutError("watchdog timeout"); + goto __ERROR; + } + } +#endif + #ifndef NDEBUG pk_print_stack(self, frame, byte); #endif diff --git a/src/interpreter/vm.c b/src/interpreter/vm.c index cbd94ed9..e9b5b6e3 100644 --- a/src/interpreter/vm.c +++ b/src/interpreter/vm.c @@ -170,6 +170,7 @@ void VM__ctor(VM* self) { INJECT_BUILTIN_EXC(IndexError, tp_Exception); INJECT_BUILTIN_EXC(ValueError, tp_Exception); INJECT_BUILTIN_EXC(RuntimeError, tp_Exception); + INJECT_BUILTIN_EXC(TimeoutError, tp_Exception); INJECT_BUILTIN_EXC(ZeroDivisionError, tp_Exception); INJECT_BUILTIN_EXC(NameError, tp_Exception); INJECT_BUILTIN_EXC(UnboundLocalError, tp_Exception); diff --git a/src/modules/pkpy.c b/src/modules/pkpy.c index 66999c99..62c34533 100644 --- a/src/modules/pkpy.c +++ b/src/modules/pkpy.c @@ -3,7 +3,6 @@ #include "pocketpy/pocketpy.h" #include "pocketpy/common/utils.h" -#include "pocketpy/objects/object.h" #include "pocketpy/common/sstream.h" #include "pocketpy/interpreter/vm.h" @@ -79,6 +78,46 @@ static bool pkpy_currentvm(int argc, py_Ref argv) { return true; } +#if PK_ENABLE_WATCHDOG +void py_watchdog_begin(py_i64 timeout) { + WatchdogInfo* info = &pk_current_vm->watchdog_info; + info->timeout = timeout; + py_watchdog_reset(); +} + +void py_watchdog_reset() { + WatchdogInfo* info = &pk_current_vm->watchdog_info; + info->last_reset_time = clock() / (CLOCKS_PER_SEC / 1000); +} + +void py_watchdog_end() { + WatchdogInfo* info = &pk_current_vm->watchdog_info; + info->timeout = 0; +} + +static bool pkpy_watchdog_begin(int argc, py_Ref argv) { + PY_CHECK_ARGC(1); + PY_CHECK_ARG_TYPE(0, tp_int); + py_watchdog_begin(py_toint(argv)); + py_newnone(py_retval()); + return true; +} + +static bool pkpy_watchdog_reset(int argc, py_Ref argv) { + PY_CHECK_ARGC(0); + py_watchdog_reset(); + py_newnone(py_retval()); + return true; +} + +static bool pkpy_watchdog_end(int argc, py_Ref argv) { + PY_CHECK_ARGC(0); + py_watchdog_end(); + py_newnone(py_retval()); + return true; +} +#endif + typedef struct c11_ComputeThread c11_ComputeThread; typedef struct { @@ -467,6 +506,12 @@ void pk__add_module_pkpy() { py_bindfunc(mod, "currentvm", pkpy_currentvm); +#if PK_ENABLE_WATCHDOG + py_bindfunc(mod, "watchdog_begin", pkpy_watchdog_begin); + py_bindfunc(mod, "watchdog_reset", pkpy_watchdog_reset); + py_bindfunc(mod, "watchdog_end", pkpy_watchdog_end); +#endif + pk_ComputeThread__register(mod); }