Delete failed Plugin/VMs via fail callbacks. (#181)

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"
 
@@ -297,12 +298,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.
@@ -312,7 +313,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

@@ -208,7 +208,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; });
   }
 }
 
@@ -222,7 +222,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; });
   }
 }
 
@@ -559,6 +559,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;
 }
 
@@ -594,6 +602,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;
 }
 

test/BUILD

@@ -78,9 +78,13 @@ cc_test(
 )
 
 cc_test(
-    name = "context_test",
-    srcs = ["context_test.cc"],
+    name = "wasm_test",
+    srcs = ["wasm_test.cc"],
+    data = [
+        "//test/test_data:abi_export.wasm",
+    ],
     deps = [
+        ":utility_lib",
         "//:lib",
         "@com_google_googletest//:gtest",
         "@com_google_googletest//:gtest_main",

test/bytecode_util_test.cc

@@ -59,7 +59,14 @@ TEST(TestBytecodeUtil, getFunctionNameIndex) {
   // OK.
   EXPECT_TRUE(BytecodeUtil::getFunctionNameIndex(source, actual));
   EXPECT_FALSE(actual.empty());
-  EXPECT_EQ(actual.find(0)->second, "proxy_abi_version_0_2_0");
+  bool abi_version_found = false;
+  for (auto it : actual) {
+    if (it.second == "proxy_abi_version_0_2_0") {
+      abi_version_found = true;
+      break;
+    }
+  }
+  EXPECT_TRUE(abi_version_found);
 
   // Fail due to the corrupted bytecode.
   // TODO(@mathetake): here we haven't covered all the parsing failure branches. Add more cases.

test/context_test.cc

@@ -49,10 +49,10 @@ class TestContext : public ContextBase {
 
 TEST_P(TestVM, Environment) {
   std::unordered_map<std::string, std::string> envs = {{"KEY1", "VALUE1"}, {"KEY2", "VALUE2"}};
-  initialize("env.wasm");
+  auto source = readTestWasmFile("env.wasm");
 
   auto wasm_base = WasmBase(std::move(vm_), "vm_id", "", "", envs, {});
-  ASSERT_TRUE(wasm_base.wasm_vm()->load(source_, {}, {}));
+  ASSERT_TRUE(wasm_base.wasm_vm()->load(source, {}, {}));
 
   TestContext context(&wasm_base);
   current_context_ = &context;
@@ -72,9 +72,9 @@ TEST_P(TestVM, Environment) {
 }
 
 TEST_P(TestVM, WithoutEnvironment) {
-  initialize("env.wasm");
+  auto source = readTestWasmFile("env.wasm");
   auto wasm_base = WasmBase(std::move(vm_), "vm_id", "", "", {}, {});
-  ASSERT_TRUE(wasm_base.wasm_vm()->load(source_, {}, {}));
+  ASSERT_TRUE(wasm_base.wasm_vm()->load(source, {}, {}));
 
   TestContext context(&wasm_base);
   current_context_ = &context;
@@ -92,9 +92,9 @@ TEST_P(TestVM, WithoutEnvironment) {
 }
 
 TEST_P(TestVM, Clock) {
-  initialize("clock.wasm");
+  auto source = readTestWasmFile("clock.wasm");
   auto wasm_base = WasmBase(std::move(vm_), "vm_id", "", "", {}, {});
-  ASSERT_TRUE(wasm_base.wasm_vm()->load(source_, {}, {}));
+  ASSERT_TRUE(wasm_base.wasm_vm()->load(source, {}, {}));
 
   TestContext context(&wasm_base);
   current_context_ = &context;

test/exports_test.cc

@@ -47,8 +47,8 @@ TEST_P(TestVM, Basic) {
 }
 
 TEST_P(TestVM, Memory) {
-  initialize("abi_export.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+  auto source = readTestWasmFile("abi_export.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
   ASSERT_TRUE(vm_->link(""));
 
   Word word;
@@ -68,8 +68,8 @@ TEST_P(TestVM, Clone) {
   if (vm_->cloneable() == proxy_wasm::Cloneable::NotCloneable) {
     return;
   }
-  initialize("abi_export.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+  auto source = readTestWasmFile("abi_export.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
   ASSERT_TRUE(vm_->link(""));
   const auto address = 0x2000;
   Word word;
@@ -113,8 +113,10 @@ TEST_P(TestVM, StraceLogLevel) {
     // See https://github.com/proxy-wasm/proxy-wasm-cpp-host/issues/120.
     return;
   }
-  initialize("callback.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+
+  auto integration = static_cast<DummyIntegration *>(vm_->integration().get());
+  auto source = readTestWasmFile("callback.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
   vm_->registerCallback("env", "callback", &nopCallback,
                         &ConvertFunctionWordToUint32<decltype(nopCallback),
                                                      nopCallback>::convertFunctionWordToUint32);
@@ -128,16 +130,16 @@ TEST_P(TestVM, StraceLogLevel) {
 
   run(nullptr);
   // no trace message found since DummyIntegration's log_level_ defaults to  LogLevel::info
-  EXPECT_EQ(integration_->trace_message_, "");
+  EXPECT_EQ(integration->trace_message_, "");
 
-  integration_->log_level_ = LogLevel::trace;
+  integration->log_level_ = LogLevel::trace;
   run(nullptr);
-  EXPECT_NE(integration_->trace_message_, "");
+  EXPECT_NE(integration->trace_message_, "");
 }
 
 TEST_P(TestVM, Callback) {
-  initialize("callback.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+  auto source = readTestWasmFile("callback.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
 
   TestContext context;
 
@@ -166,16 +168,17 @@ TEST_P(TestVM, Callback) {
 }
 
 TEST_P(TestVM, Trap) {
-  initialize("trap.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+  auto source = readTestWasmFile("trap.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
   ASSERT_TRUE(vm_->link(""));
   TestContext context;
   WasmCallVoid<0> trigger;
   vm_->getFunction("trigger", &trigger);
   EXPECT_TRUE(trigger != nullptr);
   trigger(&context);
   std::string exp_message = "Function: trigger failed";
-  ASSERT_TRUE(integration_->error_message_.find(exp_message) != std::string::npos);
+  auto integration = static_cast<DummyIntegration *>(vm_->integration().get());
+  ASSERT_TRUE(integration->error_message_.find(exp_message) != std::string::npos);
 }
 
 TEST_P(TestVM, Trap2) {
@@ -185,16 +188,17 @@ TEST_P(TestVM, Trap2) {
     // WAVM::Runtime::describeCallStack. Needs further investigation.
     return;
   }
-  initialize("trap.wasm");
-  ASSERT_TRUE(vm_->load(source_, {}, {}));
+  auto source = readTestWasmFile("trap.wasm");
+  ASSERT_TRUE(vm_->load(source, {}, {}));
   ASSERT_TRUE(vm_->link(""));
   TestContext context;
   WasmCallWord<1> trigger2;
   vm_->getFunction("trigger2", &trigger2);
   EXPECT_TRUE(trigger2 != nullptr);
   trigger2(&context, 0);
   std::string exp_message = "Function: trigger2 failed";
-  ASSERT_TRUE(integration_->error_message_.find(exp_message) != std::string::npos);
+  auto integration = static_cast<DummyIntegration *>(vm_->integration().get());
+  ASSERT_TRUE(integration->error_message_.find(exp_message) != std::string::npos);
 }
 
 } // namespace

test/runtime_test.cc

@@ -14,3 +14,21 @@
 
 #[no_mangle]
 pub extern "C" fn proxy_abi_version_0_2_0() {}
+
+#[no_mangle]
+pub extern "C" fn proxy_on_vm_start(_: u32, _: usize) -> bool {
+    true
+}
+
+#[no_mangle]
+pub extern "C" fn proxy_on_context_create(_: u32, _: u32) {}
+
+#[no_mangle]
+pub extern "C" fn proxy_on_memory_allocate(size: usize) -> *mut u8 {
+    let mut vec: Vec<u8> = Vec::with_capacity(size);
+    unsafe {
+        vec.set_len(size);
+    }
+    let slice = vec.into_boxed_slice();
+    Box::into_raw(slice) as *mut u8
+}

test/test_data/abi_export.rs

@@ -67,37 +67,39 @@ class TestVM : public testing::TestWithParam<std::string> {
 public:
   std::unique_ptr<proxy_wasm::WasmVm> vm_;
 
-  TestVM() : integration_(new DummyIntegration{}) {
+  TestVM() {
     runtime_ = GetParam();
+    vm_ = newVm();
+  }
+
+  std::unique_ptr<proxy_wasm::WasmVm> newVm() {
+    std::unique_ptr<proxy_wasm::WasmVm> vm;
     if (runtime_ == "") {
       EXPECT_TRUE(false) << "runtime must not be empty";
 #if defined(PROXY_WASM_HAS_RUNTIME_V8)
     } else if (runtime_ == "v8") {
-      vm_ = proxy_wasm::createV8Vm();
+      vm = proxy_wasm::createV8Vm();
 #endif
 #if defined(PROXY_WASM_HAS_RUNTIME_WAVM)
     } else if (runtime_ == "wavm") {
-      vm_ = proxy_wasm::createWavmVm();
+      vm = proxy_wasm::createWavmVm();
 #endif
 #if defined(PROXY_WASM_HAS_RUNTIME_WASMTIME)
     } else if (runtime_ == "wasmtime") {
-      vm_ = proxy_wasm::createWasmtimeVm();
+      vm = proxy_wasm::createWasmtimeVm();
 #endif
 #if defined(PROXY_WASM_HAS_RUNTIME_WAMR)
     } else if (runtime_ == "wamr") {
-      vm_ = proxy_wasm::createWamrVm();
+      vm = proxy_wasm::createWamrVm();
 #endif
     } else {
       EXPECT_TRUE(false) << "compiled without support for the requested \"" << runtime_
                          << "\" runtime";
     }
-    vm_->integration().reset(integration_);
-  }
-
-  void initialize(std::string filename) { source_ = readTestWasmFile(filename); }
+    vm->integration().reset(new DummyIntegration{});
+    return vm;
+  };
 
-  DummyIntegration *integration_;
-  std::string source_;
   std::string runtime_;
 };
 } // namespace proxy_wasm