add PK_ENABLE_WATCHDOG

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)

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)

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

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;

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,

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]): ...
 

src/interpreter/ceval.c

@@ -7,6 +7,7 @@
 #include "pocketpy/pocketpy.h"
 #include "pocketpy/objects/error.h"
 #include <stdbool.h>
+#include <time.h>
 
 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

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);

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);
 }