[wasm] Switch compilation to Jobs API

Use the new jobs API for WebAssembly compilation. This avoids having to
schedule as many background tasks as there are worker threads. Instead
the one job specifies the maximum concurrency, which changes dynamically
as new compile jobs become available.
This also avoids the artificial deadline we used to ensure that other
tasks get some share of the CPU resources if needed.

Even though this CL moves actual wasm function completely over to the
Jobs API, other similar tasks (like wrapper compilation) are still using
the Task API and need to be ported in a follow-up CL.
Also, we are still using the same priority for baseline compilation and
tier up. We should split this in a follow-up CL to have two jobs with
different priorities. This will also allow us to only block on baseline
compilation where we currently block on both.

R=ahaas@chromium.org
CC=gab@chromium.org

Bug: chromium:1101340
Change-Id: I5656697753346e5fdb15d578425cdb949ac6e364
Cq-Include-Trybots: luci.v8.try:v8_linux64_tsan_rel_ng
Cq-Include-Trybots: luci.v8.try:v8_linux64_tsan_isolates_rel_ng
Cq-Include-Trybots: luci.chromium.try:linux-rel
Cq-Include-Trybots: luci.v8.try:v8_linux_blink_rel
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2280100
Commit-Queue: Clemens Backes <clemensb@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/master@{#69239}

src/wasm/compilation-environment.h

@@ -13,6 +13,9 @@
 #include "src/wasm/wasm-tier.h"
 
 namespace v8 {
+
+class JobHandle;
+
 namespace internal {
 
 class Counters;
@@ -104,7 +107,7 @@ class CompilationState {
 
   ~CompilationState();
 
-  void AbortCompilation();
+  void CancelCompilation();
 
   void SetError();
 

src/wasm/module-compiler.h

@@ -67,7 +67,8 @@ bool CompileLazy(Isolate*, NativeModule*, int func_index);
 
 void TriggerTierUp(Isolate*, NativeModule*, int func_index);
 
-int GetMaxBackgroundTasks();
+// Get the maximum concurrency for parallel compilation.
+int GetMaxCompileConcurrency();
 
 template <typename Key, typename Hash>
 class WrapperQueue {

src/wasm/module-instantiate.cc

@@ -1373,7 +1373,8 @@ void InstanceBuilder::CompileImportWrappers(
   }
 
   CancelableTaskManager task_manager;
-  const int max_background_tasks = GetMaxBackgroundTasks();
+  // TODO(wasm): Switch this to the Jobs API.
+  const int max_background_tasks = GetMaxCompileConcurrency();
   for (int i = 0; i < max_background_tasks; ++i) {
     auto task = std::make_unique<CompileImportWrapperTask>(
         &task_manager, isolate_->wasm_engine(), isolate_->counters(),

src/wasm/wasm-code-manager.cc

@@ -1507,7 +1507,7 @@ NativeModule::~NativeModule() {
   TRACE_HEAP("Deleting native module: %p\n", this);
   // Cancel all background compilation before resetting any field of the
   // NativeModule or freeing anything.
-  compilation_state_->AbortCompilation();
+  compilation_state_->CancelCompilation();
   engine_->FreeNativeModule(this);
   // Free the import wrapper cache before releasing the {WasmCode} objects in
   // {owned_code_}. The destructor of {WasmImportWrapperCache} still needs to

src/wasm/wasm-engine.cc

@@ -401,8 +401,33 @@ WasmEngine::~WasmEngine() {
   gdb_server_.reset();
 #endif  // V8_ENABLE_WASM_GDB_REMOTE_DEBUGGING
 
-  // Synchronize on all background compile tasks.
-  background_compile_task_manager_.CancelAndWait();
+  // Collect the live modules into a vector first, then cancel them while
+  // releasing our lock. This will allow the background tasks to finish.
+  std::vector<std::shared_ptr<NativeModule>> live_modules;
+  {
+    base::MutexGuard guard(&mutex_);
+    for (auto& entry : native_modules_) {
+      if (auto shared_ptr = entry.second->weak_ptr.lock()) {
+        live_modules.emplace_back(std::move(shared_ptr));
+      }
+    }
+  }
+
+  for (auto& native_module : live_modules) {
+    native_module->compilation_state()->CancelCompilation();
+  }
+  live_modules.clear();
+
+  // Now wait for all background compile tasks to actually finish.
+  std::vector<std::shared_ptr<JobHandle>> compile_job_handles;
+  {
+    base::MutexGuard guard(&mutex_);
+    compile_job_handles = compile_job_handles_;
+  }
+  for (auto& job_handle : compile_job_handles) {
+    if (job_handle->IsRunning()) job_handle->Cancel();
+  }
+
   // All AsyncCompileJobs have been canceled.
   DCHECK(async_compile_jobs_.empty());
   // All Isolates have been deregistered.
@@ -1306,6 +1331,14 @@ Handle<Script> WasmEngine::GetOrCreateScript(
   }
 }
 
+void WasmEngine::ShepherdCompileJobHandle(
+    std::shared_ptr<JobHandle> job_handle) {
+  DCHECK_NOT_NULL(job_handle);
+  base::MutexGuard guard(&mutex_);
+  // TODO(clemensb): Add occasional cleanup of finished handles.
+  compile_job_handles_.emplace_back(std::move(job_handle));
+}
+
 void WasmEngine::TriggerGC(int8_t gc_sequence_index) {
   DCHECK(!mutex_.TryLock());
   DCHECK_NULL(current_gc_info_);

src/wasm/wasm-engine.h

@@ -243,12 +243,6 @@ class V8_EXPORT_PRIVATE WasmEngine {
   void AddIsolate(Isolate* isolate);
   void RemoveIsolate(Isolate* isolate);
 
-  template <typename T, typename... Args>
-  std::unique_ptr<T> NewBackgroundCompileTask(Args&&... args) {
-    return std::make_unique<T>(&background_compile_task_manager_,
-                               std::forward<Args>(args)...);
-  }
-
   // Trigger code logging for the given code objects in all Isolates which have
   // access to the NativeModule containing this code. This method can be called
   // from background threads.
@@ -338,6 +332,10 @@ class V8_EXPORT_PRIVATE WasmEngine {
                                    const std::shared_ptr<NativeModule>&,
                                    Vector<const char> source_url = {});
 
+  // Take shared ownership of a compile job handle, such that we can synchronize
+  // on that before the engine dies.
+  void ShepherdCompileJobHandle(std::shared_ptr<JobHandle>);
+
   // Call on process start and exit.
   static void InitializeOncePerProcess();
   static void GlobalTearDown();
@@ -372,10 +370,6 @@ 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_;
-
 #ifdef V8_ENABLE_WASM_GDB_REMOTE_DEBUGGING
   // Implements a GDB-remote stub for WebAssembly debugging.
   std::unique_ptr<gdb_server::GdbServer> gdb_server_;
@@ -403,6 +397,10 @@ class V8_EXPORT_PRIVATE WasmEngine {
   std::unordered_map<NativeModule*, std::unique_ptr<NativeModuleInfo>>
       native_modules_;
 
+  // Background compile jobs that are still running. We need to join them before
+  // the engine gets deleted. Otherwise we don't care when exactly they finish.
+  std::vector<std::shared_ptr<JobHandle>> compile_job_handles_;
+
   // Size of code that became dead since the last GC. If this exceeds a certain
   // threshold, a new GC is triggered.
   size_t new_potentially_dead_code_size_ = 0;

test/cctest/wasm/test-streaming-compilation.cc

@@ -31,6 +31,20 @@ class MockPlatform final : public TestPlatform {
     i::V8::SetPlatformForTesting(this);
   }
 
+  ~MockPlatform() {
+    for (auto* job_handle : job_handles_) job_handle->ResetPlatform();
+  }
+
+  std::unique_ptr<v8::JobHandle> PostJob(
+      v8::TaskPriority priority,
+      std::unique_ptr<v8::JobTask> job_task) override {
+    auto orig_job_handle = TestPlatform::PostJob(priority, std::move(job_task));
+    auto job_handle =
+        std::make_unique<MockJobHandle>(std::move(orig_job_handle), this);
+    job_handles_.insert(job_handle.get());
+    return job_handle;
+  }
+
   std::shared_ptr<TaskRunner> GetForegroundTaskRunner(
       v8::Isolate* isolate) override {
     return task_runner_;
@@ -42,7 +56,12 @@ class MockPlatform final : public TestPlatform {
 
   bool IdleTasksEnabled(v8::Isolate* isolate) override { return false; }
 
-  void ExecuteTasks() { task_runner_->ExecuteTasks(); }
+  void ExecuteTasks() {
+    for (auto* job_handle : job_handles_) {
+      if (job_handle->IsRunning()) job_handle->Join();
+    }
+    task_runner_->ExecuteTasks();
+  }
 
  private:
   class MockTaskRunner final : public TaskRunner {
@@ -75,7 +94,32 @@ class MockPlatform final : public TestPlatform {
     std::queue<std::unique_ptr<v8::Task>> tasks_;
   };
 
+  class MockJobHandle : public JobHandle {
+   public:
+    explicit MockJobHandle(std::unique_ptr<JobHandle> orig_handle,
+                           MockPlatform* platform)
+        : orig_handle_(std::move(orig_handle)), platform_(platform) {}
+
+    ~MockJobHandle() {
+      if (platform_) platform_->job_handles_.erase(this);
+    }
+
+    void ResetPlatform() { platform_ = nullptr; }
+
+    void NotifyConcurrencyIncrease() override {
+      orig_handle_->NotifyConcurrencyIncrease();
+    }
+    void Join() override { orig_handle_->Join(); }
+    void Cancel() override { orig_handle_->Cancel(); }
+    bool IsRunning() override { return orig_handle_->IsRunning(); }
+
+   private:
+    std::unique_ptr<JobHandle> orig_handle_;
+    MockPlatform* platform_;
+  };
+
   std::shared_ptr<MockTaskRunner> task_runner_;
+  std::unordered_set<MockJobHandle*> job_handles_;
 };
 
 namespace {