From e2718f4c409e2753b1529a2aceab8c5257b57535 Mon Sep 17 00:00:00 2001
From: John Plevyak <jplevyak@gmail.com>
Date: Fri, 14 Feb 2020 11:08:29 -0800
Subject: [PATCH] Simplify Wasm shutdown. (#410)

Signed-off-by: John Plevyak <jplevyak@gmail.com>
---
 source/extensions/common/wasm/context.cc | 7 -------
 source/extensions/common/wasm/wasm.cc    | 8 +++-----
 source/extensions/common/wasm/wasm.h     | 6 ------
 3 files changed, 3 insertions(+), 18 deletions(-)

diff --git a/source/extensions/common/wasm/context.cc b/source/extensions/common/wasm/context.cc
index 6e226577f045..8e95871129d0 100644
--- a/source/extensions/common/wasm/context.cc
+++ b/source/extensions/common/wasm/context.cc
@@ -1087,7 +1087,6 @@ bool Context::onConfigure(absl::string_view plugin_configuration, PluginSharedPt
       wasm_->on_configure_(this, id_, static_cast<uint32_t>(plugin_configuration.size())).u64_ != 0;
   plugin_.reset();
   configuration_ = "";
-  wasm_->checkShutdown();
   return result;
 }
 
@@ -1101,7 +1100,6 @@ void Context::onTick() {
   if (wasm_->on_tick_) {
     wasm_->on_tick_(this, id_);
   }
-  wasm_->checkShutdown();
 }
 
 void Context::onCreate(uint32_t parent_context_id) {
@@ -1272,14 +1270,12 @@ void Context::onHttpCallResponse(uint32_t token, uint32_t headers, uint32_t body
     return;
   }
   wasm_->on_http_call_response_(this, id_, token, headers, body_size, trailers);
-  wasm_->checkShutdown();
 }
 
 void Context::onQueueReady(uint32_t token) {
   if (wasm_->on_queue_ready_) {
     wasm_->on_queue_ready_(this, id_, token);
   }
-  wasm_->checkShutdown();
 }
 
 void Context::onGrpcCreateInitialMetadata(uint32_t token, Http::HeaderMap& metadata) {
@@ -1290,7 +1286,6 @@ void Context::onGrpcCreateInitialMetadata(uint32_t token, Http::HeaderMap& metad
   wasm_->on_grpc_create_initial_metadata_(this, id_, token,
                                           headerSize(grpc_create_initial_metadata_));
   grpc_create_initial_metadata_ = nullptr;
-  wasm_->checkShutdown();
 }
 
 void Context::onGrpcReceiveInitialMetadata(uint32_t token, Http::HeaderMapPtr&& metadata) {
@@ -1301,7 +1296,6 @@ void Context::onGrpcReceiveInitialMetadata(uint32_t token, Http::HeaderMapPtr&&
   wasm_->on_grpc_receive_initial_metadata_(this, id_, token,
                                            headerSize(grpc_receive_initial_metadata_));
   grpc_receive_initial_metadata_ = nullptr;
-  wasm_->checkShutdown();
 }
 
 void Context::onGrpcReceiveTrailingMetadata(uint32_t token, Http::HeaderMapPtr&& metadata) {
@@ -1312,7 +1306,6 @@ void Context::onGrpcReceiveTrailingMetadata(uint32_t token, Http::HeaderMapPtr&&
   wasm_->on_grpc_receive_trailing_metadata_(this, id_, token,
                                             headerSize(grpc_receive_trailing_metadata_));
   grpc_receive_trailing_metadata_ = nullptr;
-  wasm_->checkShutdown();
 }
 
 WasmResult Context::defineMetric(MetricType type, absl::string_view name, uint32_t* metric_id_ptr) {
diff --git a/source/extensions/common/wasm/wasm.cc b/source/extensions/common/wasm/wasm.cc
index 2da281627088..274625f662dc 100644
--- a/source/extensions/common/wasm/wasm.cc
+++ b/source/extensions/common/wasm/wasm.cc
@@ -447,6 +447,7 @@ uint32_t Wasm::allocContextId() {
 
 class Wasm::ShutdownHandle : public Envoy::Event::DeferredDeletable {
 public:
+  ~ShutdownHandle() { wasm_->finishShutdown(); }
   ShutdownHandle(WasmSharedPtr wasm) : wasm_(wasm) {}
 
 private:
@@ -474,8 +475,8 @@ WasmResult Wasm::done(Context* root_context) {
     return WasmResult::NotFound;
   }
   pending_done_.erase(it);
-  if (pending_done_.empty()) {
-    shutdown_ready_ = true;
+  if (pending_done_.empty() && shutdown_handle_) {
+    dispatcher_.deferredDelete(std::move(shutdown_handle_));
   }
   return WasmResult::Ok;
 }
@@ -484,9 +485,6 @@ void Wasm::finishShutdown() {
   for (auto& p : root_contexts_) {
     p.second->onDelete();
   }
-  if (shutdown_handle_) {
-    dispatcher_.deferredDelete(std::move(shutdown_handle_));
-  }
 }
 
 void Wasm::queueReady(uint32_t root_context_id, uint32_t token) {
diff --git a/source/extensions/common/wasm/wasm.h b/source/extensions/common/wasm/wasm.h
index 17044465e0a6..56350c8acfea 100644
--- a/source/extensions/common/wasm/wasm.h
+++ b/source/extensions/common/wasm/wasm.h
@@ -95,17 +95,11 @@ class Wasm : public Logger::Loggable<Logger::Id::wasm>, public std::enable_share
 
   void setTickPeriod(uint32_t root_context_id, std::chrono::milliseconds tick_period);
 
-  // Handlers for Root Contexts must call checkShutdown() after the call into the VM.
   void tickHandler(uint32_t root_context_id);
   void queueReady(uint32_t root_context_id, uint32_t token);
 
   void startShutdown();
   WasmResult done(Context* root_context);
-  void checkShutdown() {
-    if (shutdown_ready_) {
-      finishShutdown();
-    }
-  }
   void finishShutdown();
 
   //