[wasm] Move top-tier-finished callback to CompilationState

The AsyncCompileJob should be decoupled from tiering, hence the
top-tier-finished callback should not be delivered via the
AsyncCompileJob. Instead, store it directly on the CompilationState.

R=ahaas@chromium.org

Bug: v8:8050, v8:7921, chromium:912031
Change-Id: Iebd64655667a8078c34caea4edeb6cf5f40833fd
Reviewed-on: https://chromium-review.googlesource.com/c/1371604Reviewed-by: Andreas Haas <ahaas@chromium.org>
Commit-Queue: Clemens Hammacher <clemensh@chromium.org>
Cr-Commit-Position: refs/heads/master@{#58165}

src/wasm/compilation-environment.h

@@ -73,10 +73,20 @@ class WireBytesStorage {
   virtual Vector<const uint8_t> GetCode(WireBytesRef) const = 0;
 };
 
+// Callbacks will receive either {kFailedCompilation} or both
+// {kFinishedBaselineCompilation} and {kFinishedTopTierCompilation}, in that
+// order. If tier up is off, both events are delivered right after each other.
+enum class CompilationEvent : uint8_t {
+  kFinishedBaselineCompilation,
+  kFinishedTopTierCompilation,
+  kFailedCompilation
+};
+
 // The implementation of {CompilationState} lives in module-compiler.cc.
 // This is the PIMPL interface to that private class.
 class CompilationState {
  public:
+  using callback_t = std::function<void(CompilationEvent, const ResultBase*)>;
   ~CompilationState();
 
   void CancelAndWait();
@@ -87,6 +97,8 @@ class CompilationState {
 
   std::shared_ptr<WireBytesStorage> GetWireBytesStorage();
 
+  void AddCallback(callback_t);
+
  private:
   friend class NativeModule;
   CompilationState() = delete;

src/wasm/module-compiler.cc

@@ -53,15 +53,6 @@ namespace wasm {
 
 namespace {
 
-// Callbacks will receive either {kFailedCompilation} or both
-// {kFinishedBaselineCompilation} and {kFinishedTopTierCompilation}, in that
-// order. If tier up is off, both events are delivered right after each other.
-enum class CompilationEvent : uint8_t {
-  kFinishedBaselineCompilation,
-  kFinishedTopTierCompilation,
-  kFailedCompilation
-};
-
 enum class CompileMode : uint8_t { kRegular, kTiering };
 
 // The {CompilationStateImpl} keeps track of the compilation state of the
@@ -71,8 +62,6 @@ enum class CompileMode : uint8_t { kRegular, kTiering };
 // It's public interface {CompilationState} lives in compilation-environment.h.
 class CompilationStateImpl {
  public:
-  using callback_t = std::function<void(CompilationEvent, const VoidResult*)>;
-
   CompilationStateImpl(internal::Isolate*, NativeModule*);
   ~CompilationStateImpl();
 
@@ -85,9 +74,9 @@ class CompilationStateImpl {
   // compilation.
   void SetNumberOfFunctionsToCompile(size_t num_functions);
 
-  // Set the callback function to be called on compilation events. Needs to be
+  // Add the callback function to be called on compilation events. Needs to be
   // set before {AddCompilationUnits} is run.
-  void SetCallback(callback_t callback);
+  void AddCallback(CompilationState::callback_t);
 
   // Inserts new functions to compile and kicks off compilation.
   void AddCompilationUnits(
@@ -268,8 +257,8 @@ class CompilationStateImpl {
   // End of fields protected by {mutex_}.
   //////////////////////////////////////////////////////////////////////////////
 
-  // Callback function to be called on compilation events.
-  callback_t callback_;
+  // Callback functions to be called on compilation events.
+  std::vector<CompilationState::callback_t> callbacks_;
 
   CancelableTaskManager background_task_manager_;
   CancelableTaskManager foreground_task_manager_;
@@ -457,6 +446,10 @@ std::shared_ptr<WireBytesStorage> CompilationState::GetWireBytesStorage() {
   return Impl(this)->GetWireBytesStorage();
 }
 
+void CompilationState::AddCallback(CompilationState::callback_t callback) {
+  return Impl(this)->AddCallback(std::move(callback));
+}
+
 CompilationState::~CompilationState() { Impl(this)->~CompilationStateImpl(); }
 
 // static
@@ -2401,6 +2394,8 @@ void AsyncCompileJob::PrepareRuntimeObjects(
     module_object_ = handle(*module_object_, isolate_);
     deferred_handles_.push_back(deferred.Detach());
   }
+
+  if (stream_) stream_->NotifyRuntimeObjectsCreated(module_object_);
 }
 
 // This function assumes that it is executed in a HandleScope, and that a
@@ -2447,7 +2442,7 @@ class AsyncCompileJob::CompilationStateCallback {
  public:
   explicit CompilationStateCallback(AsyncCompileJob* job) : job_(job) {}
 
-  void operator()(CompilationEvent event, const VoidResult* error_result) {
+  void operator()(CompilationEvent event, const ResultBase* error_result) {
     // This callback is only being called from a foreground task.
     switch (event) {
       case CompilationEvent::kFinishedBaselineCompilation:
@@ -2460,10 +2455,6 @@ class AsyncCompileJob::CompilationStateCallback {
         break;
       case CompilationEvent::kFinishedTopTierCompilation:
         DCHECK_EQ(CompilationEvent::kFinishedBaselineCompilation, last_event_);
-        // Notify embedder that compilation is finished.
-        if (job_->stream_ && job_->stream_->module_compiled_callback()) {
-          job_->stream_->module_compiled_callback()(job_->module_object_);
-        }
         // If a foreground task or a finisher is pending, we rely on
         // FinishModule to remove the job.
         if (!job_->pending_foreground_task_ &&
@@ -2716,7 +2707,7 @@ class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
 
     CompilationStateImpl* compilation_state =
         Impl(job->native_module_->compilation_state());
-    compilation_state->SetCallback(CompilationStateCallback{job});
+    compilation_state->AddCallback(CompilationStateCallback{job});
     if (start_compilation_) {
       // TODO(ahaas): Try to remove the {start_compilation_} check when
       // streaming decoding is done in the background. If
@@ -3022,9 +3013,8 @@ void CompilationStateImpl::SetNumberOfFunctionsToCompile(size_t num_functions) {
   }
 }
 
-void CompilationStateImpl::SetCallback(callback_t callback) {
-  DCHECK_NULL(callback_);
-  callback_ = std::move(callback);
+void CompilationStateImpl::AddCallback(CompilationState::callback_t callback) {
+  callbacks_.emplace_back(std::move(callback));
 }
 
 void CompilationStateImpl::AddCompilationUnits(
@@ -3237,7 +3227,7 @@ void CompilationStateImpl::SetError(uint32_t func_index,
 void CompilationStateImpl::NotifyOnEvent(CompilationEvent event,
                                          const VoidResult* error_result) {
   HandleScope scope(isolate_);
-  if (callback_) callback_(event, error_result);
+  for (auto& callback : callbacks_) callback(event, error_result);
 }
 
 void CompileJsToWasmWrappers(Isolate* isolate, const WasmModule* module,

src/wasm/streaming-decoder.cc

@@ -119,6 +119,48 @@ bool StreamingDecoder::SetCompiledModuleBytes(
   return true;
 }
 
+namespace {
+
+class TopTierCompiledCallback {
+ public:
+  TopTierCompiledCallback(std::shared_ptr<NativeModule> native_module,
+                          StreamingDecoder::ModuleCompiledCallback callback)
+      : native_module_(std::move(native_module)),
+        callback_(std::move(callback)) {}
+
+  void operator()(CompilationEvent event,
+                  const ResultBase* error_result) const {
+    if (event != CompilationEvent::kFinishedTopTierCompilation) return;
+    DCHECK_NULL(error_result);
+    callback_(native_module_);
+#ifdef DEBUG
+    DCHECK(!called_);
+    called_ = true;
+#endif
+  }
+
+ private:
+  const std::shared_ptr<NativeModule> native_module_;
+  const StreamingDecoder::ModuleCompiledCallback callback_;
+#ifdef DEBUG
+  mutable bool called_ = false;
+#endif
+};
+
+}  // namespace
+
+void StreamingDecoder::NotifyRuntimeObjectsCreated(
+    Handle<WasmModuleObject> module_object) {
+  if (!module_compiled_callback_) return;
+  std::shared_ptr<NativeModule> native_module =
+      module_object->shared_native_module();
+  auto* comp_state = module_object->native_module()->compilation_state();
+  comp_state->AddCallback(TopTierCompiledCallback{
+      std::move(native_module), std::move(module_compiled_callback_)});
+  // The callback took ownership of the callback:
+  DCHECK_NULL(module_compiled_callback_);
+}
+
 // An abstract class to share code among the states which decode VarInts. This
 // class takes over the decoding of the VarInt and then calls the actual decode
 // code with the decoded value.

src/wasm/streaming-decoder.h

@@ -84,15 +84,13 @@ class V8_EXPORT_PRIVATE StreamingDecoder {
 
   // Caching support.
   // Sets the callback that is called after the module is fully compiled.
-  using ModuleCompiledCallback = std::function<void(Handle<WasmModuleObject>)>;
+  using ModuleCompiledCallback =
+      std::function<void(const std::shared_ptr<NativeModule>&)>;
   void SetModuleCompiledCallback(ModuleCompiledCallback callback);
   // Passes previously compiled module bytes from the embedder's cache.
   bool SetCompiledModuleBytes(Vector<const uint8_t> compiled_module_bytes);
-  // The callback is stored on the StreamingDecoder so it can be called by the
-  // AsyncCompileJob.
-  ModuleCompiledCallback module_compiled_callback() const {
-    return module_compiled_callback_;
-  }
+
+  void NotifyRuntimeObjectsCreated(Handle<WasmModuleObject>);
 
  private:
   // TODO(ahaas): Put the whole private state of the StreamingDecoder into the

src/wasm/wasm-js.cc

@@ -64,8 +64,9 @@ class WasmStreaming::WasmStreamingImpl {
     // Wrap the embedder callback here so we can also wrap the result as a
     // Local<WasmModuleObject> here.
     streaming_decoder_->SetModuleCompiledCallback(
-        [callback, data](i::Handle<i::WasmModuleObject> module_object) {
-          callback(data, Utils::Convert(module_object->shared_native_module()));
+        [callback,
+         data](const std::shared_ptr<i::wasm::NativeModule>& native_module) {
+          callback(data, Utils::Convert(native_module));
         });
   }