Allow execution of multiple instances of the same plugin. (#92)

Signed-off-by: Piotr Sikora <piotrsikora@google.com>

Author: PiotrSikora

include/proxy-wasm/context.h

@@ -50,20 +50,23 @@ struct PluginBase {
              std::string_view runtime, std::string_view plugin_configuration, bool fail_open)
       : name_(std::string(name)), root_id_(std::string(root_id)), vm_id_(std::string(vm_id)),
         runtime_(std::string(runtime)), plugin_configuration_(plugin_configuration),
-        fail_open_(fail_open) {}
+        fail_open_(fail_open), key_(root_id_ + "||" + plugin_configuration_) {}
 
   const std::string name_;
   const std::string root_id_;
   const std::string vm_id_;
   const std::string runtime_;
-  std::string plugin_configuration_;
+  const std::string plugin_configuration_;
   const bool fail_open_;
+
+  const std::string &key() const { return key_; }
   const std::string &log_prefix() const { return log_prefix_; }
 
 private:
   std::string makeLogPrefix() const;
 
-  std::string log_prefix_;
+  const std::string key_;
+  const std::string log_prefix_;
 };
 
 struct BufferBase : public BufferInterface {
@@ -373,16 +376,16 @@ class ContextBase : public RootInterface,
 protected:
   friend class WasmBase;
 
-  void initializeRootBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin);
   std::string makeRootLogPrefix(std::string_view vm_id) const;
 
   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_;
+  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.
   bool in_vm_context_created_ = false;
   bool destroyed_ = false;
 

include/proxy-wasm/wasm.h

@@ -59,8 +59,7 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   std::string_view vm_key() const { return vm_key_; }
   WasmVm *wasm_vm() const { return wasm_vm_.get(); }
   ContextBase *vm_context() const { return vm_context_.get(); }
-  ContextBase *getRootContext(std::string_view root_id);
-  ContextBase *getOrCreateRootContext(const std::shared_ptr<PluginBase> &plugin);
+  ContextBase *getRootContext(const std::shared_ptr<PluginBase> &plugin, bool allow_closed);
   ContextBase *getContext(uint32_t id) {
     auto it = contexts_.find(id);
     if (it != contexts_.end())
@@ -78,6 +77,7 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   void timerReady(uint32_t root_context_id);
   void queueReady(uint32_t root_context_id, uint32_t token);
 
+  void startShutdown(std::string_view plugin_key);
   void startShutdown();
   WasmResult done(ContextBase *root_context);
   void finishShutdown();
@@ -170,11 +170,12 @@ class WasmBase : public std::enable_shared_from_this<WasmBase> {
   uint32_t next_context_id_ = 1;            // 0 is reserved for the VM context.
   std::shared_ptr<ContextBase> vm_context_; // Context unrelated to any specific root or stream
                                             // (e.g. for global constructors).
-  std::unordered_map<std::string, std::unique_ptr<ContextBase>> root_contexts_;
+  std::unordered_map<std::string, std::unique_ptr<ContextBase>> root_contexts_; // Root contexts.
+  std::unordered_map<std::string, std::unique_ptr<ContextBase>> pending_done_;  // Root contexts.
+  std::unordered_set<std::unique_ptr<ContextBase>> pending_delete_;             // Root contexts.
   std::unordered_map<uint32_t, ContextBase *> contexts_;                 // Contains all contexts.
   std::unordered_map<uint32_t, std::chrono::milliseconds> timer_period_; // per root_id.
   std::unique_ptr<ShutdownHandle> shutdown_handle_;
-  std::unordered_set<ContextBase *> pending_done_; // Root contexts not done during shutdown.
 
   WasmCallVoid<0> _initialize_; /* Emscripten v1.39.17+ */
   WasmCallVoid<0> _start_;      /* Emscripten v1.39.0+ */
@@ -275,11 +276,29 @@ createWasm(std::string vm_key, std::string code, std::shared_ptr<PluginBase> plu
            WasmHandleFactory factory, WasmHandleCloneFactory clone_factory, bool allow_precompiled);
 // Get an existing ThreadLocal VM matching 'vm_id' or nullptr if there isn't one.
 std::shared_ptr<WasmHandleBase> getThreadLocalWasm(std::string_view vm_id);
+
+class PluginHandleBase : public std::enable_shared_from_this<PluginHandleBase> {
+public:
+  explicit PluginHandleBase(std::shared_ptr<WasmHandleBase> wasm_handle,
+                            std::string_view plugin_key)
+      : wasm_handle_(wasm_handle), plugin_key_(plugin_key) {}
+  ~PluginHandleBase() { wasm_handle_->wasm()->startShutdown(plugin_key_); }
+
+  std::shared_ptr<WasmBase> &wasm() { return wasm_handle_->wasm(); }
+
+protected:
+  std::shared_ptr<WasmHandleBase> wasm_handle_;
+  std::string plugin_key_;
+};
+
+using PluginHandleFactory = std::function<std::shared_ptr<PluginHandleBase>(
+    std::shared_ptr<WasmHandleBase> base_wasm, std::string_view plugin_key)>;
+
 // Get an existing ThreadLocal VM matching 'vm_id' or create one using 'base_wavm' by cloning or by
 // using it it as a template.
-std::shared_ptr<WasmHandleBase>
-getOrCreateThreadLocalWasm(std::shared_ptr<WasmHandleBase> base_wasm,
-                           std::shared_ptr<PluginBase> plugin, WasmHandleCloneFactory factory);
+std::shared_ptr<PluginHandleBase> getOrCreateThreadLocalPlugin(
+    std::shared_ptr<WasmHandleBase> base_wasm, std::shared_ptr<PluginBase> plugin,
+    WasmHandleCloneFactory clone_factory, PluginHandleFactory plugin_factory);
 
 // Clear Base Wasm cache and the thread-local Wasm sandbox cache for the calling thread.
 void clearWasmCachesForTesting();

src/context.cc

@@ -272,8 +272,10 @@ ContextBase::ContextBase(WasmBase *wasm) : wasm_(wasm), parent_context_(this) {
   wasm_->contexts_[id_] = this;
 }
 
-ContextBase::ContextBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin) {
-  initializeRootBase(wasm, plugin);
+ContextBase::ContextBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin)
+    : wasm_(wasm), id_(wasm->allocContextId()), parent_context_(this), root_id_(plugin->root_id_),
+      root_log_prefix_(makeRootLogPrefix(plugin->vm_id_)), plugin_(plugin) {
+  wasm_->contexts_[id_] = this;
 }
 
 // NB: wasm can be nullptr if it failed to be created successfully.
@@ -291,15 +293,6 @@ WasmVm *ContextBase::wasmVm() const { return wasm_->wasm_vm(); }
 
 bool ContextBase::isFailed() { return !wasm_ || wasm_->isFailed(); }
 
-void ContextBase::initializeRootBase(WasmBase *wasm, std::shared_ptr<PluginBase> plugin) {
-  wasm_ = wasm;
-  id_ = wasm->allocContextId();
-  root_id_ = plugin->root_id_;
-  root_log_prefix_ = makeRootLogPrefix(plugin->vm_id_);
-  parent_context_ = this;
-  wasm_->contexts_[id_] = this;
-}
-
 std::string ContextBase::makeRootLogPrefix(std::string_view vm_id) const {
   std::string prefix;
   if (!root_id_.empty()) {
@@ -318,10 +311,10 @@ bool ContextBase::onStart(std::shared_ptr<PluginBase> plugin) {
   DeferAfterCallActions actions(this);
   bool result = true;
   if (wasm_->on_context_create_) {
-    plugin_ = plugin;
+    temp_plugin_ = plugin;
     wasm_->on_context_create_(this, id_, 0);
     in_vm_context_created_ = true;
-    plugin_.reset();
+    temp_plugin_.reset();
   }
   if (wasm_->on_vm_start_) {
     // Do not set plugin_ as the on_vm_start handler should be independent of the plugin since the
@@ -353,11 +346,11 @@ bool ContextBase::onConfigure(std::shared_ptr<PluginBase> plugin) {
   }
 
   DeferAfterCallActions actions(this);
-  plugin_ = plugin;
+  temp_plugin_ = plugin;
   auto result =
       wasm_->on_configure_(this, id_, static_cast<uint32_t>(plugin->plugin_configuration_.size()))
           .u64_ != 0;
-  plugin_.reset();
+  temp_plugin_.reset();
   return result;
 }
 
@@ -656,8 +649,8 @@ FilterMetadataStatus ContextBase::convertVmCallResultToFilterMetadataStatus(uint
 }
 
 ContextBase::~ContextBase() {
-  // Do not remove vm or root contexts which have the same lifetime as wasm_.
-  if (parent_context_id_) {
+  // Do not remove vm context which has the same lifetime as wasm_.
+  if (id_) {
     wasm_->contexts_.erase(id_);
   }
 }