From d5a30284e3eca4b982675ab2c8220e80d46f725b Mon Sep 17 00:00:00 2001
From: jingze <daijingze.djz@alibaba-inc.com>
Date: Fri, 25 Jul 2025 15:38:15 +0800
Subject: [PATCH] fix wasm memory leak

---
 include/proxy-wasm/wasm.h    |  7 +++++++
 include/proxy-wasm/wasm_vm.h | 23 +++++++++++++++++++----
 src/wasm.cc                  |  3 ++-
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/include/proxy-wasm/wasm.h b/include/proxy-wasm/wasm.h
index c93855c8..9aad384b 100644
--- a/include/proxy-wasm/wasm.h
+++ b/include/proxy-wasm/wasm.h
@@ -399,6 +399,7 @@ class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
   ~PluginHandleBase() {
     if (wasm_handle_) {
       wasm_handle_->wasm()->startShutdown(plugin_->key());
+      wasm_handle_->wasm()->wasm_vm()->removeFailCallback(plugin_handle_key_);
     }
   }
 
@@ -406,6 +407,8 @@ 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 setPluginHandleKey(std::string_view key) { plugin_handle_key_ = std::string(key); }
+
   void setRecoverPluginCallback(
       std::function<std::shared_ptr<PluginHandleBase>(std::shared_ptr<WasmHandleBase> &)> &&f) {
     recover_plugin_callback_ = std::move(f);
@@ -433,6 +436,10 @@ class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
   std::shared_ptr<WasmHandleBase> wasm_handle_;
   std::function<std::shared_ptr<PluginHandleBase>(std::shared_ptr<WasmHandleBase> &)>
       recover_plugin_callback_;
+
+private:
+  // key for the plugin handle, used to identify the key in fail callbacks
+  std::string plugin_handle_key_;
 };
 
 using PluginHandleFactory = std::function<std::shared_ptr<PluginHandleBase>(
diff --git a/include/proxy-wasm/wasm_vm.h b/include/proxy-wasm/wasm_vm.h
index 846717ff..bc44e7f6 100644
--- a/include/proxy-wasm/wasm_vm.h
+++ b/include/proxy-wasm/wasm_vm.h
@@ -328,12 +328,27 @@ class WasmVm {
   void fail(FailState fail_state, std::string_view message) {
     integration()->error(message);
     failed_ = fail_state;
-    for (auto &callback : fail_callbacks_) {
+    for (auto & [key, callback] : fail_callbacks_) {
       callback(fail_state);
     }
   }
-  void addFailCallback(std::function<void(FailState)> fail_callback) {
-    fail_callbacks_.push_back(fail_callback);
+
+  /**
+   * Generates id for fail callbacks allowing direct insertion of the function.
+   * Note: if fail callback needs to be removed later, must provide specific key.
+   */
+   void addFailCallback(std::function<void(FailState)> fail_callback) {
+    static int id = 0;
+    std::string key = std::to_string(id++);
+    addFailCallback(key, std::move(fail_callback));
+  }
+
+  void addFailCallback(const std::string& key, std::function<void(FailState)> fail_callback) {
+    fail_callbacks_[key] = std::move(fail_callback);
+  }
+
+  void removeFailCallback(const std::string& key) {
+    fail_callbacks_.erase(key);
   }
 
   bool isHostFunctionAllowed(const std::string &name) {
@@ -353,7 +368,7 @@ class WasmVm {
 protected:
   std::unique_ptr<WasmVmIntegration> integration_;
   FailState failed_ = FailState::Ok;
-  std::vector<std::function<void(FailState)>> fail_callbacks_;
+  std::unordered_map<std::string, std::function<void(FailState)>> fail_callbacks_;
 
 private:
   bool restricted_callback_{false};
diff --git a/src/wasm.cc b/src/wasm.cc
index 080663c7..44dc988e 100644
--- a/src/wasm.cc
+++ b/src/wasm.cc
@@ -729,7 +729,7 @@ getOrCreateThreadLocalWasm(const std::shared_ptr<WasmHandleBase> &base_handle,
 
 void setPluginFailCallback(const std::string &key,
                            const std::shared_ptr<WasmHandleBase> &wasm_handle) {
-  wasm_handle->wasm()->wasm_vm()->addFailCallback([key](proxy_wasm::FailState fail_state) {
+  wasm_handle->wasm()->wasm_vm()->addFailCallback(key, [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.
@@ -819,6 +819,7 @@ std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
   }
   auto plugin_handle = plugin_factory(wasm_handle, plugin);
   cacheLocalPlugin(key, plugin_handle);
+  plugin_handle->setPluginHandleKey(key);
   setPluginFailCallback(key, wasm_handle);
   setPluginRecoverCallback(key, plugin_handle, base_handle, plugin, plugin_factory);
   return plugin_handle;