optimize wasm recover

Link: https://code.alibaba-inc.com/Ingress/proxy-wasm-cpp-host/codereview/14808697
* optimize wasm recover

* update UT

* add some assert

Author: johnlanni

include/proxy-wasm/context.h

@@ -151,7 +151,6 @@ 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

@@ -17,6 +17,7 @@
 
 #include <string.h>
 
+#include <cassert>
 #include <atomic>
 #include <deque>
 #include <map>
@@ -73,11 +74,6 @@ 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_; }
@@ -328,7 +324,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) {}
-  virtual ~WasmHandleBase() {
+  ~WasmHandleBase() {
     if (wasm_base_) {
       wasm_base_->startShutdown();
     }
@@ -341,29 +337,24 @@ 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_; }
+
+  // Recover the wasm vm and generate a new wasm handle
   bool doRecover(std::shared_ptr<WasmHandleBase> &new_handle) {
-    if (!need_recover_ || recover_vm_callback_ == nullptr) {
+    assert(new_handle == nullptr);
+    if (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_;
 };
@@ -385,7 +376,7 @@ 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) {}
-  virtual ~PluginHandleBase() {
+  ~PluginHandleBase() {
     if (wasm_handle_) {
       wasm_handle_->wasm()->startShutdown(plugin_->key());
     }
@@ -395,35 +386,33 @@ class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
   std::shared_ptr<WasmBase> &wasm() { return wasm_handle_->wasm(); }
   std::shared_ptr<WasmHandleBase> &wasmHandle() { return wasm_handle_; }
 
-  void setRecoverPluginCallback(std::function<bool(std::shared_ptr<WasmHandleBase> &)> &&f) {
+  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() {
-    if (!need_recover_ || recover_plugin_callback_ == nullptr) {
+
+  // Recover the wasm plugin and generate a new plugin handle
+  bool doRecover(std::shared_ptr<PluginHandleBase> &new_handle) {
+    assert(new_handle == nullptr);
+    if (recover_plugin_callback_ == nullptr) {
       return true;
     }
     std::shared_ptr<WasmHandleBase> new_wasm_handle;
     if (!wasm_handle_->doRecover(new_wasm_handle)) {
       return false;
     }
-    if (!recover_plugin_callback_(new_wasm_handle)) {
+    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<bool(std::shared_ptr<WasmHandleBase> &)> recover_plugin_callback_;
+  std::function<std::shared_ptr<PluginHandleBase>(std::shared_ptr<WasmHandleBase> &)>
+      recover_plugin_callback_;
 };
 
 using PluginHandleFactory = std::function<std::shared_ptr<PluginHandleBase>(

src/context.cc

@@ -13,6 +13,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include <cassert>
 #include <deque>
 #include <map>
 #include <memory>
@@ -529,7 +530,8 @@ FilterMetadataStatus ContextBase::convertVmCallResultToFilterMetadataStatus(uint
 
 ContextBase::~ContextBase() {
   // Do not remove vm context which has the same lifetime as wasm_.
-  if (id_ != 0U && wasm_ != nullptr) {
+  if (id_ != 0U) {
+    assert(wasm_ != nullptr);
     wasm_->contexts_.erase(id_);
   }
 }

src/wasm.cc

@@ -556,67 +556,52 @@ std::shared_ptr<WasmHandleBase> getThreadLocalWasm(std::string_view vm_key) {
 
 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();
-        }
-      });
+  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);
+    }
+  });
 }
 
 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;
-    }
+  wasm_handle->setRecoverVmCallback(
+      [vm_key, 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 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;
-  });
+        if (!new_handle->wasm()->initialize()) {
+          base_handle->wasm()->fail(FailState::RecoverError,
+                                    "Failed to initialize Wasm code during recover");
+          return nullptr;
+        }
+        local_wasms[vm_key] = new_handle;
+        setWasmFailCallback(vm_key, new_handle);
+        setWasmRecoverCallback(vm_key, new_handle, base_handle, clone_factory);
+        integration->trace("Wasm handle has been recovered");
+        return new_handle;
+      });
 }
 
 static std::shared_ptr<WasmHandleBase>
@@ -651,59 +636,54 @@ getOrCreateThreadLocalWasm(const std::shared_ptr<WasmHandleBase> &base_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();
-        }
-      });
+                           const std::shared_ptr<WasmHandleBase> &wasm_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);
+    }
+  });
 }
 
 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) {
-    const auto &integration = base_handle->wasm()->wasm_vm()->integration();
-    integration->trace("Start recover plugin_handle");
-    // We cannot reuse plugin that other have created because the life cycle of the plugin
-    // handle is controlled by the upper layer.
-    auto plugin_handle = plugin_handle_for_copy.lock();
-    if (!plugin_handle) {
-      base_handle->wasm()->fail(FailState::RecoverError, "Plugin handle lock failed");
-      return false;
-    }
-    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 false;
-    }
-    if (!wasm_handle->wasm()->configure(plugin_context, plugin)) {
-      base_handle->wasm()->fail(FailState::RecoverError,
-                                "Failed to configure thread-local Wasm plugin during recover");
-      return false;
-    }
-    local_plugins[key] = plugin_handle;
-    integration->trace("Plugin_handle has been recovered");
-    setPluginFailCallback(key, wasm_handle, plugin_handle);
-    return true;
-  });
+                              const std::shared_ptr<PluginBase> &plugin,
+                              const PluginHandleFactory &plugin_factory) {
+  plugin_handle->setRecoverPluginCallback(
+      [key, base_handle, plugin, plugin_factory](
+          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) {
+            return plugin_handle;
+          }
+          local_plugins.erase(key);
+        }
+        // Try to recover wasm 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;
+        }
+        auto new_handle = plugin_factory(wasm_handle, plugin);
+        local_plugins[key] = new_handle;
+        setPluginFailCallback(key, wasm_handle);
+        setPluginRecoverCallback(key, new_handle, base_handle, plugin, plugin_factory);
+        integration->trace("Plugin handle has been recovered");
+        return new_handle;
+      });
 }
 
 std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
@@ -738,8 +718,8 @@ std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
   }
   auto plugin_handle = plugin_factory(wasm_handle, plugin);
   local_plugins[key] = plugin_handle;
-  setPluginFailCallback(key, wasm_handle, plugin_handle);
-  setPluginRecoverCallback(key, plugin_handle, base_handle, plugin);
+  setPluginFailCallback(key, wasm_handle);
+  setPluginRecoverCallback(key, plugin_handle, base_handle, plugin, plugin_factory);
   return plugin_handle;
 }