[api] Add WebAssembly caching API

- Adds embedder callback to notify fully tiered compilation is finished,
  returning a WasmCompiledModule for serialization.
- Adds function to pass previously compiled bytes into WASM streaming
  compilation, for deserialization.
- Plumbs this API through StreamingDecoder.

Bug: chromium:719172
Cq-Include-Trybots: luci.chromium.try:linux_chromium_rel_ng
Change-Id: Ibe376f3a8ccfa90fda730ef4ff6628a1532da45c
Reviewed-on: https://chromium-review.googlesource.com/c/1252884Reviewed-by: Adam Klein <adamk@chromium.org>
Reviewed-by: Andreas Haas <ahaas@chromium.org>
Reviewed-by: Clemens Hammacher <clemensh@chromium.org>
Commit-Queue: Bill Budge <bbudge@chromium.org>
Cr-Commit-Position: refs/heads/master@{#56617}

include/v8.h

@@ -4364,7 +4364,7 @@ class V8_EXPORT WasmStreaming final {
   ~WasmStreaming();
 
   /**
-   * Pass a new chunck of bytes to WebAssembly streaming compilation.
+   * Pass a new chunk of bytes to WebAssembly streaming compilation.
    * The buffer passed into {OnBytesReceived} is owned by the caller.
    */
   void OnBytesReceived(const uint8_t* bytes, size_t size);
@@ -4383,6 +4383,29 @@ class V8_EXPORT WasmStreaming final {
    */
   void Abort(MaybeLocal<Value> exception);
 
+  /**
+   * Callback for module compiled notifications. |data| is the identifier
+   * passed to {SetModuleCompiledCallback}, |compiled_module| is the result.
+   */
+  typedef void (*ModuleCompiledCallback)(
+      intptr_t data, Local<WasmCompiledModule> compiled_module);
+
+  /**
+   * Sets a callback for when compilation of the Wasm module has been completed
+   * to the highest tier. |data| will be passed as the first callback parameter.
+   */
+  void SetModuleCompiledCallback(ModuleCompiledCallback callback,
+                                 intptr_t data);
+
+  /**
+   * Passes previously compiled module bytes. This must be called before calling
+   * any non-static methods of this class. Returns true if the module bytes can
+   * be used, false otherwise. The buffer passed into {SetCompiledModuleBytes}
+   * is owned by the caller. If {SetCompiledModuleBytes} returns true, the
+   * buffer must remain valid until either {Finish} or {Abort} completes.
+   */
+  bool SetCompiledModuleBytes(const uint8_t* bytes, size_t size);
+
   /**
    * Unpacks a {WasmStreaming} object wrapped in a  {Managed} for the embedder.
    * Since the embedder is on the other side of the API, it cannot unpack the

src/wasm/module-compiler.cc

@@ -25,6 +25,7 @@
 #include "src/wasm/wasm-memory.h"
 #include "src/wasm/wasm-objects-inl.h"
 #include "src/wasm/wasm-result.h"
+#include "src/wasm/wasm-serialization.h"
 
 #define TRACE(...)                                      \
   do {                                                  \
@@ -2212,6 +2213,9 @@ class AsyncStreamingProcessor final : public StreamingProcessor {
 
   void OnAbort() override;
 
+  bool Deserialize(Vector<const uint8_t> wire_bytes,
+                   Vector<const uint8_t> module_bytes) override;
+
  private:
   // Finishes the AsyncCompileJob with an error.
   void FinishAsyncCompileJobWithError(ResultBase result);
@@ -2245,7 +2249,7 @@ AsyncCompileJob::~AsyncCompileJob() {
 
 // This function assumes that it is executed in a HandleScope, and that a
 // context is set on the isolate.
-void AsyncCompileJob::FinishCompile() {
+void AsyncCompileJob::FinishCompile(bool compile_wrappers) {
   DCHECK_NOT_NULL(isolate_->context());
   // Finish the wasm script now and make it public to the debugger.
   Handle<Script> script(module_object_->script(), isolate_);
@@ -2265,8 +2269,14 @@ void AsyncCompileJob::FinishCompile() {
   compilation_state->PublishDetectedFeatures(
       isolate_, *compilation_state->detected_features());
 
-  // TODO(wasm): compiling wrappers should be made async as well.
-  DoSync<CompileWrappers>();
+  // TODO(bbudge) Allow deserialization without wrapper compilation, so we can
+  // just compile wrappers here.
+  if (compile_wrappers) {
+    DoSync<CompileWrappers>();
+  } else {
+    // TODO(wasm): compiling wrappers should be made async as well.
+    DoSync<AsyncCompileJob::FinishModule>();
+  }
 }
 
 void AsyncCompileJob::AsyncCompileFailed(Handle<Object> error_reason) {
@@ -2504,7 +2514,7 @@ class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
       job_->tiering_completed_ = true;
 
       // Degenerate case of an empty module.
-      job_->FinishCompile();
+      job_->FinishCompile(true);
       return;
     }
 
@@ -2524,10 +2534,14 @@ class AsyncCompileJob::PrepareAndStartCompile : public CompileStep {
                 if (job->DecrementAndCheckFinisherCount()) {
                   SaveContext saved_context(job->isolate());
                   job->isolate()->set_context(*job->native_context_);
-                  job->FinishCompile();
+                  job->FinishCompile(true);
                 }
                 return;
               case CompilationEvent::kFinishedTopTierCompilation:
+                // 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_ ||
@@ -2788,7 +2802,7 @@ void AsyncStreamingProcessor::OnFinishedStream(OwnedVector<uint8_t> bytes) {
       HandleScope scope(job_->isolate_);
       SaveContext saved_context(job_->isolate_);
       job_->isolate_->set_context(*job_->native_context_);
-      job_->FinishCompile();
+      job_->FinishCompile(true);
     }
   }
 }
@@ -2804,6 +2818,33 @@ void AsyncStreamingProcessor::OnAbort() {
   job_->Abort();
 }
 
+bool AsyncStreamingProcessor::Deserialize(Vector<const uint8_t> module_bytes,
+                                          Vector<const uint8_t> wire_bytes) {
+  HandleScope scope(job_->isolate_);
+  MaybeHandle<WasmModuleObject> result =
+      DeserializeNativeModule(job_->isolate_, module_bytes, wire_bytes);
+  if (result.is_null()) return false;
+
+  job_->module_object_ = result.ToHandleChecked();
+  {
+    DeferredHandleScope deferred(job_->isolate_);
+    job_->module_object_ = handle(*job_->module_object_, job_->isolate_);
+    job_->deferred_handles_.push_back(deferred.Detach());
+  }
+  job_->native_module_ = job_->module_object_->native_module();
+  auto owned_wire_bytes = OwnedVector<uint8_t>::Of(wire_bytes);
+  job_->wire_bytes_ = ModuleWireBytes(owned_wire_bytes.as_vector());
+  job_->native_module_->set_wire_bytes(std::move(owned_wire_bytes));
+
+  // Job should now behave as if it's fully compiled.
+  job_->tiering_completed_ = true;
+
+  SaveContext saved_context(job_->isolate_);
+  job_->isolate_->set_context(*job_->native_context_);
+  job_->FinishCompile(false);
+  return true;
+}
+
 void CompilationStateDeleter::operator()(
     CompilationState* compilation_state) const {
   delete compilation_state;

src/wasm/module-compiler.h

@@ -111,7 +111,7 @@ class AsyncCompileJob {
   }
   Counters* counters() const { return async_counters().get(); }
 
-  void FinishCompile();
+  void FinishCompile(bool compile_wrappers);
 
   void AsyncCompileFailed(Handle<Object> error_reason);
 

src/wasm/streaming-decoder.cc

@@ -12,6 +12,7 @@
 #include "src/wasm/decoder.h"
 #include "src/wasm/leb-helper.h"
 #include "src/wasm/module-decoder.h"
+#include "src/wasm/wasm-code-manager.h"
 #include "src/wasm/wasm-limits.h"
 #include "src/wasm/wasm-objects.h"
 #include "src/wasm/wasm-result.h"
@@ -26,7 +27,14 @@ namespace internal {
 namespace wasm {
 
 void StreamingDecoder::OnBytesReceived(Vector<const uint8_t> bytes) {
+  if (deserializing()) {
+    wire_bytes_for_deserializing_.insert(wire_bytes_for_deserializing_.end(),
+                                         bytes.begin(), bytes.end());
+    return;
+  }
+
   TRACE_STREAMING("OnBytesReceived(%zu bytes)\n", bytes.size());
+
   size_t current = 0;
   while (ok() && current < bytes.size()) {
     size_t num_bytes =
@@ -53,6 +61,19 @@ size_t StreamingDecoder::DecodingState::ReadBytes(StreamingDecoder* streaming,
 }
 
 void StreamingDecoder::Finish() {
+  if (deserializing()) {
+    Vector<const uint8_t> wire_bytes(wire_bytes_for_deserializing_.data(),
+                                     wire_bytes_for_deserializing_.size());
+    // Try to deserialize the module from wire bytes and module bytes.
+    if (processor_->Deserialize(compiled_module_bytes_, wire_bytes)) return;
+
+    // Deserialization failed. Restart decoding using |wire_bytes|.
+    compiled_module_bytes_ = {};
+    DCHECK(!deserializing());
+    OnBytesReceived(wire_bytes);
+    // The decoder has received all wire bytes; fall through and finish.
+  }
+
   TRACE_STREAMING("Finish\n");
   if (!ok()) {
     return;
@@ -89,6 +110,18 @@ void StreamingDecoder::Abort() {
   }
 }
 
+void StreamingDecoder::SetModuleCompiledCallback(
+    ModuleCompiledCallback callback) {
+  DCHECK_NULL(module_compiled_callback_);
+  module_compiled_callback_ = callback;
+}
+
+bool StreamingDecoder::SetCompiledModuleBytes(
+    Vector<const uint8_t> compiled_module_bytes) {
+  compiled_module_bytes_ = compiled_module_bytes;
+  return true;
+}
+
 // 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

@@ -49,6 +49,10 @@ class V8_EXPORT_PRIVATE StreamingProcessor {
   virtual void OnError(DecodeResult result) = 0;
   // Report the abortion of the stream.
   virtual void OnAbort() = 0;
+
+  // Attempt to deserialize the module. Supports embedder caching.
+  virtual bool Deserialize(Vector<const uint8_t> module_bytes,
+                           Vector<const uint8_t> wire_bytes) = 0;
 };
 
 // The StreamingDecoder takes a sequence of byte arrays, each received by a call
@@ -73,6 +77,18 @@ class V8_EXPORT_PRIVATE StreamingDecoder {
     ok_ = false;
   }
 
+  // Caching support.
+  // Sets the callback that is called after the module is fully compiled.
+  using ModuleCompiledCallback = std::function<void(Handle<WasmModuleObject>)>;
+  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_;
+  }
+
  private:
   // TODO(ahaas): Put the whole private state of the StreamingDecoder into the
   // cc file (PIMPL design pattern).
@@ -253,6 +269,8 @@ class V8_EXPORT_PRIVATE StreamingDecoder {
 
   uint32_t module_offset() const { return module_offset_; }
 
+  bool deserializing() const { return !compiled_module_bytes_.is_empty(); }
+
   std::unique_ptr<StreamingProcessor> processor_;
   bool ok_ = true;
   std::unique_ptr<DecodingState> state_;
@@ -261,6 +279,12 @@ class V8_EXPORT_PRIVATE StreamingDecoder {
   size_t total_size_ = 0;
   uint8_t next_section_id_ = kFirstSectionInModule;
 
+  // Caching support.
+  ModuleCompiledCallback module_compiled_callback_ = nullptr;
+  // We need wire bytes in an array for deserializing cached modules.
+  std::vector<uint8_t> wire_bytes_for_deserializing_;
+  Vector<const uint8_t> compiled_module_bytes_;
+
   DISALLOW_COPY_AND_ASSIGN(StreamingDecoder);
 };
 

src/wasm/wasm-js.cc

@@ -22,6 +22,7 @@
 #include "src/wasm/wasm-limits.h"
 #include "src/wasm/wasm-memory.h"
 #include "src/wasm/wasm-objects-inl.h"
+#include "src/wasm/wasm-serialization.h"
 
 using v8::internal::wasm::ErrorThrower;
 
@@ -58,6 +59,24 @@ class WasmStreaming::WasmStreamingImpl {
         Utils::OpenHandle(*exception.ToLocalChecked()));
   }
 
+  void SetModuleCompiledCallback(ModuleCompiledCallback callback,
+                                 intptr_t data) {
+    // Wrap the embedder callback here so we can also wrap the result as a
+    // Local<WasmCompiledModule> here.
+    streaming_decoder_->SetModuleCompiledCallback(
+        [callback, data](i::Handle<i::WasmModuleObject> module_object) {
+          callback(data, Local<WasmCompiledModule>::Cast(Utils::ToLocal(
+                             i::Handle<i::JSObject>::cast(module_object))));
+        });
+  }
+
+  bool SetCompiledModuleBytes(const uint8_t* bytes, size_t size) {
+    if (!i::wasm::IsSupportedVersion(reinterpret_cast<i::Isolate*>(isolate_),
+                                     {bytes, size}))
+      return false;
+    return streaming_decoder_->SetCompiledModuleBytes({bytes, size});
+  }
+
  private:
   Isolate* isolate_ = nullptr;
   std::shared_ptr<internal::wasm::StreamingDecoder> streaming_decoder_;
@@ -81,6 +100,15 @@ void WasmStreaming::Abort(MaybeLocal<Value> exception) {
   impl_->Abort(exception);
 }
 
+void WasmStreaming::SetModuleCompiledCallback(ModuleCompiledCallback callback,
+                                              intptr_t data) {
+  impl_->SetModuleCompiledCallback(callback, data);
+}
+
+bool WasmStreaming::SetCompiledModuleBytes(const uint8_t* bytes, size_t size) {
+  return impl_->SetCompiledModuleBytes(bytes, size);
+}
+
 // static
 std::shared_ptr<WasmStreaming> WasmStreaming::Unpack(Isolate* isolate,
                                                      Local<Value> value) {
@@ -1476,6 +1504,7 @@ void SetDummyInstanceTemplate(Isolate* isolate, Handle<JSFunction> fun) {
   fun->shared()->get_api_func_data()->set_instance_template(*instance_template);
 }
 
+// static
 void WasmJs::Install(Isolate* isolate, bool exposed_on_global_object) {
   Handle<JSGlobalObject> global = isolate->global_object();
   Handle<Context> context(global->native_context(), isolate);

src/wasm/wasm-js.h

@@ -10,6 +10,10 @@
 namespace v8 {
 namespace internal {
 
+namespace wasm {
+class StreamingDecoder;
+}
+
 // Exposes a WebAssembly API to JavaScript through the V8 API.
 class WasmJs {
  public:

src/wasm/wasm-objects.h

@@ -37,6 +37,7 @@ class JSArrayBuffer;
 class SeqOneByteString;
 class WasmDebugInfo;
 class WasmInstanceObject;
+class WasmModuleObject;
 
 template <class CppType>
 class Managed;

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

@@ -4,6 +4,7 @@
 
 #include "src/api-inl.h"
 #include "src/objects-inl.h"
+#include "src/objects/managed.h"
 #include "src/v8.h"
 #include "src/vector.h"
 
@@ -12,6 +13,9 @@
 #include "src/wasm/wasm-engine.h"
 #include "src/wasm/wasm-module-builder.h"
 #include "src/wasm/wasm-module.h"
+#include "src/wasm/wasm-objects-inl.h"
+#include "src/wasm/wasm-objects.h"
+#include "src/wasm/wasm-serialization.h"
 
 #include "test/cctest/cctest.h"
 
@@ -90,10 +94,15 @@ enum class CompilationState {
 
 class TestResolver : public CompilationResultResolver {
  public:
-  explicit TestResolver(CompilationState* state) : state_(state) {}
+  TestResolver(CompilationState* state,
+               std::shared_ptr<NativeModule>* native_module)
+      : state_(state), native_module_(native_module) {}
 
   void OnCompilationSucceeded(i::Handle<i::WasmModuleObject> module) override {
     *state_ = CompilationState::kFinished;
+    if (!module.is_null()) {
+      *native_module_ = module->managed_native_module()->get();
+    }
   }
 
   void OnCompilationFailed(i::Handle<i::Object> error_reason) override {
@@ -102,24 +111,29 @@ class TestResolver : public CompilationResultResolver {
 
  private:
   CompilationState* state_;
+  std::shared_ptr<NativeModule>* native_module_;
 };
 
 class StreamTester {
  public:
-  StreamTester() : zone_(&allocator_, "StreamTester") {
+  StreamTester()
+      : zone_(&allocator_, "StreamTester"),
+        internal_scope_(CcTest::i_isolate()) {
     v8::Isolate* isolate = CcTest::isolate();
     i::Isolate* i_isolate = CcTest::i_isolate();
-    i::HandleScope internal_scope(i_isolate);
 
     v8::Local<v8::Context> context = isolate->GetCurrentContext();
 
     stream_ = i_isolate->wasm_engine()->StartStreamingCompilation(
         i_isolate, kAllWasmFeatures, v8::Utils::OpenHandle(*context),
-        std::make_shared<TestResolver>(&state_));
+        std::make_shared<TestResolver>(&state_, &native_module_));
   }
 
   std::shared_ptr<StreamingDecoder> stream() { return stream_; }
 
+  // Compiled native module, valid after successful compile.
+  std::shared_ptr<NativeModule> native_module() { return native_module_; }
+
   // Run all compiler tasks, both foreground and background tasks.
   void RunCompilerTasks() {
     static_cast<MockPlatform*>(i::V8::GetCurrentPlatform())->ExecuteTasks();
@@ -137,12 +151,18 @@ class StreamTester {
 
   void FinishStream() { stream_->Finish(); }
 
+  void SetCompiledModuleBytes(const uint8_t* start, size_t length) {
+    stream_->SetCompiledModuleBytes(Vector<const uint8_t>(start, length));
+  }
+
   Zone* zone() { return &zone_; }
 
  private:
   AccountingAllocator allocator_;
   Zone zone_;
+  i::HandleScope internal_scope_;
   CompilationState state_ = CompilationState::kPending;
+  std::shared_ptr<NativeModule> native_module_;
   std::shared_ptr<StreamingDecoder> stream_;
 };
 }  // namespace
@@ -180,6 +200,28 @@ ZoneBuffer GetValidModuleBytes(Zone* zone) {
   return buffer;
 }
 
+// Create the same valid module as above and serialize it to test streaming
+// with compiled module caching.
+ZoneBuffer GetValidCompiledModuleBytes(Zone* zone, ZoneBuffer wire_bytes) {
+  // Use a tester to compile to a NativeModule.
+  StreamTester tester;
+  tester.OnBytesReceived(wire_bytes.begin(), wire_bytes.size());
+  tester.FinishStream();
+  tester.RunCompilerTasks();
+  CHECK(tester.IsPromiseFulfilled());
+  // Serialize the NativeModule.
+  std::shared_ptr<NativeModule> native_module = tester.native_module();
+  CHECK(native_module);
+  i::wasm::WasmSerializer serializer(
+      reinterpret_cast<i::Isolate*>(CcTest::i_isolate()), native_module.get());
+  size_t size = serializer.GetSerializedNativeModuleSize();
+  std::vector<byte> buffer(size);
+  CHECK(serializer.SerializeNativeModule({buffer.data(), size}));
+  ZoneBuffer result(zone, size);
+  result.write(buffer.data(), size);
+  return result;
+}
+
 // Test that all bytes arrive before doing any compilation. FinishStream is
 // called immediately.
 STREAM_TEST(TestAllBytesArriveImmediatelyStreamFinishesFirst) {
@@ -1010,6 +1052,40 @@ STREAM_TEST(TestModuleWithErrorAfterDataSection) {
   tester.RunCompilerTasks();
   CHECK(tester.IsPromiseRejected());
 }
+
+// Test that cached bytes work.
+STREAM_TEST(TestDeserializationBypassesCompilation) {
+  StreamTester tester;
+  ZoneBuffer wire_bytes = GetValidModuleBytes(tester.zone());
+  ZoneBuffer module_bytes =
+      GetValidCompiledModuleBytes(tester.zone(), wire_bytes);
+  tester.SetCompiledModuleBytes(module_bytes.begin(), module_bytes.size());
+  tester.OnBytesReceived(wire_bytes.begin(), wire_bytes.size());
+  tester.FinishStream();
+
+  tester.RunCompilerTasks();
+
+  CHECK(tester.IsPromiseFulfilled());
+}
+
+// Test that bad cached bytes don't cause compilation of wire bytes to fail.
+STREAM_TEST(TestDeserializationFails) {
+  StreamTester tester;
+  ZoneBuffer wire_bytes = GetValidModuleBytes(tester.zone());
+  ZoneBuffer module_bytes =
+      GetValidCompiledModuleBytes(tester.zone(), wire_bytes);
+  // corrupt header
+  byte first_byte = *module_bytes.begin();
+  module_bytes.patch_u8(0, first_byte + 1);
+  tester.SetCompiledModuleBytes(module_bytes.begin(), module_bytes.size());
+  tester.OnBytesReceived(wire_bytes.begin(), wire_bytes.size());
+  tester.FinishStream();
+
+  tester.RunCompilerTasks();
+
+  CHECK(tester.IsPromiseFulfilled());
+}
+
 #undef STREAM_TEST
 
 }  // namespace wasm

test/unittests/wasm/streaming-decoder-unittest.cc

@@ -66,6 +66,11 @@ class MockStreamingProcessor : public StreamingProcessor {
 
   void OnAbort() override {}
 
+  bool Deserialize(Vector<const uint8_t> module_bytes,
+                   Vector<const uint8_t> wire_bytes) override {
+    return false;
+  };
+
   size_t num_sections() const { return num_sections_; }
   size_t num_functions() const { return num_functions_; }
   bool ok() const { return ok_; }