Add addFailCallback to WasmBase.

Signed-off-by: Takeshi Yoneda <[email protected]>

Author: mathetake

include/proxy-wasm/wasm_vm.h

@@ -20,6 +20,7 @@
 #include <optional>
 #include <string>
 #include <unordered_map>
+#include <vector>
 
 #include "include/proxy-wasm/word.h"
 
@@ -268,12 +269,12 @@ class WasmVm {
   void fail(FailState fail_state, std::string_view message) {
     integration()->error(message);
     failed_ = fail_state;
-    if (fail_callback_) {
-      fail_callback_(fail_state);
+    for (auto &callback : fail_callbacks_) {
+      callback(fail_state);
     }
   }
-  void setFailCallback(std::function<void(FailState)> fail_callback) {
-    fail_callback_ = fail_callback;
+  void addFailCallback(std::function<void(FailState)> fail_callback) {
+    fail_callbacks_.push_back(fail_callback);
   }
 
   // Integrator operations.
@@ -283,7 +284,7 @@ class WasmVm {
 protected:
   std::unique_ptr<WasmVmIntegration> integration_;
   FailState failed_ = FailState::Ok;
-  std::function<void(FailState)> fail_callback_;
+  std::vector<std::function<void(FailState)>> fail_callbacks_;
 };
 
 // Thread local state set during a call into a WASM VM so that calls coming out of the

src/wasm.cc

@@ -210,7 +210,7 @@ WasmBase::WasmBase(const std::shared_ptr<WasmHandleBase> &base_wasm_handle, Wasm
   if (!wasm_vm_) {
     failed_ = FailState::UnableToCreateVm;
   } else {
-    wasm_vm_->setFailCallback([this](FailState fail_state) { failed_ = fail_state; });
+    wasm_vm_->addFailCallback([this](FailState fail_state) { failed_ = fail_state; });
   }
 }
 
@@ -224,7 +224,7 @@ WasmBase::WasmBase(std::unique_ptr<WasmVm> wasm_vm, std::string_view vm_id,
   if (!wasm_vm_) {
     failed_ = FailState::UnableToCreateVm;
   } else {
-    wasm_vm_->setFailCallback([this](FailState fail_state) { failed_ = fail_state; });
+    wasm_vm_->addFailCallback([this](FailState fail_state) { failed_ = fail_state; });
   }
 }
 
@@ -561,6 +561,14 @@ getOrCreateThreadLocalWasm(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);
+    };
+  });
   return wasm_handle;
 }
 
@@ -596,6 +604,14 @@ 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);
+    };
+  });
   return plugin_handle;
 }