[wasm] Cache streaming compilation result

Before compiling the code section, check whether the
bytes received so far match a cached module. If they do, delay
compilation until we receive the full bytes, since we are likely to find
a cache entry for them.

R=clemensb@chromium.org

Bug: v8:6847
Change-Id: Ie5170d1274da3da6d52ff1b408abc7cb441bbe3c
Reviewed-on: https://chromium-review.googlesource.com/c/v8/v8/+/2002823
Commit-Queue: Thibaud Michaud <thibaudm@chromium.org>
Reviewed-by: Clemens Backes <clemensb@chromium.org>
Cr-Commit-Position: refs/heads/master@{#66000}

src/wasm/module-compiler.h

@@ -149,9 +149,12 @@ class AsyncCompileJob {
 
   void CreateNativeModule(std::shared_ptr<const WasmModule> module,
                           size_t code_size_estimate);
+  // Return true for cache hit, false for cache miss.
+  bool GetOrCreateNativeModule(std::shared_ptr<const WasmModule> module,
+                               size_t code_size_estimate);
   void PrepareRuntimeObjects();
 
-  void FinishCompile();
+  void FinishCompile(bool is_after_cache_hit);
 
   void DecodeFailed(const WasmError&);
   void AsyncCompileFailed();

src/wasm/wasm-engine.cc

@@ -4,6 +4,7 @@
 
 #include "src/wasm/wasm-engine.h"
 
+#include "src/base/functional.h"
 #include "src/base/platform/time.h"
 #include "src/diagnostics/code-tracer.h"
 #include "src/diagnostics/compilation-statistics.h"
@@ -130,18 +131,27 @@ std::shared_ptr<NativeModule> NativeModuleCache::MaybeGetNativeModule(
     ModuleOrigin origin, Vector<const uint8_t> wire_bytes) {
   if (origin != kWasmOrigin) return nullptr;
   base::MutexGuard lock(&mutex_);
+  size_t prefix_hash = PrefixHash(wire_bytes);
+  NativeModuleCache::Key key{prefix_hash, wire_bytes};
   while (true) {
-    auto it = map_.find(wire_bytes);
+    auto it = map_.find(key);
     if (it == map_.end()) {
+      // Even though this exact key is not in the cache, there might be a
+      // matching prefix hash indicating that a streaming compilation is
+      // currently compiling a module with the same prefix. {OnFinishedStream}
+      // happens on the main thread too, so waiting for streaming compilation to
+      // finish would create a deadlock. Instead, compile the module twice and
+      // handle the conflict in {UpdateNativeModuleCache}.
+
       // Insert a {nullopt} entry to let other threads know that this
       // {NativeModule} is already being created on another thread.
-      map_.emplace(wire_bytes, base::nullopt);
+      auto p = map_.emplace(key, base::nullopt);
+      USE(p);
+      DCHECK(p.second);
       return nullptr;
     }
-    auto maybe_native_module = it->second;
-    if (maybe_native_module.has_value()) {
-      auto weak_ptr = maybe_native_module.value();
-      if (auto shared_native_module = weak_ptr.lock()) {
+    if (it->second.has_value()) {
+      if (auto shared_native_module = it->second.value().lock()) {
         return shared_native_module;
       }
     }
@@ -149,28 +159,62 @@ std::shared_ptr<NativeModule> NativeModuleCache::MaybeGetNativeModule(
   }
 }
 
-void NativeModuleCache::Update(std::shared_ptr<NativeModule> native_module,
-                               bool error) {
+bool NativeModuleCache::GetStreamingCompilationOwnership(size_t prefix_hash) {
+  base::MutexGuard lock(&mutex_);
+  auto it = map_.lower_bound(Key{prefix_hash, {}});
+  if (it->first.prefix_hash == prefix_hash) {
+    return false;
+  }
+  Key key{prefix_hash, {}};
+  return map_.emplace(key, base::nullopt).second;
+}
+
+void NativeModuleCache::StreamingCompilationFailed(size_t prefix_hash) {
+  base::MutexGuard lock(&mutex_);
+  map_.erase({prefix_hash, {}});
+  cache_cv_.NotifyAll();
+}
+
+std::shared_ptr<NativeModule> NativeModuleCache::Update(
+    std::shared_ptr<NativeModule> native_module, bool error) {
   DCHECK_NOT_NULL(native_module);
-  if (native_module->module()->origin != kWasmOrigin) return;
+  if (native_module->module()->origin != kWasmOrigin) return native_module;
   Vector<const uint8_t> wire_bytes = native_module->wire_bytes();
+  size_t prefix_hash = PrefixHash(native_module->wire_bytes());
   base::MutexGuard lock(&mutex_);
-  auto it = map_.find(wire_bytes);
-  DCHECK_NE(it, map_.end());
-  DCHECK(!it->second.has_value());
-  // The lifetime of the temporary entry's bytes is unknown. Use the new native
-  // module's owned copy of the bytes for the key instead.
+  map_.erase(Key{prefix_hash, {}});
+  const Key key{prefix_hash, wire_bytes};
+  auto it = map_.find(key);
+  if (it != map_.end()) {
+    if (it->second.has_value()) {
+      auto conflicting_module = it->second.value().lock();
+      if (conflicting_module != nullptr) {
+        return conflicting_module;
+      }
+    }
     map_.erase(it);
+  }
   if (!error) {
-    map_.emplace(wire_bytes, base::Optional<std::weak_ptr<NativeModule>>(
-                                 std::move(native_module)));
+    DCHECK_LT(0, native_module->wire_bytes().length());
+    // The key now points to the new native module's owned copy of the bytes,
+    // so that it stays valid until the native module is freed and erased from
+    // the map.
+    auto p = map_.emplace(
+        key, base::Optional<std::weak_ptr<NativeModule>>(native_module));
+    USE(p);
+    DCHECK(p.second);
   }
   cache_cv_.NotifyAll();
+  return native_module;
 }
 
 void NativeModuleCache::Erase(NativeModule* native_module) {
+  if (native_module->module()->origin != kWasmOrigin) return;
+  // Happens in some tests where bytes are set directly.
+  if (native_module->wire_bytes().length() == 0) return;
   base::MutexGuard lock(&mutex_);
-  auto cache_it = map_.find(native_module->wire_bytes());
+  size_t prefix_hash = PrefixHash(native_module->wire_bytes());
+  auto cache_it = map_.find(Key{prefix_hash, native_module->wire_bytes()});
   // Not all native modules are stored in the cache currently. In particular
   // streaming compilation and asmjs compilation results are not. So make
   // sure that we only delete existing and expired entries.
@@ -183,13 +227,37 @@ void NativeModuleCache::Erase(NativeModule* native_module) {
   }
 }
 
-size_t NativeModuleCache::WireBytesHasher::operator()(
-    const Vector<const uint8_t>& bytes) const {
+// static
+size_t NativeModuleCache::WireBytesHash(Vector<const uint8_t> bytes) {
   return StringHasher::HashSequentialString(
       reinterpret_cast<const char*>(bytes.begin()), bytes.length(),
       kZeroHashSeed);
 }
 
+// static
+size_t NativeModuleCache::PrefixHash(Vector<const uint8_t> wire_bytes) {
+  // Compute the hash as a combined hash of the sections up to the code section
+  // header, to mirror the way streaming compilation does it.
+  Decoder decoder(wire_bytes.begin(), wire_bytes.end());
+  decoder.consume_bytes(8, "module header");
+  size_t hash = NativeModuleCache::WireBytesHash(wire_bytes.SubVector(0, 8));
+  SectionCode section_id = SectionCode::kUnknownSectionCode;
+  while (decoder.ok() && decoder.more()) {
+    section_id = static_cast<SectionCode>(decoder.consume_u8());
+    uint32_t section_size = decoder.consume_u32v("section size");
+    if (section_id == SectionCode::kCodeSectionCode) {
+      hash = base::hash_combine(hash, section_size);
+      break;
+    }
+    const uint8_t* payload_start = decoder.pc();
+    decoder.consume_bytes(section_size, "section payload");
+    size_t section_hash = NativeModuleCache::WireBytesHash(
+        Vector<const uint8_t>(payload_start, section_size));
+    hash = base::hash_combine(hash, section_hash);
+  }
+  return hash;
+}
+
 struct WasmEngine::CurrentGCInfo {
   explicit CurrentGCInfo(int8_t gc_sequence_index)
       : gc_sequence_index(gc_sequence_index) {
@@ -762,8 +830,19 @@ std::shared_ptr<NativeModule> WasmEngine::MaybeGetNativeModule(
 }
 
 void WasmEngine::UpdateNativeModuleCache(
-    std::shared_ptr<NativeModule> native_module, bool error) {
-  native_module_cache_.Update(native_module, error);
+    bool error, std::shared_ptr<NativeModule>* native_module) {
+  // Pass {native_module} by value here to keep it alive until at least after
+  // we returned from {Update}. Otherwise, we might {Erase} it inside {Update}
+  // which would lock the mutex twice.
+  *native_module = native_module_cache_.Update(*native_module, error);
+}
+
+bool WasmEngine::GetStreamingCompilationOwnership(size_t prefix_hash) {
+  return native_module_cache_.GetStreamingCompilationOwnership(prefix_hash);
+}
+
+void WasmEngine::StreamingCompilationFailed(size_t prefix_hash) {
+  native_module_cache_.StreamingCompilationFailed(prefix_hash);
 }
 
 void WasmEngine::FreeNativeModule(NativeModule* native_module) {

src/wasm/wasm-engine.h

@@ -5,6 +5,8 @@
 #ifndef V8_WASM_WASM_ENGINE_H_
 #define V8_WASM_WASM_ENGINE_H_
 
+#include <algorithm>
+#include <map>
 #include <memory>
 #include <unordered_map>
 #include <unordered_set>
@@ -51,15 +53,43 @@ class V8_EXPORT_PRIVATE InstantiationResultResolver {
 // Native modules cached by their wire bytes.
 class NativeModuleCache {
  public:
-  struct WireBytesHasher {
-    size_t operator()(const Vector<const uint8_t>& bytes) const;
+  struct Key {
+    // Store the prefix hash as part of the key for faster lookup, and to
+    // quickly check existing prefixes for streaming compilation.
+    size_t prefix_hash;
+    Vector<const uint8_t> bytes;
+
+    bool operator==(const Key& other) const {
+      bool eq = bytes == other.bytes;
+      DCHECK_IMPLIES(eq, prefix_hash == other.prefix_hash);
+      return eq;
+    }
+
+    bool operator<(const Key& other) const {
+      if (prefix_hash != other.prefix_hash) {
+        return prefix_hash < other.prefix_hash;
+      }
+      return std::lexicographical_compare(
+          bytes.begin(), bytes.end(), other.bytes.begin(), other.bytes.end());
+    }
   };
 
   std::shared_ptr<NativeModule> MaybeGetNativeModule(
       ModuleOrigin origin, Vector<const uint8_t> wire_bytes);
-  void Update(std::shared_ptr<NativeModule> native_module, bool error);
+  bool GetStreamingCompilationOwnership(size_t prefix_hash);
+  void StreamingCompilationFailed(size_t prefix_hash);
+  std::shared_ptr<NativeModule> Update(
+      std::shared_ptr<NativeModule> native_module, bool error);
   void Erase(NativeModule* native_module);
 
+  static size_t WireBytesHash(Vector<const uint8_t> bytes);
+
+  // Hash the wire bytes up to the code section header. Used as a heuristic to
+  // avoid streaming compilation of modules that are likely already in the
+  // cache. See {GetStreamingCompilationOwnership}. Assumes that the bytes have
+  // already been validated.
+  static size_t PrefixHash(Vector<const uint8_t> wire_bytes);
+
  private:
   // Each key points to the corresponding native module's wire bytes, so they
   // should always be valid as long as the native module is alive.  When
@@ -72,10 +102,7 @@ class NativeModuleCache {
   // before trying to get it from the cache.
   // By contrast, an expired {weak_ptr} indicates that the native module died
   // and will soon be cleaned up from the cache.
-  std::unordered_map<Vector<const uint8_t>,
-                     base::Optional<std::weak_ptr<NativeModule>>,
-                     WireBytesHasher>
-      map_;
+  std::map<Key, base::Optional<std::weak_ptr<NativeModule>>> map_;
 
   base::Mutex mutex_;
 
@@ -226,21 +253,39 @@ class V8_EXPORT_PRIVATE WasmEngine {
       Isolate* isolate, const WasmFeatures& enabled_features,
       std::shared_ptr<const WasmModule> module, size_t code_size_estimate);
 
-  // Try getting a cached {NativeModule}. The {wire_bytes}' underlying array
-  // should be valid at least until the next call to {UpdateNativeModuleCache}.
-  // Return nullptr if no {NativeModule} exists for these bytes. In this case,
-  // an empty entry is added to let other threads know that a {NativeModule} for
-  // these bytes is currently being created. The caller should eventually call
-  // {UpdateNativeModuleCache} to update the entry and wake up other threads.
+  // Try getting a cached {NativeModule}, or get ownership for its creation.
+  // Return {nullptr} if no {NativeModule} exists for these bytes. In this case,
+  // a {nullopt} entry is added to let other threads know that a {NativeModule}
+  // for these bytes is currently being created. The caller should eventually
+  // call {UpdateNativeModuleCache} to update the entry and wake up other
+  // threads. The {wire_bytes}' underlying array should be valid at least until
+  // the call to {UpdateNativeModuleCache}.
   std::shared_ptr<NativeModule> MaybeGetNativeModule(
       ModuleOrigin origin, Vector<const uint8_t> wire_bytes);
 
-  // Update the temporary entry inserted by {MaybeGetNativeModule}.
-  // If {error} is true, the entry is erased. Otherwise the entry is updated to
-  // match the {native_module} argument. Wake up threads waiting for this native
+  // Replace the temporary {nullopt} with the new native module, or
+  // erase it if any error occurred. Wake up blocked threads waiting for this
   // module.
-  void UpdateNativeModuleCache(std::shared_ptr<NativeModule> native_module,
-                               bool error);
+  // To avoid a deadlock on the main thread between synchronous and streaming
+  // compilation, two compilation jobs might compile the same native module at
+  // the same time. In this case the first call to {UpdateNativeModuleCache}
+  // will insert the native module in the cache, and the last call will discard
+  // its {native_module} argument and replace it with the existing entry.
+  void UpdateNativeModuleCache(bool error,
+                               std::shared_ptr<NativeModule>* native_module);
+
+  // Register this prefix hash for a streaming compilation job.
+  // If the hash is not in the cache yet, the function returns true and the
+  // caller owns the compilation of this module.
+  // Otherwise another compilation job is currently preparing or has already
+  // prepared a module with the same prefix hash. The caller should wait until
+  // the stream is finished and call {MaybeGetNativeModule} to either get the
+  // module from the cache or get ownership for the compilation of these bytes.
+  bool GetStreamingCompilationOwnership(size_t prefix_hash);
+
+  // Remove the prefix hash from the cache when compilation failed. If
+  // compilation succeeded, {UpdateNativeModuleCache} should be called instead.
+  void StreamingCompilationFailed(size_t prefix_hash);
 
   void FreeNativeModule(NativeModule*);
 

src/wasm/wasm-serialization.cc

@@ -637,7 +637,7 @@ MaybeHandle<WasmModuleObject> DeserializeNativeModule(
 
     Reader reader(data + WasmSerializer::kHeaderSize);
     bool error = !deserializer.Read(&reader);
-    wasm_engine->UpdateNativeModuleCache(shared_native_module, error);
+    wasm_engine->UpdateNativeModuleCache(error, &shared_native_module);
     if (error) return {};
   }
 

test/cctest/cctest.status

@@ -455,7 +455,7 @@
   'test-api-wasm/WasmStreaming*': [SKIP],
   'test-backing-store/Run_WasmModule_Buffer_Externalized_Regression_UseAfterFree': [SKIP],
   'test-c-wasm-entry/*': [SKIP],
-  'test-compilation-cache/TestAsyncCache': [SKIP],
+  'test-compilation-cache/*': [SKIP],
   'test-jump-table-assembler/*': [SKIP],
   'test-grow-memory/*': [SKIP],
   'test-run-wasm-64/*': [SKIP],

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

@@ -5,6 +5,7 @@
 #include "src/api/api-inl.h"
 #include "src/init/v8.h"
 
+#include "src/wasm/streaming-decoder.h"
 #include "src/wasm/wasm-code-manager.h"
 #include "src/wasm/wasm-engine.h"
 #include "src/wasm/wasm-module-builder.h"
@@ -43,6 +44,35 @@ class TestResolver : public CompilationResultResolver {
   std::atomic<int>* pending_;
 };
 
+class StreamTester {
+ public:
+  explicit StreamTester(std::shared_ptr<TestResolver> test_resolver)
+      : internal_scope_(CcTest::i_isolate()), test_resolver_(test_resolver) {
+    i::Isolate* i_isolate = CcTest::i_isolate();
+
+    Handle<Context> context = i_isolate->native_context();
+
+    stream_ = i_isolate->wasm_engine()->StartStreamingCompilation(
+        i_isolate, WasmFeatures::All(), context,
+        "WebAssembly.compileStreaming()", test_resolver_);
+  }
+
+  void OnBytesReceived(const uint8_t* start, size_t length) {
+    stream_->OnBytesReceived(Vector<const uint8_t>(start, length));
+  }
+
+  void FinishStream() { stream_->Finish(); }
+
+  void SetCompiledModuleBytes(const uint8_t* start, size_t length) {
+    stream_->SetCompiledModuleBytes(Vector<const uint8_t>(start, length));
+  }
+
+ private:
+  i::HandleScope internal_scope_;
+  std::shared_ptr<StreamingDecoder> stream_;
+  std::shared_ptr<TestResolver> test_resolver_;
+};
+
 // Create a valid module such that the bytes depend on {n}.
 ZoneBuffer GetValidModuleBytes(Zone* zone, int n) {
   ZoneBuffer buffer(zone);
@@ -57,11 +87,51 @@ ZoneBuffer GetValidModuleBytes(Zone* zone, int n) {
   return buffer;
 }
 
+std::shared_ptr<NativeModule> SyncCompile(Vector<const uint8_t> bytes) {
+  ErrorThrower thrower(CcTest::i_isolate(), "Test");
+  auto enabled_features = WasmFeatures::FromIsolate(CcTest::i_isolate());
+  auto wire_bytes = ModuleWireBytes(bytes.begin(), bytes.end());
+  Handle<WasmModuleObject> module =
+      CcTest::i_isolate()
+          ->wasm_engine()
+          ->SyncCompile(CcTest::i_isolate(), enabled_features, &thrower,
+                        wire_bytes)
+          .ToHandleChecked();
+  return module->shared_native_module();
+}
+
+// Shared prefix.
+constexpr uint8_t kPrefix[] = {
+    WASM_MODULE_HEADER,                // module header
+    kTypeSectionCode,                  // section code
+    U32V_1(1 + SIZEOF_SIG_ENTRY_v_v),  // section size
+    U32V_1(1),                         // type count
+    SIG_ENTRY_v_v,                     // signature entry
+    kFunctionSectionCode,              // section code
+    U32V_1(2),                         // section size
+    U32V_1(1),                         // functions count
+    0,                                 // signature index
+    kCodeSectionCode,                  // section code
+    U32V_1(7),                         // section size
+    U32V_1(1),                         // functions count
+    5,                                 // body size
+};
+
+constexpr uint8_t kFunctionA[] = {
+    U32V_1(0), kExprI32Const, U32V_1(0), kExprDrop, kExprEnd,
+};
+constexpr uint8_t kFunctionB[] = {
+    U32V_1(0), kExprI32Const, U32V_1(1), kExprDrop, kExprEnd,
+};
+
+constexpr size_t kPrefixSize = arraysize(kPrefix);
+constexpr size_t kFunctionSize = arraysize(kFunctionA);
+
 }  // namespace
 
 TEST(TestAsyncCache) {
   CcTest::InitializeVM();
-  i::HandleScope internal_scope_(CcTest::i_isolate());
+  i::HandleScope internal_scope(CcTest::i_isolate());
   AccountingAllocator allocator;
   Zone zone(&allocator, "CompilationCacheTester");
 
@@ -95,6 +165,74 @@ TEST(TestAsyncCache) {
   CHECK_NE(resolverA1->native_module(), resolverB->native_module());
 }
 
+TEST(TestStreamingCache) {
+  CcTest::InitializeVM();
+
+  std::atomic<int> pending(3);
+  auto resolverA1 = std::make_shared<TestResolver>(&pending);
+  auto resolverA2 = std::make_shared<TestResolver>(&pending);
+  auto resolverB = std::make_shared<TestResolver>(&pending);
+
+  StreamTester testerA1(resolverA1);
+  StreamTester testerA2(resolverA2);
+  StreamTester testerB(resolverB);
+
+  // Start receiving kPrefix bytes.
+  testerA1.OnBytesReceived(kPrefix, kPrefixSize);
+  testerA2.OnBytesReceived(kPrefix, kPrefixSize);
+  testerB.OnBytesReceived(kPrefix, kPrefixSize);
+
+  // Receive function bytes and start streaming compilation.
+  testerA1.OnBytesReceived(kFunctionA, kFunctionSize);
+  testerA1.FinishStream();
+  testerA2.OnBytesReceived(kFunctionA, kFunctionSize);
+  testerA2.FinishStream();
+  testerB.OnBytesReceived(kFunctionB, kFunctionSize);
+  testerB.FinishStream();
+
+  while (pending > 0) {
+    v8::platform::PumpMessageLoop(i::V8::GetCurrentPlatform(),
+                                  CcTest::isolate());
+  }
+
+  std::shared_ptr<NativeModule> native_module_A1 = resolverA1->native_module();
+  std::shared_ptr<NativeModule> native_module_A2 = resolverA2->native_module();
+  std::shared_ptr<NativeModule> native_module_B = resolverB->native_module();
+  CHECK_EQ(native_module_A1, native_module_A2);
+  CHECK_NE(native_module_A1, native_module_B);
+}
+
+TEST(TestStreamingAndSyncCache) {
+  CcTest::InitializeVM();
+
+  std::atomic<int> pending(1);
+  auto resolver = std::make_shared<TestResolver>(&pending);
+  StreamTester tester(resolver);
+
+  tester.OnBytesReceived(kPrefix, kPrefixSize);
+
+  // Compile the same module synchronously to make sure we don't deadlock
+  // waiting for streaming compilation to finish.
+  auto full_bytes = OwnedVector<uint8_t>::New(kPrefixSize + kFunctionSize);
+  memcpy(full_bytes.begin(), kPrefix, kPrefixSize);
+  memcpy(full_bytes.begin() + kPrefixSize, kFunctionA, kFunctionSize);
+  auto native_module_sync = SyncCompile(full_bytes.as_vector());
+
+  // Streaming compilation should just discard its native module now and use the
+  // one inserted in the cache by sync compilation.
+  tester.OnBytesReceived(kFunctionA, kFunctionSize);
+  tester.FinishStream();
+
+  while (pending > 0) {
+    v8::platform::PumpMessageLoop(i::V8::GetCurrentPlatform(),
+                                  CcTest::isolate());
+  }
+
+  std::shared_ptr<NativeModule> native_module_streaming =
+      resolver->native_module();
+  CHECK_EQ(native_module_streaming, native_module_sync);
+}
+
 }  // namespace wasm
 }  // namespace internal
 }  // namespace v8

test/cctest/wasm/wasm-run-utils.cc

@@ -69,6 +69,12 @@ TestingModuleBuilder::TestingModuleBuilder(
   }
 }
 
+TestingModuleBuilder::~TestingModuleBuilder() {
+  // When the native module dies and is erased from the cache, it is expected to
+  // have either valid bytes or no bytes at all.
+  native_module_->SetWireBytes({});
+}
+
 byte* TestingModuleBuilder::AddMemory(uint32_t size, SharedFlag shared) {
   CHECK(!test_module_->has_memory);
   CHECK_NULL(mem_start_);

test/cctest/wasm/wasm-run-utils.h

@@ -89,6 +89,7 @@ class TestingModuleBuilder {
  public:
   TestingModuleBuilder(Zone*, ManuallyImportedJSFunction*, ExecutionTier,
                        RuntimeExceptionSupport, LowerSimd);
+  ~TestingModuleBuilder();
 
   void ChangeOriginToAsmjs() { test_module_->origin = kAsmJsSloppyOrigin; }
 

test/mjsunit/mjsunit.status

@@ -157,6 +157,9 @@
   # OOM with too many isolates/memory objects (https://crbug.com/1010272)
   # Predictable tests fail due to race between postMessage and GrowMemory
   'regress/wasm/regress-1010272': [PASS, NO_VARIANTS, ['system == android', SKIP], ['predictable', SKIP]],
+
+  # https://crbug.com/v8/10126
+  'regress/wasm/regress-789952': [SKIP]
 }],  # ALWAYS
 
 ##############################################################################