[Release-1.5] Fix memory leak in TCP (envoyproxy#187)

* Fix memory leak in TCP Signed-off-by: gargnupur <[email protected]> * fix lint Signed-off-by: gargnupur <[email protected]>
gargnupur · Mar 12, 2020 · ec96b58 · ec96b58
1 parent 759e9e1
commit ec96b58
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 2 deletions.
diff --git a/source/extensions/common/wasm/context.cc b/source/extensions/common/wasm/context.cc
@@ -269,6 +269,16 @@ WasmResult Context::setTickPeriod(std::chrono::milliseconds tick_period) {
   return WasmResult::Ok;
 }
 
+void Context::onCloseTCP() {
+  if (tcp_connection_closed_) {
+    return;
+  }
+  tcp_connection_closed_ = true;
+  onDone();
+  onLog();
+  onDelete();
+}
+
 uint64_t Context::getCurrentTimeNanoseconds() {
   return std::chrono::duration_cast<std::chrono::nanoseconds>(
              wasm_->time_source_.systemTime().time_since_epoch())
@@ -431,7 +441,6 @@ Context::FindValue(absl::string_view name, Protobuf::Arena* arena) const {
     break;
   case PropertyToken::FILTER_STATE:
     if (info) {
-
       return CelValue::CreateMap(Protobuf::Arena::Create<WasmStateWrapper>(
           arena, info->filterState(), info->upstreamFilterState().get()));
     }
@@ -1204,13 +1213,23 @@ void Context::onDownstreamConnectionClose(PeerType peer_type) {
     DeferAfterCallActions actions(this);
     wasm_->on_downstream_connection_close_(this, id_, static_cast<uint32_t>(peer_type));
   }
+  downstream_closed_ = true;
+  // Call close on TCP connection, if upstream connection closed or there was a failure seen in this
+  // connection.
+  if (upstream_closed_ || getRequestStreamInfo()->hasAnyResponseFlag()) {
+    onCloseTCP();
+  }
 }
 
 void Context::onUpstreamConnectionClose(PeerType peer_type) {
   if (wasm_->on_upstream_connection_close_) {
     DeferAfterCallActions actions(this);
     wasm_->on_upstream_connection_close_(this, id_, static_cast<uint32_t>(peer_type));
   }
+  upstream_closed_ = true;
+  if (downstream_closed_) {
+    onCloseTCP();
+  }
 }
 
 // Empty headers/trailers have zero size.

diff --git a/source/extensions/common/wasm/context.h b/source/extensions/common/wasm/context.h
@@ -334,6 +334,8 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
 protected:
   friend class Wasm;
 
+  void onCloseTCP();
+
   struct AsyncClientHandler : public Http::AsyncClient::Callbacks {
     // Http::AsyncClient::Callbacks
     void onSuccess(Envoy::Http::MessagePtr&& response) override {
@@ -487,8 +489,13 @@ class Context : public Logger::Loggable<Logger::Id::wasm>,
   std::map<uint32_t, GrpcCallClientHandler> grpc_call_request_;
   std::map<uint32_t, GrpcStreamClientHandler> grpc_stream_;
 
-  // Opaque state
+  // Opaque state.
   absl::flat_hash_map<std::string, std::unique_ptr<StorageObject>> data_storage_;
+
+  // TCP State.
+  bool upstream_closed_ = false;
+  bool downstream_closed_ = false;
+  bool tcp_connection_closed_ = false;
 };
 
 using ContextSharedPtr = std::shared_ptr<Context>;