[wasm][gc] Make import wrapper cache keep WasmCode alive

The cache also needs to keep the code alive. The code objects are
import wrappers and not wasm functions (which we will focus on first),
but eventually we would also like to collect unused import wrappers.
This CL explicitly increments the ref count when {WasmCode} is added
to the cache, and derements all ref counts in the destructor.

R=titzer@chromium.org

Bug: v8:8217
Change-Id: I1bfb276b25b359d83900147e75ec47788e1fa8de
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1535825Reviewed-by: Ben Titzer <titzer@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#60588}

BUILD.gn

@@ -2773,6 +2773,7 @@ v8_source_set("v8_base") {
     "src/wasm/wasm-feature-flags.h",
     "src/wasm/wasm-features.cc",
     "src/wasm/wasm-features.h",
+    "src/wasm/wasm-import-wrapper-cache.cc",
     "src/wasm/wasm-import-wrapper-cache.h",
     "src/wasm/wasm-interpreter.cc",
     "src/wasm/wasm-interpreter.h",

src/wasm/module-instantiate.cc

@@ -6,7 +6,9 @@
 
 #include "src/asmjs/asm-js.h"
 #include "src/conversions-inl.h"
+#include "src/counters.h"
 #include "src/property-descriptor.h"
+#include "src/tracing/trace-event.h"
 #include "src/utils.h"
 #include "src/wasm/module-compiler.h"
 #include "src/wasm/wasm-import-wrapper-cache.h"

src/wasm/wasm-code-manager.cc

@@ -10,6 +10,7 @@
 #include "src/base/adapters.h"
 #include "src/base/macros.h"
 #include "src/base/platform/platform.h"
+#include "src/counters.h"
 #include "src/disassembler.h"
 #include "src/globals.h"
 #include "src/log.h"
@@ -370,6 +371,22 @@ WasmCode::~WasmCode() {
   }
 }
 
+void WasmCode::DecrementRefCount(Vector<WasmCode*> code_vec) {
+  // Decrement the ref counter of all given code objects. Keep the ones whose
+  // ref count drops to zero.
+  std::unordered_map<NativeModule*, std::vector<WasmCode*>> dead_code;
+  for (WasmCode* code : code_vec) {
+    if (code->DecRef()) dead_code[code->native_module()].push_back(code);
+  }
+
+  // For each native module, free all its code objects at once.
+  for (auto& dead_code_entry : dead_code) {
+    NativeModule* native_module = dead_code_entry.first;
+    Vector<WasmCode*> code_vec = VectorOf(dead_code_entry.second);
+    native_module->FreeCode(code_vec);
+  }
+}
+
 NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,
                            bool can_request_more, VirtualMemory code_space,
                            std::shared_ptr<const WasmModule> module,
@@ -1053,6 +1070,10 @@ NativeModule::~NativeModule() {
   // NativeModule or freeing anything.
   compilation_state_->AbortCompilation();
   engine_->FreeNativeModule(this);
+  // Free the import wrapper cache before releasing the {WasmCode} objects in
+  // {owned_code_}. The destructor of {WasmImportWrapperCache} still needs to
+  // decrease reference counts on the {WasmCode} objects.
+  import_wrapper_cache_.reset();
 }
 
 WasmCodeManager::WasmCodeManager(WasmMemoryTracker* memory_tracker,
@@ -1420,40 +1441,6 @@ NativeModuleModificationScope::~NativeModuleModificationScope() {
 
 namespace {
 thread_local WasmCodeRefScope* current_code_refs_scope = nullptr;
-
-// Receives a vector by value which is modified in this function.
-void DecrementRefCount(std::vector<WasmCode*> code_vec) {
-  // Decrement the ref counter of all given code objects. Keep the ones whose
-  // ref count drops to zero.
-  auto remaining_elements_it = code_vec.begin();
-  for (auto it = code_vec.begin(), end = code_vec.end(); it != end; ++it) {
-    if ((*it)->DecRef()) *remaining_elements_it++ = *it;
-  }
-  code_vec.resize(remaining_elements_it - code_vec.begin());
-
-  // Sort the vector by NativeModule, then by instruction start.
-  std::sort(code_vec.begin(), code_vec.end(),
-            [](const WasmCode* a, const WasmCode* b) {
-              return a->native_module() == b->native_module()
-                         ? a->instruction_start() < b->instruction_start()
-                         : a->native_module() < b->native_module();
-            });
-  // For each native module, free all its code objects at once.
-  auto range_begin = code_vec.begin();
-  while (range_begin != code_vec.end()) {
-    NativeModule* native_module = (*range_begin)->native_module();
-    auto range_end = range_begin + 1;
-    while (range_end < code_vec.end() &&
-           (*range_end)->native_module() == native_module) {
-      ++range_end;
-    }
-    size_t range_size = static_cast<size_t>(range_end - range_begin);
-    Vector<WasmCode*> code_vec{&*range_begin, range_size};
-    native_module->FreeCode(code_vec);
-    range_begin = range_end;
-  }
-}
-
 }  // namespace
 
 WasmCodeRefScope::WasmCodeRefScope()
@@ -1464,7 +1451,10 @@ WasmCodeRefScope::WasmCodeRefScope()
 WasmCodeRefScope::~WasmCodeRefScope() {
   DCHECK_EQ(this, current_code_refs_scope);
   current_code_refs_scope = previous_scope_;
-  DecrementRefCount({code_ptrs_.begin(), code_ptrs_.end()});
+  std::vector<WasmCode*> code_ptrs;
+  code_ptrs.reserve(code_ptrs_.size());
+  code_ptrs.assign(code_ptrs_.begin(), code_ptrs_.end());
+  WasmCode::DecrementRefCount(VectorOf(code_ptrs));
 }
 
 // static

src/wasm/wasm-code-manager.h

@@ -167,6 +167,10 @@ class V8_EXPORT_PRIVATE WasmCode final {
     return old_count == 1;
   }
 
+  // Decrement the ref count on set of {WasmCode} objects, potentially belonging
+  // to different {NativeModule}s.
+  static void DecrementRefCount(Vector<WasmCode*>);
+
   enum FlushICache : bool { kFlushICache = true, kNoFlushICache = false };
 
   static constexpr uint32_t kAnonymousFuncIndex = 0xffffffff;

src/wasm/wasm-import-wrapper-cache.cc

+// Copyright 2019 the V8 project authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "src/wasm/wasm-import-wrapper-cache.h"
+
+#include <vector>
+
+#include "src/counters.h"
+#include "src/wasm/wasm-code-manager.h"
+
+namespace v8 {
+namespace internal {
+namespace wasm {
+
+WasmImportWrapperCache::~WasmImportWrapperCache() {
+  std::vector<WasmCode*> ptrs;
+  ptrs.reserve(entry_map_.size());
+  for (auto& e : entry_map_) ptrs.push_back(e.second);
+  WasmCode::DecrementRefCount(VectorOf(ptrs));
+}
+
+WasmCode* WasmImportWrapperCache::GetOrCompile(
+    WasmEngine* wasm_engine, Counters* counters,
+    compiler::WasmImportCallKind kind, FunctionSig* sig) {
+  base::MutexGuard lock(&mutex_);
+  CacheKey key(static_cast<uint8_t>(kind), *sig);
+  WasmCode*& cached = entry_map_[key];
+  if (cached == nullptr) {
+    // TODO(wasm): no need to hold the lock while compiling an import wrapper.
+    bool source_positions = native_module_->module()->origin == kAsmJsOrigin;
+    // Keep the {WasmCode} alive until we explicitly call {IncRef}.
+    WasmCodeRefScope code_ref_scope;
+    cached = compiler::CompileWasmImportCallWrapper(
+        wasm_engine, native_module_, kind, sig, source_positions);
+    cached->IncRef();
+    counters->wasm_generated_code_size()->Increment(
+        cached->instructions().length());
+    counters->wasm_reloc_size()->Increment(cached->reloc_info().length());
+  }
+  return cached;
+}
+
+}  // namespace wasm
+}  // namespace internal
+}  // namespace v8

src/wasm/wasm-import-wrapper-cache.h

@@ -5,40 +5,35 @@
 #ifndef V8_WASM_WASM_IMPORT_WRAPPER_CACHE_H_
 #define V8_WASM_WASM_IMPORT_WRAPPER_CACHE_H_
 
+#include "src/base/platform/mutex.h"
 #include "src/compiler/wasm-compiler.h"
-#include "src/counters.h"
-#include "src/wasm/value-type.h"
-#include "src/wasm/wasm-code-manager.h"
 
 namespace v8 {
 namespace internal {
+
+class Counters;
+
 namespace wasm {
 
+class WasmCode;
+class WasmEngine;
+
+using FunctionSig = Signature<ValueType>;
+
 // Implements a cache for import wrappers.
 class WasmImportWrapperCache {
  public:
+  ~WasmImportWrapperCache();
+
   WasmCode* GetOrCompile(WasmEngine* wasm_engine, Counters* counters,
-                         compiler::WasmImportCallKind kind, FunctionSig* sig) {
-    base::MutexGuard lock(&mutex_);
-    CacheKey key(static_cast<uint8_t>(kind), *sig);
-    WasmCode*& cached = entry_map_[key];
-    if (cached == nullptr) {
-      // TODO(wasm): no need to hold the lock while compiling an import wrapper.
-      bool source_positions = native_module_->module()->origin == kAsmJsOrigin;
-      cached = compiler::CompileWasmImportCallWrapper(
-          wasm_engine, native_module_, kind, sig, source_positions);
-      counters->wasm_generated_code_size()->Increment(
-          cached->instructions().length());
-      counters->wasm_reloc_size()->Increment(cached->reloc_info().length());
-    }
-    return cached;
-  }
+                         compiler::WasmImportCallKind kind, FunctionSig* sig);
 
  private:
   friend class NativeModule;
+  using CacheKey = std::pair<uint8_t, FunctionSig>;
+
   mutable base::Mutex mutex_;
   NativeModule* native_module_;
-  using CacheKey = std::pair<uint8_t, FunctionSig>;
   std::unordered_map<CacheKey, WasmCode*, base::hash<CacheKey>> entry_map_;
 
   explicit WasmImportWrapperCache(NativeModule* native_module)