support wasm vm recover

Link: https://code.alibaba-inc.com/Ingress/proxy-wasm-cpp-host/codereview/14096772
* support recover crash vm

* fix

* add log

* optimize

* add compdb repo

* optimize

* fix

* add updateWasmHandle in pluginHandle

* make dorecover be virtual

* remove useless code

* user shared_ptr

* fix

* fix bg

* optimize recover logic

* fix

* fix bug

* optimize

* optimize code

* optimize

* add UT

* optimize code

* set pluginhandle virtual

* optimize the recover logic

Author: johnlanni

.bazelrc

@@ -1,3 +1,4 @@
+build:compdb --build_tag_filters=-nocompdb
 # Pass CC, CXX and PATH from the environment.
 build --action_env=CC
 build --action_env=CXX

bazel/repositories.bzl

@@ -19,6 +19,13 @@ load("@bazel_tools//tools/build_defs/repo:utils.bzl", "maybe")
 def proxy_wasm_cpp_host_repositories():
     # Bazel extensions.
 
+    maybe(
+        http_archive,
+        name = "bazel_compdb",
+        strip_prefix = "bazel-compilation-database-0.5.2",
+        url = "https://github.com/grailbio/bazel-compilation-database/archive/0.5.2.tar.gz",
+    )
+
     maybe(
         http_archive,
         name = "bazel_skylib",

include/proxy-wasm/context.h

@@ -151,6 +151,7 @@ class ContextBase : public RootInterface,
   virtual ~ContextBase();
 
   WasmBase *wasm() const { return wasm_; }
+  void clearWasm() { wasm_ = nullptr; }
   uint32_t id() const { return id_; }
   // The VM Context used for calling "malloc" has an id_ == 0.
   bool isVmContext() const { return id_ == 0; }

include/proxy-wasm/wasm.h

@@ -73,6 +73,11 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
       return it->second;
     return nullptr;
   }
+  void clearWasmInContext() {
+    for (auto &item : contexts_) {
+      item.second->clearWasm();
+    }
+  }
   uint32_t allocContextId();
   bool isFailed() { return failed_ != FailState::Ok; }
   FailState fail_state() { return failed_; }
@@ -323,7 +328,7 @@ using WasmHandleCloneFactory =
 class WasmHandleBase : public std::enable_shared_from_this<WasmHandleBase> {
 public:
   explicit WasmHandleBase(std::shared_ptr<WasmBase> wasm_base) : wasm_base_(wasm_base) {}
-  ~WasmHandleBase() {
+  virtual ~WasmHandleBase() {
     if (wasm_base_) {
       wasm_base_->startShutdown();
     }
@@ -336,8 +341,31 @@ class WasmHandleBase : public std::enable_shared_from_this<WasmHandleBase> {
 
   std::shared_ptr<WasmBase> &wasm() { return wasm_base_; }
 
+  virtual void swap(std::shared_ptr<WasmHandleBase> &new_handle) {
+    wasm_base_.swap(new_handle->wasm_base_);
+  }
+
+  void setRecoverVmCallback(std::function<std::shared_ptr<WasmHandleBase>()> &&f) {
+    recover_vm_callback_ = std::move(f);
+  }
+  void setNeedRecover() { need_recover_ = true; }
+  bool needRecover() { return need_recover_; }
+  bool doRecover(std::shared_ptr<WasmHandleBase> &new_handle) {
+    if (!need_recover_ || recover_vm_callback_ == nullptr) {
+      return true;
+    }
+    new_handle = recover_vm_callback_();
+    if (!new_handle) {
+      return false;
+    }
+    need_recover_ = false;
+    return true;
+  }
+
 protected:
+  bool need_recover_ = false;
   std::shared_ptr<WasmBase> wasm_base_;
+  std::function<std::shared_ptr<WasmHandleBase>()> recover_vm_callback_;
 };
 
 std::string makeVmKey(std::string_view vm_id, std::string_view configuration,
@@ -357,18 +385,48 @@ class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
   explicit PluginHandleBase(std::shared_ptr<WasmHandleBase> wasm_handle,
                             std::shared_ptr<PluginBase> plugin)
       : plugin_(plugin), wasm_handle_(wasm_handle) {}
-  ~PluginHandleBase() {
+  virtual ~PluginHandleBase() {
     if (wasm_handle_) {
       wasm_handle_->wasm()->startShutdown(plugin_->key());
     }
   }
 
   std::shared_ptr<PluginBase> &plugin() { return plugin_; }
   std::shared_ptr<WasmBase> &wasm() { return wasm_handle_->wasm(); }
+  std::shared_ptr<WasmHandleBase> &wasmHandle() { return wasm_handle_; }
+
+  void setRecoverPluginCallback(
+      std::function<std::shared_ptr<PluginHandleBase>(std::shared_ptr<WasmHandleBase> &)> &&f) {
+    recover_plugin_callback_ = std::move(f);
+  }
+  void setNeedRecover() { need_recover_ = true; }
+  bool needRecover() { return need_recover_; }
+  bool doRecover(std::shared_ptr<PluginHandleBase> &new_handle) {
+    if (!need_recover_ || recover_plugin_callback_ == nullptr) {
+      return true;
+    }
+    std::shared_ptr<WasmHandleBase> new_wasm_handle;
+    if (!wasm_handle_->doRecover(new_wasm_handle)) {
+      return false;
+    }
+    new_handle = recover_plugin_callback_(new_wasm_handle);
+    if (!new_handle) {
+      return false;
+    }
+    need_recover_ = false;
+    return true;
+  }
+
+  virtual void updateWasm(std::shared_ptr<WasmHandleBase> &new_handle) {
+    wasm_handle_ = new_handle;
+  }
 
 protected:
+  bool need_recover_ = false;
   std::shared_ptr<PluginBase> plugin_;
   std::shared_ptr<WasmHandleBase> wasm_handle_;
+  std::function<std::shared_ptr<PluginHandleBase>(std::shared_ptr<WasmHandleBase> &)>
+      recover_plugin_callback_;
 };
 
 using PluginHandleFactory = std::function<std::shared_ptr<PluginHandleBase>(

include/proxy-wasm/wasm_vm.h

@@ -174,6 +174,7 @@ enum class FailState : int {
   StartFailed = 5,
   ConfigureFailed = 6,
   RuntimeError = 7,
+  RecoverError = 8,
 };
 
 // Wasm VM instance. Provides the low level WASM interface.

src/context.cc

@@ -529,7 +529,7 @@ FilterMetadataStatus ContextBase::convertVmCallResultToFilterMetadataStatus(uint
 
 ContextBase::~ContextBase() {
   // Do not remove vm context which has the same lifetime as wasm_.
-  if (id_ != 0U) {
+  if (id_ != 0U && wasm_ != nullptr) {
     wasm_->contexts_.erase(id_);
   }
 }

src/wasm.cc

@@ -554,6 +554,71 @@ std::shared_ptr<WasmHandleBase> getThreadLocalWasm(std::string_view vm_key) {
   return wasm;
 }
 
+void setWasmFailCallback(const std::string &vm_key,
+                         const std::shared_ptr<WasmHandleBase> &wasm_handle) {
+  std::weak_ptr<WasmHandleBase> wasm_handle_for_copy = wasm_handle;
+  wasm_handle->wasm()->wasm_vm()->addFailCallback(
+      [vm_key, wasm_handle_for_copy](proxy_wasm::FailState fail_state) {
+        if (fail_state == proxy_wasm::FailState::RuntimeError) {
+          // If VM failed, erase the entry so that:
+          // 1) we can recreate the new thread local VM from the same base_wasm.
+          // 2) we wouldn't reuse the failed VM for new plugins accidentally.
+          local_wasms.erase(vm_key);
+          auto wasm_handle = wasm_handle_for_copy.lock();
+          if (!wasm_handle) {
+            return;
+          }
+          wasm_handle->setNeedRecover();
+        }
+      });
+}
+
+void setWasmRecoverCallback(const std::string &vm_key,
+                            const std::shared_ptr<WasmHandleBase> &wasm_handle,
+                            const std::shared_ptr<WasmHandleBase> &base_handle,
+                            const WasmHandleCloneFactory &clone_factory) {
+  std::weak_ptr<WasmHandleBase> wasm_handle_for_copy = wasm_handle;
+  wasm_handle->setRecoverVmCallback([vm_key, wasm_handle_for_copy, base_handle,
+                                     clone_factory]() -> std::shared_ptr<WasmHandleBase> {
+    const auto &integration = base_handle->wasm()->wasm_vm()->integration();
+    integration->trace("Start recover wasm_handle");
+    auto it = local_wasms.find(vm_key);
+    if (it != local_wasms.end()) {
+      auto wasm_handle = it->second.lock();
+      if (wasm_handle) {
+        integration->trace("Wasm handle already exists");
+        return wasm_handle;
+      }
+      local_wasms.erase(vm_key);
+    }
+    // try to recover wasm vm
+    auto wasm_handle = wasm_handle_for_copy.lock();
+    if (!wasm_handle) {
+      base_handle->wasm()->fail(FailState::RecoverError, "Wasm handle lock failed");
+      return nullptr;
+    }
+    auto new_handle = clone_factory(base_handle);
+    if (!new_handle) {
+      base_handle->wasm()->fail(FailState::RecoverError,
+                                "Failed to clone Base Wasm during recover");
+      return nullptr;
+    }
+
+    if (!new_handle->wasm()->initialize()) {
+      base_handle->wasm()->fail(FailState::RecoverError,
+                                "Failed to initialize Wasm code during recover");
+      return nullptr;
+    }
+    // avoid the context use the stale wasm ptr
+    wasm_handle->wasm()->clearWasmInContext();
+    wasm_handle->swap(new_handle);
+    local_wasms[vm_key] = wasm_handle;
+    integration->trace("Wasm handle has been recovered");
+    setWasmFailCallback(vm_key, wasm_handle);
+    return wasm_handle;
+  });
+}
+
 static std::shared_ptr<WasmHandleBase>
 getOrCreateThreadLocalWasm(const std::shared_ptr<WasmHandleBase> &base_handle,
                            const WasmHandleCloneFactory &clone_factory) {
@@ -580,17 +645,75 @@ getOrCreateThreadLocalWasm(const std::shared_ptr<WasmHandleBase> &base_handle,
     return nullptr;
   }
   local_wasms[vm_key] = wasm_handle;
-  wasm_handle->wasm()->wasm_vm()->addFailCallback([vm_key](proxy_wasm::FailState fail_state) {
-    if (fail_state == proxy_wasm::FailState::RuntimeError) {
-      // If VM failed, erase the entry so that:
-      // 1) we can recreate the new thread local VM from the same base_wasm.
-      // 2) we wouldn't reuse the failed VM for new plugins accidentally.
-      local_wasms.erase(vm_key);
-    };
-  });
+  setWasmFailCallback(vm_key, wasm_handle);
+  setWasmRecoverCallback(vm_key, wasm_handle, base_handle, clone_factory);
   return wasm_handle;
 }
 
+void setPluginFailCallback(const std::string &key,
+                           const std::shared_ptr<WasmHandleBase> &wasm_handle,
+                           const std::shared_ptr<PluginHandleBase> &plugin_handle) {
+  std::weak_ptr<PluginHandleBase> plugin_handle_for_copy = plugin_handle;
+  wasm_handle->wasm()->wasm_vm()->addFailCallback(
+      [key, plugin_handle_for_copy](proxy_wasm::FailState fail_state) {
+        if (fail_state == proxy_wasm::FailState::RuntimeError) {
+          // If VM failed, erase the entry so that:
+          // 1) we can recreate the new thread local plugin from the same base_wasm.
+          // 2) we wouldn't reuse the failed VM for new plugin configs accidentally.
+          local_plugins.erase(key);
+          auto plugin_handle = plugin_handle_for_copy.lock();
+          if (!plugin_handle) {
+            return;
+          }
+          plugin_handle->setNeedRecover();
+        }
+      });
+}
+
+void setPluginRecoverCallback(const std::string &key,
+                              const std::shared_ptr<PluginHandleBase> &plugin_handle,
+                              const std::shared_ptr<WasmHandleBase> &base_handle,
+                              const std::shared_ptr<PluginBase> &plugin) {
+  std::weak_ptr<PluginHandleBase> plugin_handle_for_copy = plugin_handle;
+  plugin_handle->setRecoverPluginCallback(
+      [key, plugin_handle_for_copy, base_handle,
+       plugin](std::shared_ptr<WasmHandleBase> &wasm_handle) -> std::shared_ptr<PluginHandleBase> {
+        const auto &integration = base_handle->wasm()->wasm_vm()->integration();
+        integration->trace("Start recover plugin_handle");
+        auto it = local_plugins.find(key);
+        if (it != local_plugins.end()) {
+          auto plugin_handle = it->second.lock();
+          if (plugin_handle) {
+            integration->trace("Plugin handle already exists");
+            return plugin_handle;
+          }
+          local_plugins.erase(key);
+        }
+        auto plugin_handle = plugin_handle_for_copy.lock();
+        if (!plugin_handle) {
+          base_handle->wasm()->fail(FailState::RecoverError, "Plugin handle lock failed");
+          return nullptr;
+        }
+        plugin_handle->updateWasm(wasm_handle);
+        // Create and initialize new thread-local Plugin.
+        auto *plugin_context = wasm_handle->wasm()->start(plugin);
+        if (plugin_context == nullptr) {
+          base_handle->wasm()->fail(FailState::RecoverError,
+                                    "Failed to start thread-local Wasm during recover");
+          return nullptr;
+        }
+        if (!wasm_handle->wasm()->configure(plugin_context, plugin)) {
+          base_handle->wasm()->fail(FailState::RecoverError,
+                                    "Failed to configure thread-local Wasm plugin during recover");
+          return nullptr;
+        }
+        local_plugins[key] = plugin_handle;
+        integration->trace("Plugin_handle has been recovered");
+        setPluginFailCallback(key, wasm_handle, plugin_handle);
+        return plugin_handle;
+      });
+}
+
 std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
     const std::shared_ptr<WasmHandleBase> &base_handle, const std::shared_ptr<PluginBase> &plugin,
     const WasmHandleCloneFactory &clone_factory, const PluginHandleFactory &plugin_factory) {
@@ -623,14 +746,8 @@ std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
   }
   auto plugin_handle = plugin_factory(wasm_handle, plugin);
   local_plugins[key] = plugin_handle;
-  wasm_handle->wasm()->wasm_vm()->addFailCallback([key](proxy_wasm::FailState fail_state) {
-    if (fail_state == proxy_wasm::FailState::RuntimeError) {
-      // If VM failed, erase the entry so that:
-      // 1) we can recreate the new thread local plugin from the same base_wasm.
-      // 2) we wouldn't reuse the failed VM for new plugin configs accidentally.
-      local_plugins.erase(key);
-    };
-  });
+  setPluginFailCallback(key, wasm_handle, plugin_handle);
+  setPluginRecoverCallback(key, plugin_handle, base_handle, plugin);
   return plugin_handle;
 }