[wasm] Patch jump tables in all code spaces

If there are multiple code spaces, make sure to patch the jump tables
in all of them.

R=mstarzinger@chromium.org

Bug: v8:9477
Change-Id: I2ec3d3de913b99623fd310004555337329588da0
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1789289Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#63651}

src/wasm/wasm-code-manager.cc

@@ -839,13 +839,13 @@ void NativeModule::UseLazyStub(uint32_t func_index) {
 
   // Add jump table entry for jump to the lazy compile stub.
   uint32_t slot_index = func_index - module_->num_imported_functions;
+  DCHECK_NULL(code_table_[slot_index]);
   DCHECK_NE(runtime_stub_entry(WasmCode::kWasmCompileLazy), kNullAddress);
   Address lazy_compile_target =
       lazy_compile_table_->instruction_start() +
       JumpTableAssembler::LazyCompileSlotIndexToOffset(slot_index);
-  JumpTableAssembler::PatchJumpTableSlot(main_jump_table_->instruction_start(),
-                                         slot_index, lazy_compile_target,
-                                         WasmCode::kFlushICache);
+  base::MutexGuard guard(&allocation_mutex_);
+  PatchJumpTablesLocked(func_index, lazy_compile_target);
 }
 
 // TODO(mstarzinger): Remove {Isolate} parameter once {V8_EMBEDDED_BUILTINS}
@@ -1035,9 +1035,7 @@ WasmCode* NativeModule::PublishCodeLocked(std::unique_ptr<WasmCode> code) {
     }
 
     if (update_jump_table) {
-      JumpTableAssembler::PatchJumpTableSlot(
-          main_jump_table_->instruction_start(), slot_idx,
-          code->instruction_start(), WasmCode::kFlushICache);
+      PatchJumpTablesLocked(code->index(), code->instruction_start());
     }
   }
   WasmCodeRefScope::AddRef(code.get());
@@ -1131,6 +1129,19 @@ WasmCode* NativeModule::CreateEmptyJumpTableInRegion(
   return PublishCode(std::move(code));
 }
 
+void NativeModule::PatchJumpTablesLocked(uint32_t func_index, Address target) {
+  // The caller must hold the {allocation_mutex_}, thus we fail to lock it here.
+  DCHECK(!allocation_mutex_.TryLock());
+
+  uint32_t slot_index = func_index - module_->num_imported_functions;
+  for (auto& code_space_data : code_space_data_) {
+    if (!code_space_data.jump_table) continue;
+    Address jump_table_base = code_space_data.jump_table->instruction_start();
+    JumpTableAssembler::PatchJumpTableSlot(jump_table_base, slot_index, target,
+                                           WasmCode::kFlushICache);
+  }
+}
+
 void NativeModule::AddCodeSpace(base::AddressRegion region) {
   // Each code space must be at least twice as large as the overhead per code
   // space. Otherwise, we are wasting too much memory.

src/wasm/wasm-code-manager.h

@@ -529,6 +529,9 @@ class V8_EXPORT_PRIVATE NativeModule final {
   WasmCode* CreateEmptyJumpTableInRegion(uint32_t jump_table_size,
                                          base::AddressRegion);
 
+  // Hold the {allocation_mutex_} when calling this method.
+  void PatchJumpTablesLocked(uint32_t func_index, Address target);
+
   // Called by the {WasmCodeAllocator} to register a new code space.
   void AddCodeSpace(base::AddressRegion);