Fix Wasm VM shutdown. (envoyproxy#159)

* Move proxy_on_delete() call out of proxy_done() call from the VM (make (envoyproxy#405)

non-reentrant).

Signed-off-by: John Plevyak <jplevyak@gmail.com>

* Simplify Wasm shutdown. (envoyproxy#410)

Signed-off-by: John Plevyak <jplevyak@gmail.com>

Co-authored-by: John Plevyak <jplevyak@gmail.com>

source/extensions/common/wasm/context.cc

@@ -1096,6 +1096,12 @@ std::pair<uint32_t, absl::string_view> Context::getStatus() {
   return std::make_pair(status_code_, status_message_);
 }
 
+void Context::onTick() {
+  if (wasm_->on_tick_) {
+    wasm_->on_tick_(this, id_);
+  }
+}
+
 void Context::onCreate(uint32_t parent_context_id) {
   if (wasm_->on_context_create_) {
     wasm_->on_context_create_(this, id_, parent_context_id);

source/extensions/common/wasm/context.h

@@ -109,6 +109,7 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   virtual bool validateConfiguration(absl::string_view configuration, PluginSharedPtr plugin);
   virtual bool onStart(absl::string_view vm_configuration, PluginSharedPtr plugin);
   virtual bool onConfigure(absl::string_view plugin_configuration, PluginSharedPtr plugin);
+  virtual void onTick();
 
   //
   // Stream downcalls on Context(id > 0).

source/extensions/common/wasm/wasm.cc

@@ -424,7 +424,11 @@ void Wasm::tickHandler(uint32_t root_context_id) {
   auto& tick_period = tick_period_[root_context_id];
   auto& timer = timer_[root_context_id];
   if (on_tick_) {
-    on_tick_(getContext(root_context_id), root_context_id);
+    auto it = contexts_.find(root_context_id);
+    if (it == contexts_.end() || !it->second->isRootContext()) {
+      return;
+    }
+    it->second->onTick();
     if (timer && tick_period.count() > 0) {
       timer->enableTimer(tick_period);
     }
@@ -443,13 +447,14 @@ uint32_t Wasm::allocContextId() {
 
 class Wasm::ShutdownHandle : public Envoy::Event::DeferredDeletable {
 public:
+  ~ShutdownHandle() { wasm_->finishShutdown(); }
   ShutdownHandle(WasmSharedPtr wasm) : wasm_(wasm) {}
 
 private:
   WasmSharedPtr wasm_;
 };
 
-void Wasm::shutdown() {
+void Wasm::startShutdown() {
   bool all_done = true;
   for (auto& p : root_contexts_) {
     if (!p.second->onDone()) {
@@ -459,6 +464,8 @@ void Wasm::shutdown() {
   }
   if (!all_done) {
     shutdown_handle_ = std::make_unique<ShutdownHandle>(shared_from_this());
+  } else {
+    finishShutdown();
   }
 }
 
@@ -468,15 +475,18 @@ WasmResult Wasm::done(Context* root_context) {
     return WasmResult::NotFound;
   }
   pending_done_.erase(it);
-  if (pending_done_.empty()) {
-    for (auto& p : root_contexts_) {
-      p.second->onDelete();
-    }
+  if (pending_done_.empty() && shutdown_handle_) {
     dispatcher_.deferredDelete(std::move(shutdown_handle_));
   }
   return WasmResult::Ok;
 }
 
+void Wasm::finishShutdown() {
+  for (auto& p : root_contexts_) {
+    p.second->onDelete();
+  }
+}
+
 void Wasm::queueReady(uint32_t root_context_id, uint32_t token) {
   auto it = contexts_.find(root_context_id);
   if (it == contexts_.end() || !it->second->isRootContext()) {

source/extensions/common/wasm/wasm.h

@@ -94,11 +94,13 @@ class Wasm : public Logger::Loggable<Logger::Id::wasm>, public std::enable_share
   Event::Dispatcher& dispatcher() { return dispatcher_; }
 
   void setTickPeriod(uint32_t root_context_id, std::chrono::milliseconds tick_period);
+
   void tickHandler(uint32_t root_context_id);
   void queueReady(uint32_t root_context_id, uint32_t token);
 
-  void shutdown();
+  void startShutdown();
   WasmResult done(Context* root_context);
+  void finishShutdown();
 
   //
   // AccessLog::Instance
@@ -190,6 +192,7 @@ class Wasm : public Logger::Loggable<Logger::Id::wasm>, public std::enable_share
   std::unordered_map<uint32_t, Event::TimerPtr> timer_;                  // per root_id.
   std::unique_ptr<ShutdownHandle> shutdown_handle_;
   absl::flat_hash_set<Context*> pending_done_; // Root contexts not done during shutdown.
+  bool shutdown_ready_{false};                 // All pending RootContexts are now done.
 
   TimeSource& time_source_;
 
@@ -274,7 +277,7 @@ class WasmHandle : public Envoy::Server::Wasm,
                    public std::enable_shared_from_this<WasmHandle> {
 public:
   explicit WasmHandle(WasmSharedPtr wasm) : wasm_(wasm) {}
-  ~WasmHandle() { wasm_->shutdown(); }
+  ~WasmHandle() { wasm_->startShutdown(); }
 
   WasmSharedPtr& wasm() { return wasm_; }