reland [wasm] Compile JS to WASM wrappers asynchronously

The context was not set during streaming compilation.
The initial upload is the original CL and patch set 1 is the fix.

Original CL:

> [wasm] Compile JS to WASM wrappers asynchronously
>
> R=mstarzinger@chromium.org, ahaas@chromium.org
>
> Bug: v8:9231
> Change-Id: I9e18073bbe25bf8c9c5f9ace102316e6209d0459
> Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1669699
> Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
> Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
> Reviewed-by: Andreas Haas <ahaas@chromium.org>
> Cr-Commit-Position: refs/heads/master@{#62672}

R=mstarzinger@chromium.org, ahaas@chromium.org

Cq-Include-Trybots: luci.v8.try:v8_linux_blink_rel
Bug: v8:9231
Change-Id: I61fc11a6de54cc6e93f3600487a89fa5d2350f0e
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/1701850Reviewed-by: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Michael Starzinger <mstarzinger@chromium.org>
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Auto-Submit: Thibaud Michaud <thibaudm@chromium.org>
Cr-Commit-Position: refs/heads/master@{#62721}

src/compiler/pipeline.cc

@@ -173,12 +173,12 @@ class PipelineData {
 
   // For CodeStubAssembler and machine graph testing entry point.
   PipelineData(ZoneStats* zone_stats, OptimizedCompilationInfo* info,
-               Isolate* isolate, Graph* graph, Schedule* schedule,
-               SourcePositionTable* source_positions,
+               Isolate* isolate, AccountingAllocator* allocator, Graph* graph,
+               Schedule* schedule, SourcePositionTable* source_positions,
                NodeOriginTable* node_origins, JumpOptimizationInfo* jump_opt,
                const AssemblerOptions& assembler_options)
       : isolate_(isolate),
-        allocator_(isolate->allocator()),
+        allocator_(allocator),
         info_(info),
         debug_name_(info_->GetDebugName()),
         zone_stats_(zone_stats),
@@ -1056,16 +1056,19 @@ class WasmHeapStubCompilationJob final : public OptimizedCompilationJob {
       // we pass it to the CompilationJob constructor, but it is not
       // dereferenced there.
       : OptimizedCompilationJob(isolate->stack_guard()->real_climit(), &info_,
-                                "TurboFan"),
+                                "TurboFan", State::kReadyToExecute),
         debug_name_(std::move(debug_name)),
         info_(CStrVector(debug_name_.get()), graph->zone(), kind),
         call_descriptor_(call_descriptor),
-        zone_stats_(isolate->allocator()),
+        zone_stats_(isolate->wasm_engine()->allocator()),
         zone_(std::move(zone)),
         graph_(graph),
-        data_(&zone_stats_, &info_, isolate, graph_, nullptr, source_positions,
-              new (zone_.get()) NodeOriginTable(graph_), nullptr, options),
-        pipeline_(&data_) {}
+        data_(&zone_stats_, &info_, isolate,
+              isolate->wasm_engine()->allocator(), graph_, nullptr,
+              source_positions, new (zone_.get()) NodeOriginTable(graph_),
+              nullptr, options),
+        pipeline_(&data_),
+        wasm_engine_(isolate->wasm_engine()) {}
 
   ~WasmHeapStubCompilationJob() = default;
 
@@ -1083,6 +1086,7 @@ class WasmHeapStubCompilationJob final : public OptimizedCompilationJob {
   Graph* graph_;
   PipelineData data_;
   PipelineImpl pipeline_;
+  wasm::WasmEngine* wasm_engine_;
 
   DISALLOW_COPY_AND_ASSIGN(WasmHeapStubCompilationJob);
 };
@@ -1103,10 +1107,14 @@ Pipeline::NewWasmHeapStubCompilationJob(Isolate* isolate,
 
 CompilationJob::Status WasmHeapStubCompilationJob::PrepareJobImpl(
     Isolate* isolate) {
+  return CompilationJob::SUCCEEDED;
+}
+
+CompilationJob::Status WasmHeapStubCompilationJob::ExecuteJobImpl() {
   std::unique_ptr<PipelineStatistics> pipeline_statistics;
   if (FLAG_turbo_stats || FLAG_turbo_stats_nvp) {
     pipeline_statistics.reset(new PipelineStatistics(
-        &info_, isolate->GetTurboStatistics(), &zone_stats_));
+        &info_, wasm_engine_->GetOrCreateTurboStatistics(), &zone_stats_));
     pipeline_statistics->BeginPhaseKind("V8.WasmStubCodegen");
   }
   if (info_.trace_turbo_json_enabled() || info_.trace_turbo_graph_enabled()) {
@@ -1128,10 +1136,6 @@ CompilationJob::Status WasmHeapStubCompilationJob::PrepareJobImpl(
             << "\", \"source\":\"\",\n\"phases\":[";
   }
   pipeline_.RunPrintAndVerify("V8.WasmMachineCode", true);
-  return CompilationJob::SUCCEEDED;
-}
-
-CompilationJob::Status WasmHeapStubCompilationJob::ExecuteJobImpl() {
   pipeline_.ComputeScheduledGraph();
   if (pipeline_.SelectInstructionsAndAssemble(call_descriptor_)) {
     return CompilationJob::SUCCEEDED;
@@ -2333,8 +2337,8 @@ MaybeHandle<Code> Pipeline::GenerateCodeForCodeStub(
   JumpOptimizationInfo jump_opt;
   bool should_optimize_jumps =
       isolate->serializer_enabled() && FLAG_turbo_rewrite_far_jumps;
-  PipelineData data(&zone_stats, &info, isolate, graph, nullptr,
-                    source_positions, &node_origins,
+  PipelineData data(&zone_stats, &info, isolate, isolate->allocator(), graph,
+                    nullptr, source_positions, &node_origins,
                     should_optimize_jumps ? &jump_opt : nullptr, options);
   data.set_verify_graph(FLAG_verify_csa);
   std::unique_ptr<PipelineStatistics> pipeline_statistics;
@@ -2379,10 +2383,10 @@ MaybeHandle<Code> Pipeline::GenerateCodeForCodeStub(
   // First run code generation on a copy of the pipeline, in order to be able to
   // repeat it for jump optimization. The first run has to happen on a temporary
   // pipeline to avoid deletion of zones on the main pipeline.
-  PipelineData second_data(&zone_stats, &info, isolate, data.graph(),
-                           data.schedule(), data.source_positions(),
-                           data.node_origins(), data.jump_optimization_info(),
-                           options);
+  PipelineData second_data(&zone_stats, &info, isolate, isolate->allocator(),
+                           data.graph(), data.schedule(),
+                           data.source_positions(), data.node_origins(),
+                           data.jump_optimization_info(), options);
   second_data.set_verify_graph(FLAG_verify_csa);
   PipelineImpl second_pipeline(&second_data);
   second_pipeline.SelectInstructionsAndAssemble(call_descriptor);
@@ -2528,8 +2532,8 @@ MaybeHandle<Code> Pipeline::GenerateCodeForTesting(
   // Construct a pipeline for scheduling and code generation.
   ZoneStats zone_stats(isolate->allocator());
   NodeOriginTable* node_positions = new (info->zone()) NodeOriginTable(graph);
-  PipelineData data(&zone_stats, info, isolate, graph, schedule, nullptr,
-                    node_positions, nullptr, options);
+  PipelineData data(&zone_stats, info, isolate, isolate->allocator(), graph,
+                    schedule, nullptr, node_positions, nullptr, options);
   std::unique_ptr<PipelineStatistics> pipeline_statistics;
   if (FLAG_turbo_stats || FLAG_turbo_stats_nvp) {
     pipeline_statistics.reset(new PipelineStatistics(

src/compiler/wasm-compiler.cc

@@ -5954,7 +5954,7 @@ std::unique_ptr<OptimizedCompilationJob> NewJSToWasmCompilationJob(
   // Create the Graph.
   //----------------------------------------------------------------------------
   std::unique_ptr<Zone> zone =
-      base::make_unique<Zone>(isolate->allocator(), ZONE_NAME);
+      base::make_unique<Zone>(isolate->wasm_engine()->allocator(), ZONE_NAME);
   Graph* graph = new (zone.get()) Graph(zone.get());
   CommonOperatorBuilder common(zone.get());
   MachineOperatorBuilder machine(
@@ -6417,8 +6417,7 @@ MaybeHandle<Code> CompileCWasmEntry(Isolate* isolate, wasm::FunctionSig* sig) {
           isolate, incoming, std::move(zone), graph, Code::C_WASM_ENTRY,
           std::move(debug_name), AssemblerOptions::Default(isolate)));
 
-  if (job->PrepareJob(isolate) == CompilationJob::FAILED ||
-      job->ExecuteJob() == CompilationJob::FAILED ||
+  if (job->ExecuteJob() == CompilationJob::FAILED ||
       job->FinalizeJob(isolate) == CompilationJob::FAILED) {
     return {};
   }

src/wasm/function-compiler.cc

@@ -265,15 +265,12 @@ void WasmCompilationUnit::CompileWasmFunction(Isolate* isolate,
 JSToWasmWrapperCompilationUnit::JSToWasmWrapperCompilationUnit(Isolate* isolate,
                                                                FunctionSig* sig,
                                                                bool is_import)
-    : job_(compiler::NewJSToWasmCompilationJob(isolate, sig, is_import)) {}
+    : is_import_(is_import),
+      sig_(sig),
+      job_(compiler::NewJSToWasmCompilationJob(isolate, sig, is_import)) {}
 
 JSToWasmWrapperCompilationUnit::~JSToWasmWrapperCompilationUnit() = default;
 
-void JSToWasmWrapperCompilationUnit::Prepare(Isolate* isolate) {
-  CompilationJob::Status status = job_->PrepareJob(isolate);
-  CHECK_EQ(status, CompilationJob::SUCCEEDED);
-}
-
 void JSToWasmWrapperCompilationUnit::Execute() {
   TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "CompileJSToWasmWrapper");
   DCHECK_EQ(job_->state(), CompilationJob::State::kReadyToExecute);
@@ -297,7 +294,6 @@ Handle<Code> JSToWasmWrapperCompilationUnit::CompileJSToWasmWrapper(
     Isolate* isolate, FunctionSig* sig, bool is_import) {
   // Run the compilation unit synchronously.
   JSToWasmWrapperCompilationUnit unit(isolate, sig, is_import);
-  unit.Prepare(isolate);
   unit.Execute();
   return unit.Finalize(isolate);
 }

src/wasm/function-compiler.h

@@ -108,15 +108,19 @@ class V8_EXPORT_PRIVATE JSToWasmWrapperCompilationUnit final {
                                  bool is_import);
   ~JSToWasmWrapperCompilationUnit();
 
-  void Prepare(Isolate* isolate);
   void Execute();
   Handle<Code> Finalize(Isolate* isolate);
 
+  bool is_import() const { return is_import_; }
+  FunctionSig* sig() const { return sig_; }
+
   // Run a compilation unit synchronously.
   static Handle<Code> CompileJSToWasmWrapper(Isolate* isolate, FunctionSig* sig,
                                              bool is_import);
 
  private:
+  bool is_import_;
+  FunctionSig* sig_;
   std::unique_ptr<OptimizedCompilationJob> job_;
 };
 

src/wasm/module-compiler.cc

@@ -381,7 +381,7 @@ class CompilationStateImpl {
   // Initialize compilation progress. Set compilation tiers to expect for
   // baseline and top tier compilation. Must be set before {AddCompilationUnits}
   // is invoked which triggers background compilation.
-  void InitializeCompilationProgress(bool lazy_module, int num_import_wrappers);
+  void InitializeCompilationProgress(bool lazy_module, int num_wrappers);
 
   // Add the callback function to be called on compilation events. Needs to be
   // set before {AddCompilationUnits} is run to ensure that it receives all
@@ -389,13 +389,24 @@ class CompilationStateImpl {
   void AddCallback(CompilationState::callback_t);
 
   // Inserts new functions to compile and kicks off compilation.
-  void AddCompilationUnits(Vector<WasmCompilationUnit> baseline_units,
-                           Vector<WasmCompilationUnit> top_tier_units);
+  void AddCompilationUnits(
+      Vector<WasmCompilationUnit> baseline_units,
+      Vector<WasmCompilationUnit> top_tier_units,
+      Vector<std::shared_ptr<JSToWasmWrapperCompilationUnit>>
+          js_to_wasm_wrapper_units);
   void AddTopTierCompilationUnit(WasmCompilationUnit);
   base::Optional<WasmCompilationUnit> GetNextCompilationUnit(
       int task_id, CompileBaselineOnly baseline_only);
 
+  std::shared_ptr<JSToWasmWrapperCompilationUnit>
+  GetNextJSToWasmWrapperCompilationUnit();
+  void FinalizeJSToWasmWrappers(Isolate* isolate, const WasmModule* module,
+                                Handle<FixedArray> export_wrappers);
+
   void OnFinishedUnits(Vector<WasmCode*>);
+  void OnFinishedJSToWasmWrapperUnits(int num);
+  void TriggerCallbacks(bool completes_baseline_compilation,
+                        bool completes_top_tier_compilation);
 
   void OnBackgroundTaskStopped(int task_id, const WasmFeatures& detected);
   void UpdateDetectedFeatures(const WasmFeatures& detected);
@@ -483,6 +494,13 @@ class CompilationStateImpl {
   // tasks a fair chance to utilize the worker threads on a regular basis.
   std::atomic<double> next_compilation_deadline_{0};
 
+  // Index of the next wrapper to compile in {js_to_wasm_wrapper_units_}.
+  std::atomic<int> js_to_wasm_wrapper_id_{0};
+  // Wrapper compilation units are stored in shared_ptrs so that they are kept
+  // alive by the tasks even if the NativeModule dies.
+  std::vector<std::shared_ptr<JSToWasmWrapperCompilationUnit>>
+      js_to_wasm_wrapper_units_;
+
   // This mutex protects all information of this {CompilationStateImpl} which is
   // being accessed concurrently.
   mutable base::Mutex mutex_;
@@ -711,6 +729,11 @@ class CompilationUnitBuilder {
     }
   }
 
+  void AddJSToWasmWrapperUnit(
+      std::shared_ptr<JSToWasmWrapperCompilationUnit> unit) {
+    js_to_wasm_wrapper_units_.emplace_back(std::move(unit));
+  }
+
   void AddTopTierUnit(int func_index) {
     ExecutionTierPair tiers = GetRequestedExecutionTiers(
         native_module_->module(), compilation_state()->compile_mode(),
@@ -729,9 +752,13 @@ class CompilationUnitBuilder {
   }
 
   bool Commit() {
-    if (baseline_units_.empty() && tiering_units_.empty()) return false;
-    compilation_state()->AddCompilationUnits(VectorOf(baseline_units_),
-                                             VectorOf(tiering_units_));
+    if (baseline_units_.empty() && tiering_units_.empty() &&
+        js_to_wasm_wrapper_units_.empty()) {
+      return false;
+    }
+    compilation_state()->AddCompilationUnits(
+        VectorOf(baseline_units_), VectorOf(tiering_units_),
+        VectorOf(js_to_wasm_wrapper_units_));
     Clear();
     return true;
   }
@@ -739,6 +766,7 @@ class CompilationUnitBuilder {
   void Clear() {
     baseline_units_.clear();
     tiering_units_.clear();
+    js_to_wasm_wrapper_units_.clear();
   }
 
  private:
@@ -750,6 +778,8 @@ class CompilationUnitBuilder {
   const ExecutionTier default_tier_;
   std::vector<WasmCompilationUnit> baseline_units_;
   std::vector<WasmCompilationUnit> tiering_units_;
+  std::vector<std::shared_ptr<JSToWasmWrapperCompilationUnit>>
+      js_to_wasm_wrapper_units_;
 };
 
 void SetCompileError(ErrorThrower* thrower, ModuleWireBytes wire_bytes,
@@ -909,6 +939,33 @@ void RecordStats(const Code code, Counters* counters) {
 
 constexpr int kMainThreadTaskId = -1;
 
+bool ExecuteJSToWasmWrapperCompilationUnits(
+    const std::shared_ptr<BackgroundCompileToken>& token) {
+  std::shared_ptr<JSToWasmWrapperCompilationUnit> wrapper_unit = nullptr;
+  int num_processed_wrappers = 0;
+  do {
+    // TODO(thibaudm): Reschedule the compilation task if it takes too long, so
+    // that the background thread is not blocked.
+    {
+      BackgroundCompileScope compile_scope(token);
+      if (compile_scope.cancelled()) return false;
+      wrapper_unit = compile_scope.compilation_state()
+                         ->GetNextJSToWasmWrapperCompilationUnit();
+    }
+    if (wrapper_unit) {
+      wrapper_unit->Execute();
+      ++num_processed_wrappers;
+    }
+  } while (wrapper_unit);
+  {
+    BackgroundCompileScope compile_scope(token);
+    if (compile_scope.cancelled()) return false;
+    compile_scope.compilation_state()->OnFinishedJSToWasmWrapperUnits(
+        num_processed_wrappers);
+  }
+  return true;
+}
+
 // Run by the main thread and background tasks to take part in compilation.
 // Returns whether any units were executed.
 bool ExecuteCompilationUnits(
@@ -917,6 +974,13 @@ bool ExecuteCompilationUnits(
   TRACE_COMPILE("Compiling (task %d)...\n", task_id);
   TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "ExecuteCompilationUnits");
 
+  // Execute JS to WASM wrapper units first, so that they are ready to be
+  // finalized by the main thread when the kFinishedBaselineCompilation event is
+  // triggered.
+  if (!ExecuteJSToWasmWrapperCompilationUnits(token)) {
+    return false;
+  }
+
   const bool is_foreground = task_id == kMainThreadTaskId;
   // The main thread uses task id 0, which might collide with one of the
   // background tasks. This is fine, as it will only cause some contention on
@@ -1049,6 +1113,26 @@ bool ExecuteCompilationUnits(
   return true;
 }
 
+using JSToWasmWrapperKey = std::pair<bool, FunctionSig>;
+
+// Returns the number of units added.
+int AddExportWrapperUnits(Isolate* isolate, NativeModule* native_module,
+                          CompilationUnitBuilder* builder) {
+  std::unordered_set<JSToWasmWrapperKey, base::hash<JSToWasmWrapperKey>> keys;
+  for (auto exp : native_module->module()->export_table) {
+    if (exp.kind != kExternalFunction) continue;
+    auto& function = native_module->module()->functions[exp.index];
+    JSToWasmWrapperKey key(function.imported, *function.sig);
+    if (keys.insert(key).second) {
+      auto unit = std::make_shared<JSToWasmWrapperCompilationUnit>(
+          isolate, function.sig, function.imported);
+      builder->AddJSToWasmWrapperUnit(std::move(unit));
+    }
+  }
+
+  return static_cast<int>(keys.size());
+}
+
 // Returns the number of units added.
 int AddImportWrapperUnits(NativeModule* native_module,
                           CompilationUnitBuilder* builder) {
@@ -1074,7 +1158,7 @@ int AddImportWrapperUnits(NativeModule* native_module,
   return static_cast<int>(keys.size());
 }
 
-void InitializeCompilationUnits(NativeModule* native_module) {
+void InitializeCompilationUnits(Isolate* isolate, NativeModule* native_module) {
   CompilationStateImpl* compilation_state =
       Impl(native_module->compilation_state());
   const bool lazy_module = IsLazyModule(native_module->module());
@@ -1098,8 +1182,10 @@ void InitializeCompilationUnits(NativeModule* native_module) {
     }
   }
   int num_import_wrappers = AddImportWrapperUnits(native_module, &builder);
-  compilation_state->InitializeCompilationProgress(lazy_module,
-                                                   num_import_wrappers);
+  int num_export_wrappers =
+      AddExportWrapperUnits(isolate, native_module, &builder);
+  compilation_state->InitializeCompilationProgress(
+      lazy_module, num_import_wrappers + num_export_wrappers);
   builder.Commit();
 }
 
@@ -1201,7 +1287,7 @@ void CompileNativeModule(Isolate* isolate, ErrorThrower* thrower,
   }
 
   // Initialize the compilation units and kick off background compile tasks.
-  InitializeCompilationUnits(native_module);
+  InitializeCompilationUnits(isolate, native_module);
 
   // If tiering is disabled, the main thread can execute any unit (all of them
   // are part of initial compilation). Otherwise, just execute baseline units.
@@ -1291,8 +1377,9 @@ std::shared_ptr<NativeModule> CompileToNativeModule(
   int num_wrappers = MaxNumExportWrappers(native_module->module());
   *export_wrappers_out =
       isolate->factory()->NewFixedArray(num_wrappers, AllocationType::kOld);
-  CompileJsToWasmWrappers(isolate, native_module->module(),
-                          *export_wrappers_out);
+  Impl(native_module->compilation_state())
+      ->FinalizeJSToWasmWrappers(isolate, native_module->module(),
+                                 *export_wrappers_out);
 
   // Log the code within the generated module for profiling.
   native_module->LogWasmCodes(isolate);
@@ -1466,8 +1553,12 @@ void AsyncCompileJob::FinishCompile() {
   // TODO(bbudge) Allow deserialization without wrapper compilation, so we can
   // just compile wrappers here.
   if (!is_after_deserialization) {
-    // TODO(wasm): compiling wrappers should be made async.
-    CompileWrappers();
+    CompilationStateImpl* compilation_state =
+        Impl(native_module_->compilation_state());
+    Handle<FixedArray> export_wrappers =
+        handle(module_object_->export_wrappers(), isolate_);
+    compilation_state->FinalizeJSToWasmWrappers(
+        isolate_, module_object_->module(), export_wrappers);
   }
 
   FinishModule();
@@ -1788,7 +1879,7 @@ class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
       // then DoAsync would do the same as NextStep already.
 
       // Add compilation units and kick off compilation.
-      InitializeCompilationUnits(job->native_module_.get());
+      InitializeCompilationUnits(job->isolate(), job->native_module_.get());
     }
   }
 };
@@ -1849,17 +1940,8 @@ class AsyncCompileJob::CompileFinished : public CompileStep {
   }
 };
 
-void AsyncCompileJob::CompileWrappers() {
-  // TODO(wasm): Compile all wrappers here, including the start function wrapper
-  // and the wrappers for the function table elements.
-  TRACE_COMPILE("(5) Compile wrappers...\n");
-  // Compile JS->wasm wrappers for exported functions.
-  CompileJsToWasmWrappers(isolate_, module_object_->native_module()->module(),
-                          handle(module_object_->export_wrappers(), isolate_));
-}
-
 void AsyncCompileJob::FinishModule() {
-  TRACE_COMPILE("(6) Finish module...\n");
+  TRACE_COMPILE("(4) Finish module...\n");
   AsyncCompileSucceeded(module_object_);
   isolate_->wasm_engine()->RemoveCompileJob(this);
 }
@@ -1969,8 +2051,12 @@ bool AsyncStreamingProcessor::ProcessCodeSectionHeader(
 
   int num_import_wrappers =
       AddImportWrapperUnits(native_module, compilation_unit_builder_.get());
-  compilation_state->InitializeCompilationProgress(lazy_module,
-                                                   num_import_wrappers);
+  HandleScope scope(job_->isolate_);
+  SaveAndSwitchContext saved_context(job_->isolate_, *job_->native_context_);
+  int num_export_wrappers = AddExportWrapperUnits(
+      job_->isolate_, native_module, compilation_unit_builder_.get());
+  compilation_state->InitializeCompilationProgress(
+      lazy_module, num_import_wrappers + num_export_wrappers);
   return true;
 }
 
@@ -2141,8 +2227,8 @@ void CompilationStateImpl::AbortCompilation() {
   callbacks_.clear();
 }
 
-void CompilationStateImpl::InitializeCompilationProgress(
-    bool lazy_module, int num_import_wrappers) {
+void CompilationStateImpl::InitializeCompilationProgress(bool lazy_module,
+                                                         int num_wrappers) {
   DCHECK(!failed());
   auto enabled_features = native_module_->enabled_features();
   auto* module = native_module_->module();
@@ -2186,7 +2272,7 @@ void CompilationStateImpl::InitializeCompilationProgress(
   DCHECK_IMPLIES(lazy_module, outstanding_top_tier_functions_ == 0);
   DCHECK_LE(0, outstanding_baseline_units_);
   DCHECK_LE(outstanding_baseline_units_, outstanding_top_tier_functions_);
-  outstanding_baseline_units_ += num_import_wrappers;
+  outstanding_baseline_units_ += num_wrappers;
 
   // Trigger callbacks if module needs no baseline or top tier compilation. This
   // can be the case for an empty or fully lazy module.
@@ -2211,15 +2297,50 @@ void CompilationStateImpl::AddCallback(CompilationState::callback_t callback) {
 
 void CompilationStateImpl::AddCompilationUnits(
     Vector<WasmCompilationUnit> baseline_units,
-    Vector<WasmCompilationUnit> top_tier_units) {
-  compilation_unit_queues_.AddUnits(baseline_units, top_tier_units,
-                                    native_module_->module());
+    Vector<WasmCompilationUnit> top_tier_units,
+    Vector<std::shared_ptr<JSToWasmWrapperCompilationUnit>>
+        js_to_wasm_wrapper_units) {
+  if (!baseline_units.empty() || !top_tier_units.empty()) {
+    compilation_unit_queues_.AddUnits(baseline_units, top_tier_units,
+                                      native_module_->module());
+  }
+  js_to_wasm_wrapper_units_.insert(js_to_wasm_wrapper_units_.end(),
+                                   js_to_wasm_wrapper_units.begin(),
+                                   js_to_wasm_wrapper_units.end());
 
   RestartBackgroundTasks();
 }
 
 void CompilationStateImpl::AddTopTierCompilationUnit(WasmCompilationUnit unit) {
-  AddCompilationUnits({}, {&unit, 1});
+  AddCompilationUnits({}, {&unit, 1}, {});
+}
+
+std::shared_ptr<JSToWasmWrapperCompilationUnit>
+CompilationStateImpl::GetNextJSToWasmWrapperCompilationUnit() {
+  int wrapper_id =
+      js_to_wasm_wrapper_id_.fetch_add(1, std::memory_order_relaxed);
+  if (wrapper_id < static_cast<int>(js_to_wasm_wrapper_units_.size())) {
+    return js_to_wasm_wrapper_units_[wrapper_id];
+  }
+  return nullptr;
+}
+
+void CompilationStateImpl::FinalizeJSToWasmWrappers(
+    Isolate* isolate, const WasmModule* module,
+    Handle<FixedArray> export_wrappers) {
+  // TODO(6792): Wrappers below are allocated with {Factory::NewCode}. As an
+  // optimization we keep the code space unlocked to avoid repeated unlocking
+  // because many such wrapper are allocated in sequence below.
+  TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"),
+               "FinalizeJSToWasmWrappers");
+  CodeSpaceMemoryModificationScope modification_scope(isolate->heap());
+  for (auto& unit : js_to_wasm_wrapper_units_) {
+    Handle<Code> code = unit->Finalize(isolate);
+    int wrapper_index =
+        GetExportWrapperIndex(module, unit->sig(), unit->is_import());
+    export_wrappers->set(wrapper_index, *code);
+    RecordStats(*code, isolate->counters());
+  }
 }
 
 base::Optional<WasmCompilationUnit>
@@ -2309,25 +2430,38 @@ void CompilationStateImpl::OnFinishedUnits(Vector<WasmCode*> code_vector) {
       DCHECK_LE(0, outstanding_baseline_units_);
     }
 
-    // Trigger callbacks.
-    if (completes_baseline_compilation) {
-      TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "BaselineFinished");
-      for (auto& callback : callbacks_) {
-        callback(CompilationEvent::kFinishedBaselineCompilation);
-      }
-      if (outstanding_top_tier_functions_ == 0) {
-        completes_top_tier_compilation = true;
-      }
+    TriggerCallbacks(completes_baseline_compilation,
+                     completes_top_tier_compilation);
+  }
+}
+
+void CompilationStateImpl::OnFinishedJSToWasmWrapperUnits(int num) {
+  if (num == 0) return;
+  base::MutexGuard guard(&callbacks_mutex_);
+  outstanding_baseline_units_ -= num;
+  bool completes_baseline_compilation = outstanding_baseline_units_ == 0;
+  TriggerCallbacks(completes_baseline_compilation, false);
+}
+
+void CompilationStateImpl::TriggerCallbacks(
+    bool completes_baseline_compilation, bool completes_top_tier_compilation) {
+  if (completes_baseline_compilation) {
+    TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "BaselineFinished");
+    for (auto& callback : callbacks_) {
+      callback(CompilationEvent::kFinishedBaselineCompilation);
     }
-    if (outstanding_baseline_units_ == 0 && completes_top_tier_compilation) {
-      TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "TopTierFinished");
-      for (auto& callback : callbacks_) {
-        callback(CompilationEvent::kFinishedTopTierCompilation);
-      }
-      // Clear the callbacks because no more events will be delivered.
-      callbacks_.clear();
+    if (outstanding_top_tier_functions_ == 0) {
+      completes_top_tier_compilation = true;
     }
   }
+  if (outstanding_baseline_units_ == 0 && completes_top_tier_compilation) {
+    TRACE_EVENT0(TRACE_DISABLED_BY_DEFAULT("v8.wasm"), "TopTierFinished");
+    for (auto& callback : callbacks_) {
+      callback(CompilationEvent::kFinishedTopTierCompilation);
+    }
+    // Clear the callbacks because no more events will be delivered.
+    callbacks_.clear();
+  }
 }
 
 void CompilationStateImpl::OnBackgroundTaskStopped(
@@ -2375,6 +2509,11 @@ void CompilationStateImpl::RestartBackgroundTasks() {
     if (failed()) return;
 
     size_t max_num_restart = compilation_unit_queues_.GetTotalSize();
+    if (js_to_wasm_wrapper_id_ <
+        static_cast<int>(js_to_wasm_wrapper_units_.size())) {
+      max_num_restart +=
+          js_to_wasm_wrapper_units_.size() - js_to_wasm_wrapper_id_;
+    }
 
     while (!available_task_ids_.empty() && max_num_restart-- > 0) {
       int task_id = available_task_ids_.back();
@@ -2414,7 +2553,6 @@ void CompilationStateImpl::SetError() {
 }
 
 namespace {
-using JSToWasmWrapperKey = std::pair<bool, FunctionSig>;
 using JSToWasmWrapperQueue =
     WrapperQueue<JSToWasmWrapperKey, base::hash<JSToWasmWrapperKey>>;
 using JSToWasmWrapperUnitMap =
@@ -2457,7 +2595,6 @@ void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,
     if (queue.insert(key)) {
       auto unit = base::make_unique<JSToWasmWrapperCompilationUnit>(
           isolate, function.sig, function.imported);
-      unit->Prepare(isolate);
       compilation_units.emplace(key, std::move(unit));
     }
   }

src/wasm/module-compiler.h

@@ -153,8 +153,6 @@ class AsyncCompileJob {
 
   void AsyncCompileSucceeded(Handle<WasmModuleObject> result);
 
-  void CompileWrappers();
-
   void FinishModule();
 
   void StartForegroundTask();

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

@@ -154,10 +154,18 @@ class StreamTester {
   bool IsPromisePending() { return state_ == CompilationState::kPending; }
 
   void OnBytesReceived(const uint8_t* start, size_t length) {
+    // Streaming compiler is expected to set its own context and handle scope.
+    i::SaveAndSwitchContext saved_context(CcTest::i_isolate(), i::Context{});
+    v8::SealHandleScope seal_handle_scope(CcTest::isolate());
     stream_->OnBytesReceived(Vector<const uint8_t>(start, length));
   }
 
-  void FinishStream() { stream_->Finish(); }
+  void FinishStream() {
+    // Streaming compiler is expected to set its own context and handle scope.
+    i::SaveAndSwitchContext saved_context(CcTest::i_isolate(), i::Context{});
+    v8::SealHandleScope seal_handle_scope(CcTest::isolate());
+    stream_->Finish();
+  }
 
   void SetCompiledModuleBytes(const uint8_t* start, size_t length) {
     stream_->SetCompiledModuleBytes(Vector<const uint8_t>(start, length));