Use hotspot-style checkpoint mechanism

Dockerfile

@@ -1,4 +1,4 @@
-FROM debian:bookworm-slim AS builder
+FROM ubuntu:22.04 AS builder
 
 WORKDIR /wanco
 

a.c

@@ -0,0 +1,54 @@
+#include <execinfo.h>
+#include <pthread.h>
+#include <signal.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/mman.h>
+#include <sys/ucontext.h>
+#include <unistd.h>
+
+int *safepoint = NULL;
+
+// シグナルハンドラ
+void signal_segv_handler(int signo, siginfo_t *info, void *context) {
+  void *buffer[1000];
+  int nptrs;
+  ucontext_t *ucontext;
+  ucontext = (ucontext_t *)context;
+
+  nptrs = backtrace(buffer, sizeof(buffer) / sizeof(void *));
+
+  // バックトレースを標準出力に表示
+  backtrace_symbols_fd(buffer, nptrs, fileno(stdout));
+  exit(0);
+}
+
+void signal_usr1_handler(int signo) { mprotect(safepoint, 4096, PROT_NONE); }
+
+// スレッドの処理
+void cause_segfault() {
+  while (1) {
+    int hoge = *safepoint;
+    printf("hoge: %d\n", hoge);
+    sleep(1);
+  }
+}
+
+int main() {
+  // mmap
+  safepoint = mmap(NULL, 4096, PROT_READ | PROT_WRITE,
+                   MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  *safepoint = 42;
+
+  struct sigaction sa;
+
+  sa.sa_sigaction = signal_segv_handler;
+  sa.sa_flags = SA_SIGINFO;
+  sigaction(SIGSEGV, &sa, NULL);
+
+  signal(SIGUSR1, signal_usr1_handler);
+
+  cause_segfault();
+
+  return 0;
+}

b.c

@@ -0,0 +1,32 @@
+#include <stdio.h>
+#include <stdint.h>
+#include <unistd.h>
+#include <poll.h>
+#include <sys/eventfd.h>
+
+int main() {
+    int efd = eventfd(0, 0);  // eventfd作成
+    if (efd == -1) {
+        perror("eventfd");
+        return 1;
+    }
+
+    struct pollfd pfd = { .fd = efd, .events = POLLIN };
+
+    printf("Waiting for event...\n");
+
+    if (fork() == 0) {  // 子プロセスで通知
+        sleep(1);
+        uint64_t u = 1;
+        write(efd, &u, sizeof(u));  // カウンタを増やす
+        return 0;
+    }
+
+    poll(&pfd, 1, -1);  // イベント待機
+    uint64_t value;
+    read(efd, &value, sizeof(value));  // カウンタをリセット
+    printf("Event received!\n");
+
+    close(efd);
+    return 0;
+}

lib-rt/aot.h

@@ -19,7 +19,7 @@ extern "C" Checkpoint chkpt;
 
 } // namespace wanco
 
-extern "C" struct ExecEnv {
+extern "C" struct __attribute__((packed)) ExecEnv {
   int8_t *memory_base;
   int32_t memory_size;
   wanco::MigrationState migration_state;
@@ -30,6 +30,8 @@ extern "C" struct ExecEnv {
 // defined in wasm AOT module
 extern "C" const int32_t INIT_MEMORY_SIZE;
 extern "C" void aot_main(ExecEnv *);
+extern "C" void store_globals(ExecEnv *);
+extern "C" void store_table(ExecEnv *);
 
 // defined in wrt.c
 extern "C" ExecEnv exec_env;

lib-rt/api.cc

@@ -1,17 +1,11 @@
 #include "aot.h"
-#include "arch/arch.h"
 #include "chkpt/chkpt.h"
-#include "elf/elf.h"
-#include "osr/wasm_stacktrace.h"
-#include "stackmap/stackmap.h"
-#include "stacktrace/stacktrace.h"
 #include "wanco.h"
 #include <chrono>
 #include <cstdint>
 #include <cstdlib>
 #include <fstream>
 #include <iostream>
-#include <optional>
 #include <sys/mman.h>
 #include <thread>
 #include <ucontext.h>
@@ -24,10 +18,6 @@ int32_t extend_memory(ExecEnv *exec_env, int32_t inc_pages);
 
 } // namespace wanco
 
-// defined by AOT module
-extern "C" void store_globals(ExecEnv *);
-extern "C" void store_table(ExecEnv *);
-
 extern "C" int32_t memory_grow(ExecEnv *exec_env, int32_t inc_pages) {
   return wanco::extend_memory(exec_env, inc_pages);
 }
@@ -51,68 +41,6 @@ extern "C" void sleep_msec(ExecEnv *exec_env, int32_t ms) {
 ** checkpoint related functions
 */
 
-extern "C" void start_checkpoint(ExecEnv *exec_env) {
-  WANCO_SAVE_REGISTERS();
-  wanco::CallerSavedRegisters regs{};
-  WANCO_RESTORE_REGISTERS(regs);
-
-  // override migration state
-  exec_env->migration_state = wanco::MigrationState::STATE_CHECKPOINT_CONTINUE;
-
-  // auto regs = wanco::stackmap::CallerSavedRegisters{};
-  Info() << "Checkpoint started" << std::endl;
-
-  wanco::ElfFile elf_file{"/proc/self/exe"};
-  auto stackmap_section = elf_file.get_section_data(".llvm_stackmaps");
-  if (!stackmap_section.has_value()) {
-    Fatal() << "Failed to get stackmap section" << std::endl;
-    exit(1);
-  }
-
-  wanco::stackmap::Stackmap stackmap =
-      wanco::stackmap::parse_stackmap(stackmap_section.value());
-
-  const auto native_trace = wanco::get_stack_trace();
-
-  const auto wasm_trace = wanco::asr_exit(regs, native_trace, stackmap);
-
-  Debug() << "Wasm trace:" << std::endl;
-  for (const auto &frame : wasm_trace) {
-    Debug() << frame.to_string() << std::endl;
-  }
-
-  // store the call stack
-  for (const auto &frame : wasm_trace) {
-    wanco::chkpt.frames.push_front(wanco::Frame{
-        .fn_index = frame.loc.get_func(),
-        .pc = frame.loc.get_insn(),
-        .locals = frame.locals,
-        .stack = frame.stack,
-    });
-  }
-
-  // store the globals, table, and memory
-  store_globals(exec_env);
-  store_table(exec_env);
-  wanco::chkpt.memory_size = exec_env->memory_size;
-
-  // write snapshot
-  {
-    std::ofstream ofs("checkpoint.pb");
-    encode_checkpoint_proto(ofs, wanco::chkpt, exec_env->memory_base);
-    Info() << "Snapshot has been saved to checkpoint.pb" << '\n';
-
-    auto time = std::chrono::duration_cast<std::chrono::microseconds>(
-                    std::chrono::system_clock::now().time_since_epoch())
-                    .count();
-    time = time - wanco::CHKPT_START_TIME;
-    // TODO(tamaron): remove this (research purpose)
-    std::ofstream chktime("chkpt-time.txt");
-    chktime << time << '\n';
-  }
-  exit(0);
-}
-
 extern "C" void push_frame(ExecEnv *exec_env) {
   ASSERT(exec_env->migration_state ==
              wanco::MigrationState::STATE_CHECKPOINT_CONTINUE &&
@@ -288,11 +216,11 @@ void dump_checkpoint(wanco::Checkpoint &chkpt) {
 */
 
 static void check_restore_finished(ExecEnv *exec_env, bool cond) {
-  DEBUG_LOG << "Rest frame size: " << std::dec << wanco::chkpt.frames.size()
-            << std::endl;
+  DEBUG_LOG << "Rest frame size to restore: " << std::dec
+            << wanco::chkpt.frames.size() << std::endl;
   if (cond) {
     exec_env->migration_state = wanco::MigrationState::STATE_NONE;
-    DEBUG_LOG << " Restore completed" << std::endl;
+    DEBUG_LOG << "Restore completed" << std::endl;
     ASSERT(wanco::chkpt.restore_stack.empty() && "Stack not empty");
     ASSERT(wanco::chkpt.frames.empty() && "Frames not empty");
     // equivalent to date +%s.%N

lib-rt/arch/x86_64.h

@@ -2,27 +2,10 @@
 #pragma once
 #include "wanco.h"
 #include <cstdint>
+#include <libunwind-x86_64.h>
+#include <libunwind.h>
 #include <string>
-
-#define WANCO_SAVE_REGISTERS()                                                 \
-  asm volatile(".intel_syntax noprefix \n\t"                                   \
-               "push rbx \n\t"                                                 \
-               "push r12 \n\t"                                                 \
-               "push r13 \n\t"                                                 \
-               "push r14 \n\t"                                                 \
-               "push r15 \n\t"                                                 \
-               ".att_syntax \n\t");
-
-#define WANCO_RESTORE_REGISTERS(regs)                                          \
-  asm volatile(".intel_syntax noprefix \n\t"                                   \
-               "pop %0 \n\t"                                                   \
-               "pop %1 \n\t"                                                   \
-               "pop %2 \n\t"                                                   \
-               "pop %3 \n\t"                                                   \
-               "pop %4 \n\t"                                                   \
-               ".att_syntax \n\t"                                              \
-               : "=r"((regs).r15), "=r"((regs).r14), "=r"((regs).r13),         \
-                 "=r"((regs).r12), "=r"((regs).rbx));
+#include <sys/ucontext.h>
 
 namespace wanco {
 
@@ -218,6 +201,49 @@ struct CallerSavedRegisters {
       exit(1);
     }
   }
+
+  static CallerSavedRegisters from_unw_cursor(unw_cursor_t *cursor) {
+    CallerSavedRegisters regs{};
+    unw_get_reg(cursor, UNW_X86_64_RBX, &regs.rbx);
+    unw_get_reg(cursor, UNW_X86_64_R12, &regs.r12);
+    unw_get_reg(cursor, UNW_X86_64_R13, &regs.r13);
+    unw_get_reg(cursor, UNW_X86_64_R14, &regs.r14);
+    unw_get_reg(cursor, UNW_X86_64_R15, &regs.r15);
+    return regs;
+  }
 };
 
+inline unw_context_t convert_ucontext(ucontext_t *uc) {
+  unw_context_t unw_ctx;
+  unw_getcontext(&unw_ctx);
+  unw_ctx.uc_mcontext.gregs[REG_RAX] = uc->uc_mcontext.gregs[REG_RAX];
+  unw_ctx.uc_mcontext.gregs[REG_RDX] = uc->uc_mcontext.gregs[REG_RDX];
+  unw_ctx.uc_mcontext.gregs[REG_RCX] = uc->uc_mcontext.gregs[REG_RCX];
+  unw_ctx.uc_mcontext.gregs[REG_RBX] = uc->uc_mcontext.gregs[REG_RBX];
+  unw_ctx.uc_mcontext.gregs[REG_RSI] = uc->uc_mcontext.gregs[REG_RSI];
+  unw_ctx.uc_mcontext.gregs[REG_RDI] = uc->uc_mcontext.gregs[REG_RDI];
+  unw_ctx.uc_mcontext.gregs[REG_RBP] = uc->uc_mcontext.gregs[REG_RBP];
+  unw_ctx.uc_mcontext.gregs[REG_RSP] = uc->uc_mcontext.gregs[REG_RSP];
+  unw_ctx.uc_mcontext.gregs[REG_R8] = uc->uc_mcontext.gregs[REG_R8];
+  unw_ctx.uc_mcontext.gregs[REG_R9] = uc->uc_mcontext.gregs[REG_R9];
+  unw_ctx.uc_mcontext.gregs[REG_R10] = uc->uc_mcontext.gregs[REG_R10];
+  unw_ctx.uc_mcontext.gregs[REG_R11] = uc->uc_mcontext.gregs[REG_R11];
+  unw_ctx.uc_mcontext.gregs[REG_R12] = uc->uc_mcontext.gregs[REG_R12];
+  unw_ctx.uc_mcontext.gregs[REG_R13] = uc->uc_mcontext.gregs[REG_R13];
+  unw_ctx.uc_mcontext.gregs[REG_R14] = uc->uc_mcontext.gregs[REG_R14];
+  unw_ctx.uc_mcontext.gregs[REG_R15] = uc->uc_mcontext.gregs[REG_R15];
+  unw_ctx.uc_mcontext.gregs[REG_RIP] = uc->uc_mcontext.gregs[REG_RIP];
+  unw_ctx.uc_mcontext.gregs[REG_RSP] = uc->uc_mcontext.gregs[REG_RSP];
+  unw_ctx.uc_mcontext.fpregs = &uc->__fpregs_mem;
+
+  unw_ctx.uc_stack.ss_sp = uc->uc_stack.ss_sp;
+  unw_ctx.uc_stack.ss_size = uc->uc_stack.ss_size;
+  unw_ctx.uc_stack.ss_flags = uc->uc_stack.ss_flags;
+
+  unw_ctx.uc_link = uc->uc_link;
+  unw_ctx.uc_sigmask = uc->uc_sigmask;
+
+  return unw_ctx;
+}
+
 } // namespace wanco

lib-rt/chkpt/chkpt.h

@@ -80,6 +80,7 @@ class Checkpoint {
   }
 
   void prepare_restore() {
+    Info() << "Loading value stack" << '\n';
     restore_stack.clear();
     for (auto &frame : frames) {
       for (auto &value : frame.stack) {
@@ -92,6 +93,7 @@ class Checkpoint {
 std::pair<wanco::Checkpoint, int8_t *>
 decode_checkpoint_proto(std::ifstream &f);
 
-void encode_checkpoint_proto(std::ofstream &ofs, Checkpoint &chkpt, int8_t *memory_base);
+void encode_checkpoint_proto(std::ofstream &ofs, Checkpoint &chkpt,
+                             int8_t *memory_base);
 
 } // namespace wanco

lib-rt/osr/asr_exit.cc

@@ -24,6 +24,8 @@ using stackmap_table =
 // func => stackmap_record[]
 static stackmap_table populate_stackmap(const stackmap::Stackmap &stackmap) {
   stackmap_table map;
+  Info() << "Populating stackmaps: " << std::dec
+         << stackmap.stkmap_records.size() << " records" << '\n';
 
   for (auto &record : stackmap.stkmap_records) {
     auto id = record->patchpoint_id;
@@ -38,6 +40,8 @@ static stackmap_table populate_stackmap(const stackmap::Stackmap &stackmap) {
       auto &ent = it->second;
       ent.push_back(record);
     }
+    Debug() << "Found stackmap record for func_" << std::dec << loc.get_func()
+            << ", pc_offset=" << std::dec << record->instruction_offset << '\n';
   }
 
   // sort all entries by the first element
@@ -72,6 +76,7 @@ lookup_stackmap(stackmap_table &map, int32_t func_index, int32_t pc_offset) {
       Debug() << "Instead, found a record for pc_offset=" << std::dec
               << record->instruction_offset << '\n';
     }
+    Warn() << "Assuming the record is empty" << '\n';
     return std::nullopt;
   }
 

lib-rt/stackmap/stackmap.cc

@@ -110,6 +110,8 @@ auto parse_stackmap(const std::span<const uint8_t> data) -> Stackmap {
   uint32_t const num_functions = parse_u32(ptr);
   uint32_t const num_constants = parse_u32(ptr);
   uint32_t const num_records = parse_u32(ptr);
+  Info() << "Expects " << num_functions << " functions, " << num_constants
+         << " constants, and " << num_records << " records" << '\n';
 
   std::vector<StkSizeRecord> stksize_records;
   stksize_records.reserve(num_functions);
@@ -130,6 +132,11 @@ auto parse_stackmap(const std::span<const uint8_t> data) -> Stackmap {
         std::make_shared<StkMapRecord>(parse_stk_map_record(ptr)));
   }
 
+  ASSERT(num_functions == stksize_records.size() &&
+         "Invalid number of records");
+  ASSERT(num_constants == constants.size() && "Invalid number of constants");
+  ASSERT(num_records == stkmap_records.size() && "Invalid number of records");
+
   return Stackmap{.header = header,
                   .num_functions = num_functions,
                   .num_constants = num_constants,