Enhance wasm with checkpoint and restore support (#2333)

- Add wasm_runtime_checkpoint/wasm_runtime_restore API
- Support AOT and Classic Interpreter mode checkpoint and debug through OS signal, tested on windows/mac/linux aarch64/x64
- Static instrument the AOT to have the checkpoint and restore switches
- Add sub extra library folder for implementing the ckpt-restore
- Include extra dependency of yalantinglib

Co-authored-by: Aibo Hu <[email protected]>
Co-authored-by: kikispace <[email protected]>
Co-authored-by: Brian Zhao <[email protected]>
Signed-off-by: victoryang00 <[email protected]>

Author: 4 people

build-scripts/config_common.cmake

@@ -304,6 +304,10 @@ if (WAMR_BUILD_AOT_STACK_FRAME EQUAL 1)
   add_definitions (-DWASM_ENABLE_AOT_STACK_FRAME=1)
   message ("     AOT stack frame enabled")
 endif ()
+if (WAMR_BUILD_CHECKPOINT_RESTORE EQUAL 1)
+  add_definitions (-DWASM_ENABLE_CHECKPOINT_RESTORE=1)
+  message ("     Checkpoint Restore enabled")
+endif ()
 if (WAMR_BUILD_MEMORY_PROFILING EQUAL 1)
   add_definitions (-DWASM_ENABLE_MEMORY_PROFILING=1)
   message ("     Memory profiling enabled")

build-scripts/runtime_lib.cmake

@@ -86,6 +86,10 @@ if (WAMR_BUILD_LIBC_BUILTIN EQUAL 1)
     include (${IWASM_DIR}/libraries/libc-builtin/libc_builtin.cmake)
 endif ()
 
+if (WAMR_BUILD_CHECKPOINT_RESTORE EQUAL 1)
+    include (${IWASM_DIR}/libraries/ckpt-restore/ckpt_restore.cmake)
+endif ()
+
 if (WAMR_BUILD_LIBC_UVWASI EQUAL 1)
     include (${IWASM_DIR}/libraries/libc-uvwasi/libc_uvwasi.cmake)
     set (WAMR_BUILD_MODULE_INST_CONTEXT 1)
@@ -193,6 +197,7 @@ set (source_all
     ${LIBC_EMCC_SOURCE}
     ${LIB_RATS_SOURCE}
     ${DEBUG_ENGINE_SOURCE}
+    ${CKPT_RESTORE_SOURCE}
 )
 
 set (WAMR_RUNTIME_LIB_SOURCE ${source_all})

core/config.h

@@ -335,6 +335,11 @@
 #define WASM_ENABLE_AOT_STACK_FRAME 0
 #endif
 
+/* Checkpoint Restore */
+#ifndef WASM_ENABLE_CHECKPOINT_RESTORE
+#define WASM_ENABLE_CHECKPOINT_RESTORE 0
+#endif
+
 /* Heap verification */
 #ifndef BH_ENABLE_GC_VERIFY
 #define BH_ENABLE_GC_VERIFY 0

core/iwasm/aot/aot_reloc.h

@@ -57,11 +57,13 @@ typedef struct {
 #define REG_AOT_TRACE_SYM()               \
     REG_SYM(aot_alloc_frame),             \
     REG_SYM(aot_free_frame),              \
+    REG_SYM(aot_raise),                   \
     REG_SYM(aot_frame_update_profile_info),
 #else
 #define REG_AOT_TRACE_SYM()
 #endif
 
+#if WASM_ENABLE_AOT_INTRINSICS != 0
 #define REG_INTRINSIC_SYM()               \
     REG_SYM(aot_intrinsic_fabs_f32),      \
     REG_SYM(aot_intrinsic_fabs_f64),      \
@@ -124,7 +126,10 @@ typedef struct {
     REG_SYM(aot_intrinsic_i32_div_s),     \
     REG_SYM(aot_intrinsic_i32_div_u),     \
     REG_SYM(aot_intrinsic_i32_rem_s),     \
-    REG_SYM(aot_intrinsic_i32_rem_u),     \
+    REG_SYM(aot_intrinsic_i32_rem_u),
+#else
+#define REG_INTRINSIC_SYM()
+#endif
 
 #if WASM_ENABLE_STATIC_PGO != 0
 #define REG_LLVM_PGO_SYM()               \
@@ -241,4 +246,4 @@ apply_relocation(AOTModule *module,
 }
 #endif
 
-#endif /* end of _AOT_RELOC_H_ */
+#endif /* end of _AOT_RELOC_H_ */

core/iwasm/aot/aot_runtime.c

@@ -16,6 +16,11 @@
 #include "../libraries/thread-mgr/thread_manager.h"
 #endif
 
+#include "wasm_interp.h"
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+#include "../libraries/ckpt-restore/ckpt_restore.h"
+#endif
+
 /*
  * Note: These offsets need to match the values hardcoded in
  * AoT compilation code: aot_create_func_context, check_suspend_flags.
@@ -72,6 +77,11 @@ bh_static_assert(offsetof(AOTFrame, sp) == sizeof(uintptr_t) * 5);
 bh_static_assert(offsetof(AOTFrame, frame_ref) == sizeof(uintptr_t) * 6);
 bh_static_assert(offsetof(AOTFrame, lp) == sizeof(uintptr_t) * 7);
 
+bh_static_assert(offsetof(AOTFrame, ip_offset) == sizeof(uintptr_t) * 4);
+bh_static_assert(offsetof(AOTFrame, sp) == sizeof(uintptr_t) * 5);
+bh_static_assert(offsetof(AOTFrame, frame_ref) == sizeof(uintptr_t) * 6);
+bh_static_assert(offsetof(AOTFrame, lp) == sizeof(uintptr_t) * 7);
+
 static void
 set_error_buf(char *error_buf, uint32 error_buf_size, const char *string)
 {
@@ -804,6 +814,7 @@ memory_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
         bh_assert(memory_idx == 0);
         bh_assert(parent->memory_count > memory_idx);
         shared_memory_instance = parent->memories[memory_idx];
+        shared_memory_instance->ref_count++;
         shared_memory_inc_reference(shared_memory_instance);
         return shared_memory_instance;
     }
@@ -992,6 +1003,9 @@ memories_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
     AOTMemoryInstance *memories, *memory_inst;
     AOTMemInitData *data_seg;
     uint64 total_size;
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    bool is_shared_memory;
+#endif
 
     module_inst->memory_count = memory_count;
     total_size = sizeof(AOTMemoryInstance *) * (uint64)memory_count;
@@ -1019,6 +1033,15 @@ memories_instantiate(AOTModuleInstance *module_inst, AOTModuleInstance *parent,
         return true;
     }
 
+#if WASM_ENABLE_SHARED_MEMORY != 0
+    /* Currently we have only one memory instance */
+    is_shared_memory = module->memories[0].memory_flags & 0x02 ? true : false;
+    if (is_shared_memory && parent != NULL) {
+        /* Ignore setting memory init data if the memory has been initialized */
+        return true;
+    }
+#endif
+
     for (i = 0; i < module->mem_init_data_count; i++) {
         data_seg = module->mem_init_data_list[i];
 #if WASM_ENABLE_BULK_MEMORY != 0
@@ -1860,6 +1883,7 @@ destroy_c_api_frames(Vector *frames)
 void
 aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE == 0
     WASMModuleInstanceExtraCommon *common =
         &((AOTModuleInstanceExtra *)module_inst->e)->common;
     if (module_inst->exec_env_singleton) {
@@ -1934,6 +1958,7 @@ aot_deinstantiate(AOTModuleInstance *module_inst, bool is_sub_inst)
 #endif
 
     wasm_runtime_free(module_inst);
+#endif
 }
 
 AOTFunctionInstance *
@@ -2225,6 +2250,7 @@ aot_call_function(WASMExecEnv *exec_env, AOTFunctionInstance *function,
         while (exec_env->cur_frame != prev_frame)
             aot_free_frame(exec_env);
 #endif
+        // checkpoint
         if (!ret) {
             if (argv1 != argv1_buf)
                 wasm_runtime_free(argv1);
@@ -2789,6 +2815,17 @@ aot_call_indirect(WASMExecEnv *exec_env, uint32 tbl_idx, uint32 table_elem_idx,
     }
 
     tbl_elem_val = ((table_elem_type_t *)tbl_inst->elems)[table_elem_idx];
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    if (exec_env->is_restore && exec_env->restore_call_chain) {
+        struct AOTFrame *rcc = *(exec_env->restore_call_chain);
+        while (rcc->prev_frame) {
+            rcc = rcc->prev_frame;
+        }
+        LOG_DEBUG("func_idx: %d instead of %d of thread %ld\n", rcc->func_index,
+                  func_idx, exec_env->handle);
+        func_idx = rcc->func_index;
+    }
+#endif
     if (tbl_elem_val == NULL_REF) {
         aot_set_exception_with_id(module_inst, EXCE_UNINITIALIZED_ELEMENT);
         goto fail;
@@ -3449,6 +3486,12 @@ get_func_name_from_index(const AOTModuleInstance *module_inst,
 #endif /* end of WASM_ENABLE_DUMP_CALL_STACK != 0 || \
           WASM_ENABLE_PERF_PROFILING != 0 */
 
+void
+aot_raise(WASMExecEnv *exec_env, int sig)
+{
+    raise(sig);
+}
+
 #if WASM_ENABLE_GC == 0
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
@@ -3528,6 +3571,9 @@ aot_free_frame(WASMExecEnv *exec_env)
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    LOG_DEBUG("aot_alloc_frame %u thread %d\n", func_index, exec_env->handle);
+#endif
     AOTModuleInstance *module_inst = (AOTModuleInstance *)exec_env->module_inst;
     AOTModule *module = (AOTModule *)module_inst->module;
 #if WASM_ENABLE_PERF_PROFILING != 0
@@ -3537,6 +3583,27 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
     AOTFrame *frame;
     uint32 max_local_cell_num, max_stack_cell_num, all_cell_num;
     uint32 aot_func_idx, frame_size;
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    if (exec_env->restore_call_chain) {
+        frame = exec_env->restore_call_chain[exec_env->call_chain_size - 1];
+        LOG_DEBUG("frame restored, func idx %zu\n", frame->func_index);
+        exec_env->call_chain_size--;
+        frame->prev_frame = (AOTFrame *)exec_env->cur_frame;
+        exec_env->cur_frame = (struct WASMInterpFrame *)frame;
+        if (exec_env->call_chain_size == 0) {
+            // TODO: fix memory leak
+            exec_env->restore_call_chain = NULL;
+                    }
+        LOG_DEBUG("restore call chain %zu==%u, %p, %p, %d\n",
+                ((AOTFrame *)exec_env->cur_frame)->func_index, func_index,
+                exec_env, exec_env->restore_call_chain, exec_env->handle);
+        if (((AOTFrame *)exec_env->cur_frame)->func_index != func_index) {
+            LOG_DEBUG("NOT MATCH!!!\n");
+            exit(1);
+        }
+        return true;
+    }
+#endif
 
     if (func_index >= module->import_func_count) {
         aot_func_idx = func_index - module->import_func_count;
@@ -3568,6 +3635,11 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
     frame->time_started = (uintptr_t)os_time_thread_cputime_us();
     frame->func_perf_prof_info = func_perf_prof;
 #endif
+    frame->ip_offset = 0;
+    frame->sp = frame->lp + max_local_cell_num;
+#if WASM_ENABLE_GC != 0
+    frame->frame_ref = frame->sp + max_stack_cell_num;
+#endif
 
 #if WASM_ENABLE_GC != 0
     frame->sp = frame->lp + max_local_cell_num;
@@ -3584,6 +3656,10 @@ aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index)
 static inline void
 aot_free_frame_internal(WASMExecEnv *exec_env)
 {
+#if WASM_ENABLE_CHECKPOINT_RESTORE != 0
+    int func_index = ((AOTFrame *)exec_env->cur_frame)->func_index;
+    LOG_DEBUG("aot_free_frame %zu %d\n", func_index, exec_env->handle);
+#endif
     AOTFrame *cur_frame = (AOTFrame *)exec_env->cur_frame;
     AOTFrame *prev_frame = cur_frame->prev_frame;
 
@@ -3598,7 +3674,6 @@ aot_free_frame_internal(WASMExecEnv *exec_env)
     if (prev_frame)
         prev_frame->func_perf_prof_info->children_exec_time += time_elapsed;
 #endif
-
     wasm_exec_env_free_wasm_frame(exec_env, cur_frame);
     exec_env->cur_frame = (struct WASMInterpFrame *)prev_frame;
 }

core/iwasm/aot/aot_runtime.h

@@ -673,6 +673,9 @@ aot_table_grow(AOTModuleInstance *module_inst, uint32 tbl_idx,
 bool
 aot_alloc_frame(WASMExecEnv *exec_env, uint32 func_index);
 
+void
+aot_raise(WASMExecEnv *exec_env, int exception);
+
 void
 aot_free_frame(WASMExecEnv *exec_env);
 

core/iwasm/common/wasm_exec_env.c

@@ -72,6 +72,7 @@ wasm_exec_env_create_internal(struct WASMModuleInstanceCommon *module_inst,
     exec_env->wasm_stack.top_boundary =
         exec_env->wasm_stack.bottom + stack_size;
     exec_env->wasm_stack.top = exec_env->wasm_stack.bottom;
+    exec_env->is_checkpoint = false;
 
 #if WASM_ENABLE_AOT != 0
     if (module_inst->module_type == Wasm_Module_AoT) {
@@ -85,6 +86,10 @@ wasm_exec_env_create_internal(struct WASMModuleInstanceCommon *module_inst,
     wasm_runtime_dump_exec_env_mem_consumption(exec_env);
 #endif
 
+    exec_env->is_checkpoint = false;
+    exec_env->is_restore = false;
+    exec_env->call_chain_size = 0;
+    exec_env->restore_call_chain = NULL;
     return exec_env;
 
 #ifdef OS_ENABLE_HW_BOUND_CHECK
@@ -174,6 +179,7 @@ wasm_exec_env_create(struct WASMModuleInstanceCommon *module_inst,
     }
 #endif
 
+
 #if WASM_ENABLE_THREAD_MGR != 0
     /* Create a new cluster for this exec_env */
     if (!(cluster = wasm_cluster_create(exec_env))) {

core/iwasm/common/wasm_exec_env.h

@@ -158,6 +158,12 @@ typedef struct WASMExecEnv {
 
     /* The WASM stack size */
     uint32 wasm_stack_size;
+    /* Whether is checkpoint */
+    bool is_checkpoint;
+    /* Whether is restore */
+    bool is_restore;
+    size_t call_chain_size;
+    struct AOTFrame **restore_call_chain;
 
     /* The WASM stack of current thread */
     union {

core/iwasm/common/wasm_memory.c

@@ -280,6 +280,7 @@ wasm_runtime_get_mem_alloc_info(mem_alloc_info_t *mem_alloc_info)
     return false;
 }
 
+
 bool
 wasm_runtime_validate_app_addr(WASMModuleInstanceCommon *module_inst_comm,
                                uint32 app_offset, uint32 size)
@@ -553,6 +554,8 @@ wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
         return false;
     }
 
+    SHARED_MEMORY_LOCK(memory_inst);
+
     native_addr = memory_inst->memory_data + app_buf_addr;
 
     bounds_checks = is_bounds_checks_enabled((wasm_module_inst_t)module_inst);
@@ -566,32 +569,34 @@ wasm_check_app_addr_and_convert(WASMModuleInstance *module_inst, bool is_str,
 
     /* No need to check the app_offset and buf_size if memory access
        boundary check with hardware trap is enabled */
-#ifndef OS_ENABLE_HW_BOUND_CHECK
-    SHARED_MEMORY_LOCK(memory_inst);
-
-    if (app_buf_addr >= memory_inst->memory_data_size) {
-        goto fail;
-    }
-
-    if (!is_str) {
-        if (app_buf_size > memory_inst->memory_data_size - app_buf_addr) {
-            goto fail;
-        }
-    }
-    else {
-        const char *str, *str_end;
-
-        /* The whole string must be in the linear memory */
-        str = (const char *)native_addr;
-        str_end = (const char *)memory_inst->memory_data_end;
-        while (str < str_end && *str != '\0')
-            str++;
-        if (str == str_end)
-            goto fail;
-    }
+// #ifndef OS_ENABLE_HW_BOUND_CHECK
+//     SHARED_MEMORY_LOCK(memory_inst);
+
+//     if (app_buf_addr >= memory_inst->memory_data_size) {
+//         goto fail;
+//     }
+
+//     if (!is_str) {
+//         if (app_buf_size > memory_inst->memory_data_size - app_buf_addr) {
+//             goto fail;
+//         }
+//     }
+//     else {
+//         const char *str, *str_end;
+
+//         /* The whole string must be in the linear memory */
+//         str = (const char *)native_addr;
+//         str_end = (const char *)memory_inst->memory_data_end;
+//         while (str < str_end && *str != '\0')
+//             str++;
+//         if (str == str_end)
+//             goto fail;
+//     }
+
+//     SHARED_MEMORY_UNLOCK(memory_inst);
+// #endif
 
     SHARED_MEMORY_UNLOCK(memory_inst);
-#endif
 
 success:
     *p_native_addr = (void *)native_addr;