[wasm] Update native module info on cache hit

The set of isolates known to a native module and the set of native
modules known to an isolate were not updated on cache hit. This caused
the wasm engine to collect code when it was still live in some isolate.

R=clemensb@chromium.org

Bug: chromium:1055131
Change-Id: I56682509b284c9c0dce7c95ee20ec3929e2e8c9f
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2078583
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#66488}

src/wasm/module-compiler.cc

@@ -1382,8 +1382,8 @@ std::shared_ptr<NativeModule> CompileToNativeModule(
     Handle<FixedArray>* export_wrappers_out) {
   const WasmModule* wasm_module = module.get();
   std::shared_ptr<NativeModule> native_module =
-      isolate->wasm_engine()->MaybeGetNativeModule(wasm_module->origin,
-                                                   wire_bytes.module_bytes());
+      isolate->wasm_engine()->MaybeGetNativeModule(
+          wasm_module->origin, wire_bytes.module_bytes(), isolate);
   if (native_module) {
     // TODO(thibaudm): Look into sharing export wrappers.
     CompileJsToWasmWrappers(isolate, wasm_module, export_wrappers_out);
@@ -1411,7 +1411,7 @@ std::shared_ptr<NativeModule> CompileToNativeModule(
 
   CompileNativeModule(isolate, thrower, wasm_module, native_module.get());
   bool cache_hit = !isolate->wasm_engine()->UpdateNativeModuleCache(
-      thrower->error(), &native_module);
+      thrower->error(), &native_module, isolate);
   if (thrower->error()) return {};
 
   if (cache_hit) {
@@ -1594,7 +1594,7 @@ void AsyncCompileJob::CreateNativeModule(
 bool AsyncCompileJob::GetOrCreateNativeModule(
     std::shared_ptr<const WasmModule> module, size_t code_size_estimate) {
   native_module_ = isolate_->wasm_engine()->MaybeGetNativeModule(
-      module->origin, wire_bytes_.module_bytes());
+      module->origin, wire_bytes_.module_bytes(), isolate_);
   if (native_module_ == nullptr) {
     CreateNativeModule(std::move(module), code_size_estimate);
     return false;
@@ -1720,7 +1720,7 @@ class AsyncCompileJob::CompilationStateCallback {
           std::shared_ptr<NativeModule> native_module = job_->native_module_;
           bool cache_hit =
               !job_->isolate_->wasm_engine()->UpdateNativeModuleCache(
-                  false, &native_module);
+                  false, &native_module, job_->isolate_);
           DCHECK_EQ(cache_hit, native_module != job_->native_module_);
           job_->DoSync<CompileFinished>(cache_hit ? std::move(native_module)
                                                   : nullptr);
@@ -1735,7 +1735,7 @@ class AsyncCompileJob::CompilationStateCallback {
         DCHECK(!last_event_.has_value());
         if (job_->DecrementAndCheckFinisherCount()) {
           job_->isolate_->wasm_engine()->UpdateNativeModuleCache(
-              true, &job_->native_module_);
+              true, &job_->native_module_, job_->isolate_);
           job_->DoSync<CompileFailed>();
         }
         break;
@@ -2371,7 +2371,7 @@ void AsyncStreamingProcessor::OnFinishedStream(OwnedVector<uint8_t> bytes) {
     const bool failed = job_->native_module_->compilation_state()->failed();
     if (!cache_hit) {
       cache_hit = !job_->isolate_->wasm_engine()->UpdateNativeModuleCache(
-          failed, &job_->native_module_);
+          failed, &job_->native_module_, job_->isolate_);
     }
     if (failed) {
       job_->AsyncCompileFailed();

src/wasm/wasm-engine.cc

@@ -450,6 +450,17 @@ MaybeHandle<WasmModuleObject> WasmEngine::SyncCompile(
                             std::move(result).value(), bytes, &export_wrappers);
   if (!native_module) return {};
 
+#ifdef DEBUG
+  // Ensure that code GC will check this isolate for live code.
+  {
+    base::MutexGuard lock(&mutex_);
+    DCHECK_EQ(1, isolates_.count(isolate));
+    DCHECK_EQ(1, isolates_[isolate]->native_modules.count(native_module.get()));
+    DCHECK_EQ(1, native_modules_.count(native_module.get()));
+    DCHECK_EQ(1, native_modules_[native_module.get()]->isolates.count(isolate));
+  }
+#endif
+
   Handle<Script> script =
       CreateWasmScript(isolate, bytes.module_bytes(),
                        VectorOf(native_module->module()->source_map_url),
@@ -866,18 +877,40 @@ std::shared_ptr<NativeModule> WasmEngine::NewNativeModule(
 }
 
 std::shared_ptr<NativeModule> WasmEngine::MaybeGetNativeModule(
-    ModuleOrigin origin, Vector<const uint8_t> wire_bytes) {
-  return native_module_cache_.MaybeGetNativeModule(origin, wire_bytes);
+    ModuleOrigin origin, Vector<const uint8_t> wire_bytes, Isolate* isolate) {
+  std::shared_ptr<NativeModule> native_module =
+      native_module_cache_.MaybeGetNativeModule(origin, wire_bytes);
+  if (native_module) {
+    base::MutexGuard guard(&mutex_);
+    auto& native_module_info = native_modules_[native_module.get()];
+    if (!native_module_info) {
+      native_module_info = std::make_unique<NativeModuleInfo>();
+    }
+    native_module_info->isolates.insert(isolate);
+    isolates_[isolate]->native_modules.insert(native_module.get());
+  }
+  return native_module;
 }
 
 bool WasmEngine::UpdateNativeModuleCache(
-    bool error, std::shared_ptr<NativeModule>* native_module) {
+    bool error, std::shared_ptr<NativeModule>* native_module,
+    Isolate* isolate) {
   // Pass {native_module} by value here to keep it alive until at least after
   // we returned from {Update}. Otherwise, we might {Erase} it inside {Update}
   // which would lock the mutex twice.
   auto prev = native_module->get();
   *native_module = native_module_cache_.Update(*native_module, error);
-  return prev == native_module->get();
+
+  if (prev == native_module->get()) return true;
+
+  base::MutexGuard guard(&mutex_);
+  auto& native_module_info = native_modules_[native_module->get()];
+  if (!native_module_info) {
+    native_module_info = std::make_unique<NativeModuleInfo>();
+  }
+  native_module_info->isolates.insert(isolate);
+  isolates_[isolate]->native_modules.insert((*native_module).get());
+  return false;
 }
 
 bool WasmEngine::GetStreamingCompilationOwnership(size_t prefix_hash) {

src/wasm/wasm-engine.h

@@ -281,7 +281,7 @@ class V8_EXPORT_PRIVATE WasmEngine {
   // threads. The {wire_bytes}' underlying array should be valid at least until
   // the call to {UpdateNativeModuleCache}.
   std::shared_ptr<NativeModule> MaybeGetNativeModule(
-      ModuleOrigin origin, Vector<const uint8_t> wire_bytes);
+      ModuleOrigin origin, Vector<const uint8_t> wire_bytes, Isolate* isolate);
 
   // Replace the temporary {nullopt} with the new native module, or
   // erase it if any error occurred. Wake up blocked threads waiting for this
@@ -293,7 +293,8 @@ class V8_EXPORT_PRIVATE WasmEngine {
   // its {native_module} argument and replace it with the existing entry.
   // Return true in the former case, and false in the latter.
   bool UpdateNativeModuleCache(bool error,
-                               std::shared_ptr<NativeModule>* native_module);
+                               std::shared_ptr<NativeModule>* native_module,
+                               Isolate* isolate);
 
   // Register this prefix hash for a streaming compilation job.
   // If the hash is not in the cache yet, the function returns true and the

src/wasm/wasm-serialization.cc

@@ -620,8 +620,8 @@ MaybeHandle<WasmModuleObject> DeserializeNativeModule(
                                            VectorOf(module->source_map_url),
                                            module->name, source_url);
 
-  auto shared_native_module =
-      wasm_engine->MaybeGetNativeModule(module->origin, wire_bytes_vec);
+  auto shared_native_module = wasm_engine->MaybeGetNativeModule(
+      module->origin, wire_bytes_vec, isolate);
   if (shared_native_module == nullptr) {
     const bool kIncludeLiftoff = false;
     size_t code_size_estimate =
@@ -637,7 +637,7 @@ MaybeHandle<WasmModuleObject> DeserializeNativeModule(
 
     Reader reader(data + WasmSerializer::kHeaderSize);
     bool error = !deserializer.Read(&reader);
-    wasm_engine->UpdateNativeModuleCache(error, &shared_native_module);
+    wasm_engine->UpdateNativeModuleCache(error, &shared_native_module, isolate);
     if (error) return {};
   }