Pass PluginHandle to stream contexts.

include/proxy-wasm/context.h

@@ -32,6 +32,7 @@ namespace proxy_wasm {
 #include "proxy_wasm_common.h"
 #include "proxy_wasm_enums.h"
 
+class PluginHandleBase;
 class WasmBase;
 class WasmVm;
 
@@ -143,7 +144,7 @@ class ContextBase : public RootInterface,
   ContextBase(WasmBase *wasm);                                     // Vm Context.
   ContextBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin); // Root Context.
   ContextBase(WasmBase *wasm, uint32_t parent_context_id,
-              std::shared_ptr<PluginBase> plugin); // Stream context.
+              std::shared_ptr<PluginHandleBase> plugin_handle); // Stream context.
   virtual ~ContextBase();
 
   WasmBase *wasm() const { return wasm_; }
@@ -394,12 +395,13 @@ class ContextBase : public RootInterface,
 
   WasmBase *wasm_{nullptr};
   uint32_t id_{0};
-  uint32_t parent_context_id_{0};           // 0 for roots and the general context.
-  ContextBase *parent_context_{nullptr};    // set in all contexts.
-  std::string root_id_;                     // set only in root context.
-  std::string root_log_prefix_;             // set only in root context.
-  std::shared_ptr<PluginBase> plugin_;      // set in root and stream contexts.
-  std::shared_ptr<PluginBase> temp_plugin_; // Remove once ABI v0.1.0 is gone.
+  uint32_t parent_context_id_{0};                   // 0 for roots and the general context.
+  ContextBase *parent_context_{nullptr};            // set in all contexts.
+  std::string root_id_;                             // set only in root context.
+  std::string root_log_prefix_;                     // set only in root context.
+  std::shared_ptr<PluginBase> plugin_;              // set in root and stream contexts.
+  std::shared_ptr<PluginHandleBase> plugin_handle_; // set only in stream context.
+  std::shared_ptr<PluginBase> temp_plugin_;         // Remove once ABI v0.1.0 is gone.
   bool in_vm_context_created_ = false;
   bool destroyed_ = false;
   bool stream_failed_ = false; // Set true after failStream is called in case of VM failure.

include/proxy-wasm/wasm.h

@@ -321,14 +321,19 @@ class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
 public:
   explicit PluginHandleBase(std::shared_ptr<WasmHandleBase> wasm_handle,
                             std::shared_ptr<PluginBase> plugin)
-      : wasm_handle_(wasm_handle), plugin_key_(plugin->key()) {}
-  ~PluginHandleBase() { wasm_handle_->wasm()->startShutdown(plugin_key_); }
+      : plugin_(plugin), wasm_handle_(wasm_handle) {}
+  ~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(); }
 
 protected:
+  std::shared_ptr<PluginBase> plugin_;
   std::shared_ptr<WasmHandleBase> wasm_handle_;
-  std::string plugin_key_;
 };
 
 using PluginHandleFactory = std::function<std::shared_ptr<PluginHandleBase>(

src/context.cc

@@ -98,9 +98,9 @@ ContextBase::ContextBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin)
 
 // NB: wasm can be nullptr if it failed to be created successfully.
 ContextBase::ContextBase(WasmBase *wasm, uint32_t parent_context_id,
-                         std::shared_ptr<PluginBase> plugin)
+                         std::shared_ptr<PluginHandleBase> plugin_handle)
     : wasm_(wasm), id_(wasm ? wasm->allocContextId() : 0), parent_context_id_(parent_context_id),
-      plugin_(plugin) {
+      plugin_(plugin_handle->plugin()), plugin_handle_(plugin_handle) {
   if (wasm_) {
     wasm_->contexts_[id_] = this;
     parent_context_ = wasm_->contexts_[parent_context_id_];