[wasm] Abort wrapper compilation on isolate shutdown

JS-to-Wasm wrappers embed heap constants (like the undefined value), and
those heap values are being accessed during compilation for tracing.
This is not a data race, since those values are read-only. But if the
isolate dies while we are compiling those wrappers, we might read from
the heap after it has been free'd.

Ideally we would not access the isolate or the heap at all during
compilation, but delaying all tracing until the "finalization" phase is
not feasible, and removing the heap value printing from tracing would
significantly regress quality of this tracing.

Hence this CL only fixes the actual issue: That we keep compiling
wrappers when the isolate is already gone. It does so by introducing an
{OperationsBarrier} per isolate that is being taken by each thread that
executes wrapper compilation, and is used for waiting for background
threads to finish before the isolate shuts down.
Additionally, we actually cancel all compilation if a module dies (or
the isolate shuts down) before it finished baseline compilation. In this
state, the module cannot be shared between isolates yet, so it's safe to
fully cancel all compilation. This cancellation is not strictly
necessary, but it will reduce the time we are blocked while waiting for
wrapper compilation to finish (because no new compilation will start).

R=thibaudm@chromium.org
CC=manoskouk@chromium.org

Bug: v8:11626, chromium:1200231
Change-Id: I5b19141d22bd0cb00ba84ffa53fb07cf001e13cc
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2846881Reviewed-by: Thibaud Michaud <thibaudm@chromium.org>
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#74142}

src/wasm/compilation-environment.h

@@ -119,6 +119,8 @@ class V8_EXPORT_PRIVATE CompilationState {
 
   void CancelCompilation();
 
+  void CancelInitialCompilation();
+
   void SetError();
 
   void SetWireBytesStorage(std::shared_ptr<WireBytesStorage>);

src/wasm/function-compiler.cc

@@ -306,7 +306,8 @@ JSToWasmWrapperCompilationUnit::JSToWasmWrapperCompilationUnit(
     Isolate* isolate, WasmEngine* wasm_engine, const FunctionSig* sig,
     const WasmModule* module, bool is_import,
     const WasmFeatures& enabled_features, AllowGeneric allow_generic)
-    : is_import_(is_import),
+    : isolate_(isolate),
+      is_import_(is_import),
       sig_(sig),
       use_generic_wrapper_(allow_generic && UseGenericWrapper(sig) &&
                            !is_import),
@@ -326,19 +327,19 @@ void JSToWasmWrapperCompilationUnit::Execute() {
   }
 }
 
-Handle<Code> JSToWasmWrapperCompilationUnit::Finalize(Isolate* isolate) {
+Handle<Code> JSToWasmWrapperCompilationUnit::Finalize() {
   Handle<Code> code;
   if (use_generic_wrapper_) {
     code =
-        isolate->builtins()->builtin_handle(Builtins::kGenericJSToWasmWrapper);
+        isolate_->builtins()->builtin_handle(Builtins::kGenericJSToWasmWrapper);
   } else {
-    CompilationJob::Status status = job_->FinalizeJob(isolate);
+    CompilationJob::Status status = job_->FinalizeJob(isolate_);
     CHECK_EQ(status, CompilationJob::SUCCEEDED);
     code = job_->compilation_info()->code();
   }
-  if (!use_generic_wrapper_ && must_record_function_compilation(isolate)) {
+  if (!use_generic_wrapper_ && must_record_function_compilation(isolate_)) {
     RecordWasmHeapStubCompilation(
-        isolate, code, "%s", job_->compilation_info()->GetDebugName().get());
+        isolate_, code, "%s", job_->compilation_info()->GetDebugName().get());
   }
   return code;
 }
@@ -353,7 +354,7 @@ Handle<Code> JSToWasmWrapperCompilationUnit::CompileJSToWasmWrapper(
                                       module, is_import, enabled_features,
                                       kAllowGeneric);
   unit.Execute();
-  return unit.Finalize(isolate);
+  return unit.Finalize();
 }
 
 // static
@@ -366,7 +367,7 @@ Handle<Code> JSToWasmWrapperCompilationUnit::CompileSpecificJSToWasmWrapper(
                                       module, is_import, enabled_features,
                                       kDontAllowGeneric);
   unit.Execute();
-  return unit.Finalize(isolate);
+  return unit.Finalize();
 }
 
 }  // namespace wasm

src/wasm/function-compiler.h

@@ -127,8 +127,10 @@ class V8_EXPORT_PRIVATE JSToWasmWrapperCompilationUnit final {
                                  AllowGeneric allow_generic);
   ~JSToWasmWrapperCompilationUnit();
 
+  Isolate* isolate() const { return isolate_; }
+
   void Execute();
-  Handle<Code> Finalize(Isolate* isolate);
+  Handle<Code> Finalize();
 
   bool is_import() const { return is_import_; }
   const FunctionSig* sig() const { return sig_; }
@@ -146,6 +148,11 @@ class V8_EXPORT_PRIVATE JSToWasmWrapperCompilationUnit final {
                                                      const WasmModule* module);
 
  private:
+  // Wrapper compilation is bound to an isolate. Concurrent accesses to the
+  // isolate (during the "Execute" phase) must be audited carefully, i.e. we
+  // should only access immutable information (like the root table). The isolate
+  // is guaranteed to be alive when this unit executes.
+  Isolate* isolate_;
   bool is_import_;
   const FunctionSig* sig_;
   bool use_generic_wrapper_;

src/wasm/module-compiler.cc

@@ -531,17 +531,19 @@ class CompilationStateImpl {
   CompilationStateImpl(const std::shared_ptr<NativeModule>& native_module,
                        std::shared_ptr<Counters> async_counters);
   ~CompilationStateImpl() {
-    DCHECK(compile_job_->IsValid());
-    compile_job_->CancelAndDetach();
+    if (compile_job_->IsValid()) compile_job_->CancelAndDetach();
   }
 
   // Call right after the constructor, after the {compilation_state_} field in
   // the {NativeModule} has been initialized.
   void InitCompileJob(WasmEngine*);
 
-  // Cancel all background compilation, without waiting for compile tasks to
-  // finish.
-  void CancelCompilation();
+  // {kCancelUnconditionally}: Cancel all compilation.
+  // {kCancelInitialCompilation}: Cancel all compilation if initial (baseline)
+  // compilation is not finished yet.
+  enum CancellationPolicy { kCancelUnconditionally, kCancelInitialCompilation };
+  void CancelCompilation(CancellationPolicy);
+
   bool cancelled() const;
 
   // Initialize compilation progress. Set compilation tiers to expect for
@@ -790,7 +792,14 @@ void CompilationState::InitCompileJob(WasmEngine* engine) {
   Impl(this)->InitCompileJob(engine);
 }
 
-void CompilationState::CancelCompilation() { Impl(this)->CancelCompilation(); }
+void CompilationState::CancelCompilation() {
+  Impl(this)->CancelCompilation(CompilationStateImpl::kCancelUnconditionally);
+}
+
+void CompilationState::CancelInitialCompilation() {
+  Impl(this)->CancelCompilation(
+      CompilationStateImpl::kCancelInitialCompilation);
+}
 
 void CompilationState::SetError() { Impl(this)->SetError(); }
 
@@ -1201,16 +1210,25 @@ CompilationExecutionResult ExecuteJSToWasmWrapperCompilationUnits(
   std::shared_ptr<JSToWasmWrapperCompilationUnit> wrapper_unit = nullptr;
   int num_processed_wrappers = 0;
 
+  OperationsBarrier::Token wrapper_compilation_token;
+  Isolate* isolate;
+
   {
     BackgroundCompileScope compile_scope(native_module);
     if (compile_scope.cancelled()) return kYield;
     wrapper_unit = compile_scope.compilation_state()
                        ->GetNextJSToWasmWrapperCompilationUnit();
     if (!wrapper_unit) return kNoMoreUnits;
+    isolate = wrapper_unit->isolate();
+    wrapper_compilation_token =
+        compile_scope.native_module()->engine()->StartWrapperCompilation(
+            isolate);
+    if (!wrapper_compilation_token) return kNoMoreUnits;
   }
 
   TRACE_EVENT0("v8.wasm", "wasm.JSToWasmWrapperCompilation");
   while (true) {
+    DCHECK_EQ(isolate, wrapper_unit->isolate());
     wrapper_unit->Execute();
     ++num_processed_wrappers;
     bool yield = delegate && delegate->ShouldYield();
@@ -1828,10 +1846,10 @@ std::shared_ptr<StreamingDecoder> AsyncCompileJob::CreateStreamingDecoder() {
 AsyncCompileJob::~AsyncCompileJob() {
   // Note: This destructor always runs on the foreground thread of the isolate.
   background_task_manager_.CancelAndWait();
-  // If the runtime objects were not created yet, then initial compilation did
-  // not finish yet. In this case we can abort compilation.
-  if (native_module_ && module_object_.is_null()) {
-    Impl(native_module_->compilation_state())->CancelCompilation();
+  // If initial compilation did not finish yet we can abort it.
+  if (native_module_) {
+    Impl(native_module_->compilation_state())
+        ->CancelCompilation(CompilationStateImpl::kCancelInitialCompilation);
   }
   // Tell the streaming decoder that the AsyncCompileJob is not available
   // anymore.
@@ -2458,7 +2476,8 @@ void AsyncStreamingProcessor::FinishAsyncCompileJobWithError(
   // Check if there is already a CompiledModule, in which case we have to clean
   // up the CompilationStateImpl as well.
   if (job_->native_module_) {
-    Impl(job_->native_module_->compilation_state())->CancelCompilation();
+    Impl(job_->native_module_->compilation_state())
+        ->CancelCompilation(CompilationStateImpl::kCancelUnconditionally);
 
     job_->DoSync<AsyncCompileJob::DecodeFail,
                  AsyncCompileJob::kUseExistingForegroundTask>(error);
@@ -2782,13 +2801,22 @@ void CompilationStateImpl::InitCompileJob(WasmEngine* engine) {
                                              async_counters_));
 }
 
-void CompilationStateImpl::CancelCompilation() {
+void CompilationStateImpl::CancelCompilation(
+    CompilationStateImpl::CancellationPolicy cancellation_policy) {
+  base::MutexGuard callbacks_guard(&callbacks_mutex_);
+
+  if (cancellation_policy == kCancelInitialCompilation &&
+      finished_events_.contains(
+          CompilationEvent::kFinishedBaselineCompilation)) {
+    // Initial compilation already finished; cannot be cancelled.
+    return;
+  }
+
   // std::memory_order_relaxed is sufficient because no other state is
   // synchronized with |compile_cancelled_|.
   compile_cancelled_.store(true, std::memory_order_relaxed);
 
   // No more callbacks after abort.
-  base::MutexGuard callbacks_guard(&callbacks_mutex_);
   callbacks_.clear();
 }
 
@@ -3039,7 +3067,8 @@ void CompilationStateImpl::FinalizeJSToWasmWrappers(
                js_to_wasm_wrapper_units_.size());
   CodeSpaceMemoryModificationScope modification_scope(isolate->heap());
   for (auto& unit : js_to_wasm_wrapper_units_) {
-    Handle<Code> code = unit->Finalize(isolate);
+    DCHECK_EQ(isolate, unit->isolate());
+    Handle<Code> code = unit->Finalize();
     int wrapper_index =
         GetExportWrapperIndex(module, unit->sig(), unit->is_import());
     (*export_wrappers_out)->set(wrapper_index, *code);
@@ -3448,7 +3477,8 @@ void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,
   for (auto& pair : compilation_units) {
     JSToWasmWrapperKey key = pair.first;
     JSToWasmWrapperCompilationUnit* unit = pair.second.get();
-    Handle<Code> code = unit->Finalize(isolate);
+    DCHECK_EQ(isolate, unit->isolate());
+    Handle<Code> code = unit->Finalize();
     int wrapper_index = GetExportWrapperIndex(module, &key.second, key.first);
     (*export_wrappers_out)->set(wrapper_index, *code);
     RecordStats(*code, isolate->counters());

src/wasm/wasm-engine.cc

@@ -347,7 +347,8 @@ struct WasmEngine::CurrentGCInfo {
 struct WasmEngine::IsolateInfo {
   explicit IsolateInfo(Isolate* isolate)
       : log_codes(WasmCode::ShouldBeLogged(isolate)),
-        async_counters(isolate->async_counters()) {
+        async_counters(isolate->async_counters()),
+        wrapper_compilation_barrier_(std::make_shared<OperationsBarrier>()) {
     v8::Isolate* v8_isolate = reinterpret_cast<v8::Isolate*>(isolate);
     v8::Platform* platform = V8::GetCurrentPlatform();
     foreground_task_runner = platform->GetForegroundTaskRunner(v8_isolate);
@@ -398,6 +399,12 @@ struct WasmEngine::IsolateInfo {
   int throw_count = 0;
   int rethrow_count = 0;
   int catch_count = 0;
+
+  // Operations barrier to synchronize on wrapper compilation on isolate
+  // shutdown.
+  // TODO(wasm): Remove this once we can use the generic js-to-wasm wrapper
+  // everywhere.
+  std::shared_ptr<OperationsBarrier> wrapper_compilation_barrier_;
 };
 
 struct WasmEngine::NativeModuleInfo {
@@ -934,9 +941,10 @@ void WasmEngine::DeleteCompileJobsOnIsolate(Isolate* isolate) {
   // Under the mutex get all jobs to delete. Then delete them without holding
   // the mutex, such that deletion can reenter the WasmEngine.
   std::vector<std::unique_ptr<AsyncCompileJob>> jobs_to_delete;
+  std::vector<std::weak_ptr<NativeModule>> modules_in_isolate;
+  std::shared_ptr<OperationsBarrier> wrapper_compilation_barrier;
   {
     base::MutexGuard guard(&mutex_);
-    DCHECK_EQ(1, isolates_.count(isolate));
     for (auto it = async_compile_jobs_.begin();
          it != async_compile_jobs_.end();) {
       if (it->first->isolate() != isolate) {
@@ -946,7 +954,34 @@ void WasmEngine::DeleteCompileJobsOnIsolate(Isolate* isolate) {
       jobs_to_delete.push_back(std::move(it->second));
       it = async_compile_jobs_.erase(it);
     }
+    DCHECK_EQ(1, isolates_.count(isolate));
+    auto* isolate_info = isolates_[isolate].get();
+    wrapper_compilation_barrier = isolate_info->wrapper_compilation_barrier_;
+    for (auto* native_module : isolate_info->native_modules) {
+      DCHECK_EQ(1, native_modules_.count(native_module));
+      modules_in_isolate.emplace_back(native_modules_[native_module]->weak_ptr);
+    }
   }
+
+  // All modules that have not finished initial compilation yet cannot be
+  // shared with other isolates. Hence we cancel their compilation. In
+  // particular, this will cancel wrapper compilation which is bound to this
+  // isolate (this would be a UAF otherwise).
+  for (auto& weak_module : modules_in_isolate) {
+    if (auto shared_module = weak_module.lock()) {
+      shared_module->compilation_state()->CancelInitialCompilation();
+    }
+  }
+
+  // After cancelling, wait for all current wrapper compilation to actually
+  // finish.
+  wrapper_compilation_barrier->CancelAndWait();
+}
+
+OperationsBarrier::Token WasmEngine::StartWrapperCompilation(Isolate* isolate) {
+  base::MutexGuard guard(&mutex_);
+  DCHECK_EQ(1, isolates_.count(isolate));
+  return isolates_[isolate]->wrapper_compilation_barrier_->TryLock();
 }
 
 void WasmEngine::AddIsolate(Isolate* isolate) {

src/wasm/wasm-engine.h

@@ -246,6 +246,12 @@ class V8_EXPORT_PRIVATE WasmEngine {
   // for tearing down an isolate, or to clean it up to be reused.
   void DeleteCompileJobsOnIsolate(Isolate* isolate);
 
+  // Get a token for compiling wrappers for an Isolate. The token is used to
+  // synchronize background tasks on isolate shutdown. The caller should only
+  // hold the token while compiling export wrappers. If the isolate is already
+  // shutting down, this method will return an invalid token.
+  OperationsBarrier::Token StartWrapperCompilation(Isolate*);
+
   // Manage the set of Isolates that use this WasmEngine.
   void AddIsolate(Isolate* isolate);
   void RemoveIsolate(Isolate* isolate);

test/mjsunit/regress/wasm/regress-1200231.js

+// Copyright 2021 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.
+
+// Flags: --experimental-wasm-reftypes --trace-turbo-graph
+
+load('test/mjsunit/wasm/wasm-module-builder.js');
+
+const builder = new WasmModuleBuilder();
+
+builder.addFunction(`get`, makeSig([], [kWasmExternRef]))
+    .addBody([kExprTableGet])
+    .exportFunc();
+builder.addFunction(`fill`, makeSig([kWasmI32, kWasmAnyFunc, kWasmI32], []))
+    .addBody([])
+    .exportFunc();
+try {
+  builder.toModule();
+} catch {}