[wasm] Decouple background compile jobs from NativeModule

Background compile jobs should not keep the NativeModule alive, for two
reasons:
1) We sometimes have to wait for background compilation to finish (from
   a foreground task!). This introduces unnecessary latency.
2) Giving the background compile tasks shared ownership of the
   NativeModule causes the NativeModule (and the CompilationState) to
   be freed from background tasks, which is error-prone (see
   https://crrev.com/c/1400420).

Instead, this CL introduces a BackgroundCompileToken which is held
alive by the NativeModule and all background compile jobs. The initial
and the final phase of compilation (getting and submitting work)
synchronize on this token to check and ensure that the NativeModule is
and stays alive. During compilation itself, the mutex is released, such
that the NativeModule can die.
The destructor of the NativeModule cancels the BackgroundCompileToken.
Immediately afterwards, the NativeModule and the CompilationState can
die.

This change allows to remove two hacks introduced previously: The atomic
{aborted_} flag and the {FreeCallbacksTask}.

R=mstarzinger@chromium.org
CC=titzer@chromium.org

Bug: v8:8689, v8:7921
Change-Id: I42e06eab3c944b0988286f2ce18e3c294535dfb6
Cq-Include-Trybots: luci.v8.try:v8_linux64_tsan_rel
Reviewed-on: https://chromium-review.googlesource.com/c/1421364
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
Cr-Commit-Position: refs/heads/master@{#59020}

src/v8.cc

@@ -48,10 +48,10 @@ bool V8::Initialize() {
 
 
 void V8::TearDown() {
+  wasm::WasmEngine::GlobalTearDown();
 #if defined(USE_SIMULATOR)
   Simulator::GlobalTearDown();
 #endif
-  wasm::WasmEngine::GlobalTearDown();
   CallDescriptors::TearDown();
   Bootstrapper::TearDownExtensions();
   ElementsAccessor::TearDown();

src/wasm/compilation-environment.h

@@ -15,6 +15,7 @@ namespace internal {
 namespace wasm {
 
 class NativeModule;
+class WasmCode;
 class WasmError;
 
 enum RuntimeExceptionSupport : bool {
@@ -112,6 +113,8 @@ class CompilationState {
 
   bool failed() const;
 
+  void OnFinishedUnit(ExecutionTier, WasmCode*);
+
  private:
   friend class NativeModule;
   friend class WasmCompilationUnit;

src/wasm/module-compiler.cc

@@ -53,6 +53,64 @@ namespace {
 
 enum class CompileMode : uint8_t { kRegular, kTiering };
 
+// Background compile jobs hold a shared pointer to this token. The token is
+// used to notify them that they should stop. As soon as they see this (after
+// finishing their current compilation unit), they will stop.
+// This allows to already remove the NativeModule without having to synchronize
+// on background compile jobs.
+class BackgroundCompileToken {
+ public:
+  explicit BackgroundCompileToken(NativeModule* native_module)
+      : native_module_(native_module) {}
+
+  void Cancel() {
+    base::MutexGuard mutex_guard(&mutex_);
+    native_module_ = nullptr;
+  }
+
+  // Only call this while holding the {mutex_}.
+  void CancelLocked() { native_module_ = nullptr; }
+
+ private:
+  friend class BackgroundCompileScope;
+  base::Mutex mutex_;
+  NativeModule* native_module_;
+
+  NativeModule* StartScope() {
+    mutex_.Lock();
+    return native_module_;
+  }
+
+  void ExitScope() { mutex_.Unlock(); }
+};
+
+class CompilationStateImpl;
+
+// Keep these scopes short, as they hold the mutex of the token, which
+// sequentializes all these scopes. The mutex is also acquired from foreground
+// tasks, which should not be blocked for a long time.
+class BackgroundCompileScope {
+ public:
+  explicit BackgroundCompileScope(
+      const std::shared_ptr<BackgroundCompileToken>& token)
+      : token_(token.get()), native_module_(token->StartScope()) {}
+
+  ~BackgroundCompileScope() { token_->ExitScope(); }
+
+  bool cancelled() const { return native_module_ == nullptr; }
+
+  NativeModule* native_module() {
+    DCHECK(!cancelled());
+    return native_module_;
+  }
+
+  inline CompilationStateImpl* compilation_state();
+
+ private:
+  BackgroundCompileToken* const token_;
+  NativeModule* const native_module_;
+};
+
 // The {CompilationStateImpl} keeps track of the compilation state of the
 // owning NativeModule, i.e. which functions are left to be compiled.
 // It contains a task manager to allow parallel and asynchronous background
@@ -193,18 +251,6 @@ class CompilationStateImpl {
     std::vector<WasmCode*> code_to_log_;
   };
 
-  class FreeCallbacksTask : public CancelableTask {
-   public:
-    explicit FreeCallbacksTask(CompilationStateImpl* comp_state)
-        : CancelableTask(&comp_state->foreground_task_manager_),
-          compilation_state_(comp_state) {}
-
-    void RunInternal() override { compilation_state_->callbacks_.clear(); }
-
-   private:
-    CompilationStateImpl* const compilation_state_;
-  };
-
   void NotifyOnEvent(CompilationEvent event, const WasmError* error);
 
   std::vector<std::unique_ptr<WasmCompilationUnit>>& finish_units() {
@@ -215,7 +261,9 @@ class CompilationStateImpl {
   // TODO(mstarzinger): Get rid of the Isolate field to make sure the
   // {CompilationStateImpl} can be shared across multiple Isolates.
   Isolate* const isolate_;
+  WasmEngine* const engine_;
   NativeModule* const native_module_;
+  const std::shared_ptr<BackgroundCompileToken> background_compile_token_;
   const CompileMode compile_mode_;
   // Store the value of {WasmCode::ShouldBeLogged()} at creation time of the
   // compilation state.
@@ -267,25 +315,12 @@ class CompilationStateImpl {
   // the foreground thread.
   std::vector<CompilationState::callback_t> callbacks_;
 
-  // Remember whether {Abort()} was called. When set from the foreground this
-  // ensures no more callbacks will be called afterwards. No guarantees when set
-  // from the background. Only needs to be atomic so that it can be set from
-  // foreground and background.
-  std::atomic<bool> aborted_{false};
-
-  CancelableTaskManager background_task_manager_;
   CancelableTaskManager foreground_task_manager_;
   std::shared_ptr<v8::TaskRunner> foreground_task_runner_;
 
   const size_t max_background_tasks_ = 0;
 };
 
-void UpdateFeatureUseCounts(Isolate* isolate, const WasmFeatures& detected) {
-  if (detected.threads) {
-    isolate->CountUsage(v8::Isolate::UseCounterFeature::kWasmThreadOpcodes);
-  }
-}
-
 CompilationStateImpl* Impl(CompilationState* compilation_state) {
   return reinterpret_cast<CompilationStateImpl*>(compilation_state);
 }
@@ -293,6 +328,16 @@ const CompilationStateImpl* Impl(const CompilationState* compilation_state) {
   return reinterpret_cast<const CompilationStateImpl*>(compilation_state);
 }
 
+CompilationStateImpl* BackgroundCompileScope::compilation_state() {
+  return Impl(native_module()->compilation_state());
+}
+
+void UpdateFeatureUseCounts(Isolate* isolate, const WasmFeatures& detected) {
+  if (detected.threads) {
+    isolate->CountUsage(v8::Isolate::UseCounterFeature::kWasmThreadOpcodes);
+  }
+}
+
 }  // namespace
 
 //////////////////////////////////////////////////////
@@ -322,6 +367,10 @@ void CompilationState::AddCallback(CompilationState::callback_t callback) {
 
 bool CompilationState::failed() const { return Impl(this)->failed(); }
 
+void CompilationState::OnFinishedUnit(ExecutionTier tier, WasmCode* code) {
+  Impl(this)->OnFinishedUnit(tier, code);
+}
+
 // static
 std::unique_ptr<CompilationState> CompilationState::New(
     Isolate* isolate, NativeModule* native_module) {
@@ -755,41 +804,81 @@ class FinishCompileTask : public CancelableTask {
 // The runnable task that performs compilations in the background.
 class BackgroundCompileTask : public CancelableTask {
  public:
-  explicit BackgroundCompileTask(CancelableTaskManager* task_manager,
-                                 NativeModule* native_module,
-                                 Counters* counters)
-      : CancelableTask(task_manager),
-        native_module_(native_module),
-        counters_(counters) {}
+  explicit BackgroundCompileTask(CancelableTaskManager* manager,
+                                 std::shared_ptr<BackgroundCompileToken> token,
+                                 std::shared_ptr<Counters> async_counters)
+      : CancelableTask(manager),
+        token_(std::move(token)),
+        async_counters_(std::move(async_counters)) {}
 
   void RunInternal() override {
     TRACE_COMPILE("(3b) Compiling...\n");
     TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"),
                  "BackgroundCompileTask::RunInternal");
-    // The number of currently running background tasks is reduced in
-    // {OnBackgroundTaskStopped}.
-    CompilationEnv env = native_module_->CreateCompilationEnv();
-    auto* compilation_state = Impl(native_module_->compilation_state());
+
+    // These fields are initialized before getting the first unit of work.
+    base::Optional<CompilationEnv> env;
+    std::shared_ptr<WireBytesStorage> wire_bytes;
+    std::shared_ptr<const WasmModule> module;
+
     WasmFeatures detected_features = kNoWasmFeatures;
     double deadline = MonotonicallyIncreasingTimeInMs() + 50.0;
-    while (!compilation_state->failed()) {
-      if (!FetchAndExecuteCompilationUnit(&env, native_module_,
-                                          compilation_state, &detected_features,
-                                          counters_)) {
-        break;
+    while (true) {
+      // Step 1 (synchronized): Get a WasmCompilationUnit, and initialize some
+      // fields if this is the first unit executed by this task.
+      std::unique_ptr<WasmCompilationUnit> unit;
+      {
+        BackgroundCompileScope compile_scope(token_);
+        if (compile_scope.cancelled()) return;
+        if (!env.has_value()) {
+          env.emplace(compile_scope.native_module()->CreateCompilationEnv());
+          wire_bytes = compile_scope.compilation_state()->GetWireBytesStorage();
+          module = compile_scope.native_module()->shared_module();
+        }
+        unit = compile_scope.compilation_state()->GetNextCompilationUnit();
+        if (unit == nullptr) {
+          compile_scope.compilation_state()->OnBackgroundTaskStopped(
+              detected_features);
+          return;
+        }
       }
-      if (deadline < MonotonicallyIncreasingTimeInMs()) {
-        compilation_state->ReportDetectedFeatures(detected_features);
-        compilation_state->RestartBackgroundCompileTask();
-        return;
+
+      // Step 2: Execute the compilation.
+
+      // Get the tier before starting compilation, as compilation can switch
+      // tiers if baseline bails out.
+      ExecutionTier tier = unit->tier();
+      WasmCompilationResult result = unit->ExecuteCompilation(
+          &env.value(), wire_bytes, async_counters_.get(), &detected_features);
+
+      // Step 3 (synchronized): Publish the compilation result.
+      {
+        BackgroundCompileScope compile_scope(token_);
+        if (compile_scope.cancelled()) return;
+        WasmCode* code =
+            unit->Publish(std::move(result), compile_scope.native_module());
+        if (code == nullptr) {
+          compile_scope.compilation_state()->OnBackgroundTaskStopped(
+              detected_features);
+          // Also, cancel all remaining compilation.
+          token_->CancelLocked();
+          return;
+        }
+        compile_scope.compilation_state()->OnFinishedUnit(tier, code);
+        if (deadline < MonotonicallyIncreasingTimeInMs()) {
+          compile_scope.compilation_state()->ReportDetectedFeatures(
+              detected_features);
+          compile_scope.compilation_state()->RestartBackgroundCompileTask();
+          return;
+        }
       }
     }
-    compilation_state->OnBackgroundTaskStopped(detected_features);
+    UNREACHABLE();  // Loop exits via explicit return.
   }
 
  private:
-  NativeModule* const native_module_;
-  Counters* const counters_;
+  std::shared_ptr<BackgroundCompileToken> token_;
+  std::shared_ptr<Counters> async_counters_;
 };
 
 }  // namespace
@@ -1556,7 +1645,10 @@ bool AsyncStreamingProcessor::Deserialize(Vector<const uint8_t> module_bytes,
 CompilationStateImpl::CompilationStateImpl(internal::Isolate* isolate,
                                            NativeModule* native_module)
     : isolate_(isolate),
+      engine_(isolate->wasm_engine()),
       native_module_(native_module),
+      background_compile_token_(
+          std::make_shared<BackgroundCompileToken>(native_module)),
       compile_mode_(FLAG_wasm_tier_up &&
                             native_module->module()->origin == kWasmOrigin
                         ? CompileMode::kTiering
@@ -1571,14 +1663,13 @@ CompilationStateImpl::CompilationStateImpl(internal::Isolate* isolate,
 }
 
 CompilationStateImpl::~CompilationStateImpl() {
-  DCHECK(background_task_manager_.canceled());
   DCHECK(foreground_task_manager_.canceled());
   CompilationError* error = compile_error_.load(std::memory_order_acquire);
   if (error != nullptr) delete error;
 }
 
 void CompilationStateImpl::CancelAndWait() {
-  background_task_manager_.CancelAndWait();
+  Abort();
   foreground_task_manager_.CancelAndWait();
 }
 
@@ -1719,8 +1810,8 @@ void CompilationStateImpl::OnFinishedUnit(ExecutionTier tier, WasmCode* code) {
 }
 
 void CompilationStateImpl::RestartBackgroundCompileTask() {
-  auto task = base::make_unique<BackgroundCompileTask>(
-      &background_task_manager_, native_module_, isolate_->counters());
+  auto task = engine_->NewBackgroundCompileTask<BackgroundCompileTask>(
+      background_compile_token_, isolate_->async_counters());
 
   // If --wasm-num-compilation-tasks=0 is passed, do only spawn foreground
   // tasks. This is used to make timing deterministic.
@@ -1795,16 +1886,9 @@ void CompilationStateImpl::ScheduleFinisherTask() {
 }
 
 void CompilationStateImpl::Abort() {
-  SetError(0, WasmError{0, "Compilation aborted"});
-  background_task_manager_.CancelAndWait();
-  // No more callbacks after abort. Don't free the std::function objects here,
-  // since this might clear references in the embedder, which is only allowed on
-  // the main thread.
-  aborted_.store(true);
-  if (!callbacks_.empty()) {
-    foreground_task_runner_->PostTask(
-        base::make_unique<FreeCallbacksTask>(this));
-  }
+  background_compile_token_->Cancel();
+  // No more callbacks after abort.
+  callbacks_.clear();
 }
 
 void CompilationStateImpl::SetError(uint32_t func_index,
@@ -1831,7 +1915,6 @@ void CompilationStateImpl::SetError(uint32_t func_index,
 
 void CompilationStateImpl::NotifyOnEvent(CompilationEvent event,
                                          const WasmError* error) {
-  if (aborted_.load()) return;
   HandleScope scope(isolate_);
   for (auto& callback : callbacks_) callback(event, error);
   // If no more events are expected after this one, clear the callbacks to free

src/wasm/wasm-code-manager.h

@@ -30,7 +30,6 @@ namespace wasm {
 
 class NativeModule;
 class WasmCodeManager;
-class WasmEngine;
 class WasmMemoryTracker;
 class WasmImportWrapperCache;
 struct WasmModule;
@@ -326,8 +325,9 @@ class V8_EXPORT_PRIVATE NativeModule final {
 
   CompilationState* compilation_state() { return compilation_state_.get(); }
 
-  // Create a {CompilationEnv} object for compilation. Only valid as long as
-  // this {NativeModule} is alive.
+  // Create a {CompilationEnv} object for compilation. The caller has to ensure
+  // that the {WasmModule} pointer stays valid while the {CompilationEnv} is
+  // being used.
   CompilationEnv CreateCompilationEnv() const;
 
   uint32_t num_functions() const {
@@ -341,6 +341,7 @@ class V8_EXPORT_PRIVATE NativeModule final {
   bool lazy_compile_frozen() const { return lazy_compile_frozen_; }
   Vector<const uint8_t> wire_bytes() const { return wire_bytes_->as_vector(); }
   const WasmModule* module() const { return module_.get(); }
+  std::shared_ptr<const WasmModule> shared_module() const { return module_; }
   size_t committed_code_space() const { return committed_code_space_.load(); }
 
   void SetWireBytes(OwnedVector<const uint8_t> wire_bytes);
@@ -423,8 +424,8 @@ class V8_EXPORT_PRIVATE NativeModule final {
   // to be consistent across asynchronous compilations later.
   const WasmFeatures enabled_features_;
 
-  // TODO(clemensh): Make this a unique_ptr (requires refactoring
-  // AsyncCompileJob).
+  // The decoded module, stored in a shared_ptr such that background compile
+  // tasks can keep this alive.
   std::shared_ptr<const WasmModule> module_;
 
   // Wire bytes, held in a shared_ptr so they can be kept alive by the

src/wasm/wasm-engine.cc

@@ -25,6 +25,8 @@ WasmEngine::WasmEngine()
     : code_manager_(&memory_tracker_, FLAG_wasm_max_code_space * MB) {}
 
 WasmEngine::~WasmEngine() {
+  // Synchronize on all background compile tasks.
+  background_compile_task_manager_.CancelAndWait();
   // All AsyncCompileJobs have been canceled.
   DCHECK(jobs_.empty());
   // All Isolates have been deregistered.

src/wasm/wasm-engine.h

@@ -8,6 +8,7 @@
 #include <memory>
 #include <unordered_set>
 
+#include "src/cancelable-task.h"
 #include "src/wasm/wasm-code-manager.h"
 #include "src/wasm/wasm-memory.h"
 #include "src/wasm/wasm-tier.h"
@@ -154,6 +155,12 @@ class V8_EXPORT_PRIVATE WasmEngine {
   // engines this might be a pointer to a new instance or to a shared one.
   static std::shared_ptr<WasmEngine> GetWasmEngine();
 
+  template <typename T, typename... Args>
+  std::unique_ptr<T> NewBackgroundCompileTask(Args&&... args) {
+    return base::make_unique<T>(&background_compile_task_manager_,
+                                std::forward<Args>(args)...);
+  }
+
  private:
   AsyncCompileJob* CreateAsyncCompileJob(
       Isolate* isolate, const WasmFeatures& enabled,
@@ -165,6 +172,10 @@ class V8_EXPORT_PRIVATE WasmEngine {
   WasmCodeManager code_manager_;
   AccountingAllocator allocator_;
 
+  // Task manager managing all background compile jobs. Before shut down of the
+  // engine, they must all be finished because they access the allocator.
+  CancelableTaskManager background_compile_task_manager_;
+
   // This mutex protects all information which is mutated concurrently or
   // fields that are initialized lazily on the first access.
   base::Mutex mutex_;