Reland "[wasm] Decouple background compile jobs from NativeModule"

This is a reland of 92d9b09c0e68382d1427ad25c70db016bb9e5b80.
Patch unchanged, errors fixed by https://crrev.com/c/1430059.

Original change's description:
> [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}

TBR=mstarzinger@chromium.org

Bug: v8:8689, v8:7921
Change-Id: Iead972ef77c8503da7246cab48e7693b176d8f02
Cq-Include-Trybots: luci.v8.try:v8_linux64_tsan_rel
Reviewed-on: https://chromium-review.googlesource.com/c/1429862
Reviewed-by: Clemens Hammacher <clemensh@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#59035}
diff --git a/src/v8.cc b/src/v8.cc
index ee2a3ba..fa28a21 100644
--- a/src/v8.cc
+++ b/src/v8.cc
@@ -48,10 +48,10 @@
 
 
 void V8::TearDown() {
+  wasm::WasmEngine::GlobalTearDown();
 #if defined(USE_SIMULATOR)
   Simulator::GlobalTearDown();
 #endif
-  wasm::WasmEngine::GlobalTearDown();
   CallDescriptors::TearDown();
   Bootstrapper::TearDownExtensions();
   ElementsAccessor::TearDown();
diff --git a/src/wasm/compilation-environment.h b/src/wasm/compilation-environment.h
index c6bed6c..d72d205 100644
--- a/src/wasm/compilation-environment.h
+++ b/src/wasm/compilation-environment.h
@@ -15,6 +15,7 @@
 namespace wasm {
 
 class NativeModule;
+class WasmCode;
 class WasmError;
 
 enum RuntimeExceptionSupport : bool {
@@ -112,6 +113,8 @@
 
   bool failed() const;
 
+  void OnFinishedUnit(ExecutionTier, WasmCode*);
+
  private:
   friend class NativeModule;
   friend class WasmCompilationUnit;
diff --git a/src/wasm/module-compiler.cc b/src/wasm/module-compiler.cc
index bdceb0b..90240c9 100644
--- a/src/wasm/module-compiler.cc
+++ b/src/wasm/module-compiler.cc
@@ -53,6 +53,64 @@
 
 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 @@
     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 @@
   // 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 @@
   // 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 @@
   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 @@
 
 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 @@
 // 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 @@
 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() {
-  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::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::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::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
diff --git a/src/wasm/wasm-code-manager.h b/src/wasm/wasm-code-manager.h
index 4247350..9e40d1d 100644
--- a/src/wasm/wasm-code-manager.h
+++ b/src/wasm/wasm-code-manager.h
@@ -30,7 +30,6 @@
 
 class NativeModule;
 class WasmCodeManager;
-class WasmEngine;
 class WasmMemoryTracker;
 class WasmImportWrapperCache;
 struct WasmModule;
@@ -326,8 +325,9 @@
 
   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 @@
   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 @@
   // 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
diff --git a/src/wasm/wasm-engine.cc b/src/wasm/wasm-engine.cc
index d948157..67e29a0 100644
--- a/src/wasm/wasm-engine.cc
+++ b/src/wasm/wasm-engine.cc
@@ -25,6 +25,8 @@
     : 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.
diff --git a/src/wasm/wasm-engine.h b/src/wasm/wasm-engine.h
index 4aa9331..1d7a9be 100644
--- a/src/wasm/wasm-engine.h
+++ b/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 @@
   // 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 @@
   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_;