[wasm][gc] Don't give up mutex during GC finalization

Between determining the set of wasm code objects to free, and actually
freeing them, we should not give up the mutex of the wasm engine.
Otherwise, a NativeModule can die in-between, and we would access a
stale pointer.
This fixes some flakes seen on the TSan bots with --stress-wasm-code-gc.

R=mstarzinger@chromium.org

Bug: v8:8217, v8:9200
Change-Id: Iad5b47379b5be6269180094cfeb2a2f2dfefb425
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1593340Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#61204}

src/wasm/wasm-code-manager.cc

@@ -377,31 +377,24 @@ V8_WARN_UNUSED_RESULT bool WasmCode::DecRefOnPotentiallyDeadCode() {
   }
   // If we reach here, the code was already potentially dead. Decrement the ref
   // count, and return true if it drops to zero.
-  int old_count = ref_count_.load(std::memory_order_relaxed);
-  while (true) {
-    DCHECK_LE(1, old_count);
-    if (ref_count_.compare_exchange_weak(old_count, old_count - 1,
-                                         std::memory_order_relaxed)) {
-      return old_count == 1;
-    }
-  }
+  return DecRefOnDeadCode();
 }
 
 // static
 void WasmCode::DecrementRefCount(Vector<WasmCode* const> 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;
+  WasmEngine::DeadCodeMap dead_code;
+  WasmEngine* engine = nullptr;
   for (WasmCode* code : code_vec) {
-    if (code->DecRef()) dead_code[code->native_module()].push_back(code);
+    if (!code->DecRef()) continue;  // Remaining references.
+    dead_code[code->native_module()].push_back(code);
+    if (!engine) engine = code->native_module()->engine();
+    DCHECK_EQ(engine, code->native_module()->engine());
   }
 
-  // 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->engine()->FreeDeadCode(native_module, code_vec);
-  }
+  DCHECK_EQ(dead_code.empty(), engine == nullptr);
+  if (engine) engine->FreeDeadCode(dead_code);
 }
 
 NativeModule::NativeModule(WasmEngine* engine, const WasmFeatures& enabled,

src/wasm/wasm-code-manager.h

@@ -168,6 +168,13 @@ class V8_EXPORT_PRIVATE WasmCode final {
     }
   }
 
+  // Decrement the ref count on code that is known to be dead, even though there
+  // might still be C++ references. Returns whether this drops the last
+  // reference and the code needs to be freed.
+  V8_WARN_UNUSED_RESULT bool DecRefOnDeadCode() {
+    return ref_count_.fetch_sub(1, std::memory_order_relaxed) == 1;
+  }
+
   // Decrement the ref count on a set of {WasmCode} objects, potentially
   // belonging to different {NativeModule}s. Dead code will be deleted.
   static void DecrementRefCount(Vector<WasmCode* const>);
@@ -229,7 +236,7 @@ class V8_EXPORT_PRIVATE WasmCode final {
 
   // Slow path for {DecRef}: The code becomes potentially dead.
   // Returns whether this code becomes dead and needs to be freed.
-  bool DecRefOnPotentiallyDeadCode();
+  V8_NOINLINE bool DecRefOnPotentiallyDeadCode();
 
   Vector<byte> instructions_;
   OwnedVector<const byte> reloc_info_;

src/wasm/wasm-engine.cc

@@ -648,10 +648,10 @@ void WasmEngine::FreeNativeModule(NativeModule* native_module) {
         } else {
           ++it;
         }
-        TRACE_CODE_GC(
-            "Native module %p died, reducing dead code objects to %zu.\n",
-            native_module, current_gc_info_->dead_code.size());
       }
+      TRACE_CODE_GC(
+          "Native module %p died, reducing dead code objects to %zu.\n",
+          native_module, current_gc_info_->dead_code.size());
     }
     native_modules_.erase(it);
   }
@@ -696,45 +696,44 @@ void WasmEngine::ReportLiveCodeForGC(Isolate* isolate,
                                      Vector<WasmCode*> live_code) {
   TRACE_CODE_GC("Isolate %d reporting %zu live code objects.\n", isolate->id(),
                 live_code.size());
-  // Get the set of dead code under the mutex, but decrement the ref count after
-  // releasing the mutex to avoid deadlocks.
-  OwnedVector<WasmCode*> dead_code;
-  {
-    base::MutexGuard guard(&mutex_);
-    DCHECK_NOT_NULL(current_gc_info_);
-    auto outstanding_isolate_it =
-        current_gc_info_->outstanding_isolates.find(isolate);
-    DCHECK_NE(current_gc_info_->outstanding_isolates.end(),
-              outstanding_isolate_it);
-    auto* fg_task = outstanding_isolate_it->second;
-    if (fg_task) fg_task->Cancel();
-    current_gc_info_->outstanding_isolates.erase(outstanding_isolate_it);
-    for (WasmCode* code : live_code) current_gc_info_->dead_code.erase(code);
-    TRACE_CODE_GC(
-        "Remaining dead code objects: %zu; outstanding isolates: %zu.\n",
-        current_gc_info_->dead_code.size(),
-        current_gc_info_->outstanding_isolates.size());
-
-    // If there are more outstanding isolates, return here.
-    if (!current_gc_info_->outstanding_isolates.empty()) return;
-
-    // All remaining code in {current_gc_info->dead_code} is really dead.
-    // Move it from the set of potentially dead code to the set of dead code,
-    // and decrement its ref count.
-    TRACE_CODE_GC("Decrementing ref count on %zu code objects.\n",
-                  current_gc_info_->dead_code.size());
-    dead_code = OwnedVector<WasmCode*>::Of(current_gc_info_->dead_code);
-    for (WasmCode* code : dead_code) {
-      DCHECK_EQ(1, native_modules_.count(code->native_module()));
-      auto* native_module_info = native_modules_[code->native_module()].get();
-      DCHECK_EQ(1, native_module_info->potentially_dead_code.count(code));
-      native_module_info->potentially_dead_code.erase(code);
-      DCHECK_EQ(0, native_module_info->dead_code.count(code));
-      native_module_info->dead_code.insert(code);
+  DeadCodeMap dead_code;
+  base::MutexGuard guard(&mutex_);
+  DCHECK_NOT_NULL(current_gc_info_);
+  auto outstanding_isolate_it =
+      current_gc_info_->outstanding_isolates.find(isolate);
+  DCHECK_NE(current_gc_info_->outstanding_isolates.end(),
+            outstanding_isolate_it);
+  auto* fg_task = outstanding_isolate_it->second;
+  if (fg_task) fg_task->Cancel();
+  current_gc_info_->outstanding_isolates.erase(outstanding_isolate_it);
+  for (WasmCode* code : live_code) current_gc_info_->dead_code.erase(code);
+  TRACE_CODE_GC(
+      "Remaining dead code objects: %zu; outstanding isolates: %zu.\n",
+      current_gc_info_->dead_code.size(),
+      current_gc_info_->outstanding_isolates.size());
+
+  // If there are more outstanding isolates, return here.
+  if (!current_gc_info_->outstanding_isolates.empty()) return;
+
+  // All remaining code in {current_gc_info->dead_code} is really dead.
+  // Move it from the set of potentially dead code to the set of dead code,
+  // and decrement its ref count.
+  TRACE_CODE_GC("Decrementing ref count on %zu code objects.\n",
+                current_gc_info_->dead_code.size());
+  for (WasmCode* code : current_gc_info_->dead_code) {
+    DCHECK_EQ(1, native_modules_.count(code->native_module()));
+    auto* native_module_info = native_modules_[code->native_module()].get();
+    DCHECK_EQ(1, native_module_info->potentially_dead_code.count(code));
+    native_module_info->potentially_dead_code.erase(code);
+    DCHECK_EQ(0, native_module_info->dead_code.count(code));
+    native_module_info->dead_code.insert(code);
+    if (code->DecRefOnDeadCode()) {
+      dead_code[code->native_module()].push_back(code);
     }
-    current_gc_info_.reset();
   }
-  if (!dead_code.empty()) WasmCode::DecrementRefCount(dead_code.as_vector());
+  current_gc_info_.reset();
+
+  FreeDeadCodeLocked(dead_code);
 }
 
 bool WasmEngine::AddPotentiallyDeadCode(WasmCode* code) {
@@ -763,21 +762,26 @@ bool WasmEngine::AddPotentiallyDeadCode(WasmCode* code) {
   return true;
 }
 
-void WasmEngine::FreeDeadCode(NativeModule* native_module,
-                              Vector<WasmCode* const> codes) {
-  {
-    base::MutexGuard guard(&mutex_);
+void WasmEngine::FreeDeadCode(const DeadCodeMap& dead_code) {
+  base::MutexGuard guard(&mutex_);
+  FreeDeadCodeLocked(dead_code);
+}
+
+void WasmEngine::FreeDeadCodeLocked(const DeadCodeMap& dead_code) {
+  DCHECK(!mutex_.TryLock());
+  for (auto& dead_code_entry : dead_code) {
+    NativeModule* native_module = dead_code_entry.first;
+    const std::vector<WasmCode*>& code_vec = dead_code_entry.second;
     DCHECK_EQ(1, native_modules_.count(native_module));
     auto* info = native_modules_[native_module].get();
-    for (WasmCode* code : codes) {
+    TRACE_CODE_GC("Freeing %zu code object%s of module %p.\n", code_vec.size(),
+                  code_vec.size() == 1 ? "" : "s", native_module);
+    for (WasmCode* code : code_vec) {
       DCHECK_EQ(1, info->dead_code.count(code));
       info->dead_code.erase(code);
     }
+    native_module->FreeCode(VectorOf(code_vec));
   }
-
-  TRACE_CODE_GC("Freeing %zu code object%s of module %p.\n", codes.size(),
-                codes.size() == 1 ? "" : "s", native_module);
-  native_module->FreeCode(codes);
 }
 
 void WasmEngine::TriggerGC() {

src/wasm/wasm-engine.h

@@ -196,8 +196,10 @@ class V8_EXPORT_PRIVATE WasmEngine {
   // case.
   V8_WARN_UNUSED_RESULT bool AddPotentiallyDeadCode(WasmCode*);
 
-  // Free dead code on a native module.
-  void FreeDeadCode(NativeModule*, Vector<WasmCode* const>);
+  // Free dead code.
+  using DeadCodeMap = std::unordered_map<NativeModule*, std::vector<WasmCode*>>;
+  void FreeDeadCode(const DeadCodeMap&);
+  void FreeDeadCodeLocked(const DeadCodeMap&);
 
   // Call on process start and exit.
   static void InitializeOncePerProcess();