Bug 1487113 - Use alt-data to cache stream-compiled WebAssembly modules. r=necko-reviewers,valentin,dragana

Depends on D117360

Differential Revision: https://phabricator.services.mozilla.com/D26731

dom/fetch/FetchDriver.cpp

@@ -821,11 +821,11 @@ nsresult FetchDriver::HttpFetch(
   } else {
     // Integrity check cannot be done on alt-data yet.
     if (mRequest->GetIntegrity().IsEmpty()) {
+      MOZ_ASSERT(!FetchUtil::WasmAltDataType.IsEmpty());
       nsCOMPtr<nsICacheInfoChannel> cic = do_QueryInterface(chan);
       if (cic) {
         cic->PreferAlternativeDataType(
-            nsLiteralCString(WASM_ALT_DATA_TYPE_V1),
-            nsLiteralCString(WASM_CONTENT_TYPE),
+            FetchUtil::WasmAltDataType, nsLiteralCString(WASM_CONTENT_TYPE),
             nsICacheInfoChannel::PreferredAlternativeDataDeliveryType::
                 SERIALIZE);
       }
@@ -1079,8 +1079,8 @@ FetchDriver::OnStartRequest(nsIRequest* aRequest) {
       }
     } else if (!cic->PreferredAlternativeDataTypes().IsEmpty()) {
       MOZ_ASSERT(cic->PreferredAlternativeDataTypes().Length() == 1);
-      MOZ_ASSERT(cic->PreferredAlternativeDataTypes()[0].type().EqualsLiteral(
-          WASM_ALT_DATA_TYPE_V1));
+      MOZ_ASSERT(cic->PreferredAlternativeDataTypes()[0].type().Equals(
+          FetchUtil::WasmAltDataType));
       MOZ_ASSERT(
           cic->PreferredAlternativeDataTypes()[0].contentType().EqualsLiteral(
               WASM_CONTENT_TYPE));

dom/fetch/FetchUtil.cpp

@@ -10,12 +10,15 @@
 #include "nsCRT.h"
 #include "nsError.h"
 #include "nsIAsyncInputStream.h"
+#include "nsICloneableInputStream.h"
 #include "nsIHttpChannel.h"
 #include "nsNetUtil.h"
 #include "nsStreamUtils.h"
 #include "nsString.h"
+#include "js/BuildId.h"
 #include "mozilla/dom/Document.h"
 
+#include "mozilla/ClearOnShutdown.h"
 #include "mozilla/dom/DOMException.h"
 #include "mozilla/dom/InternalRequest.h"
 #include "mozilla/dom/Response.h"
@@ -162,6 +165,41 @@ nsresult FetchUtil::SetRequestReferrer(nsIPrincipal* aPrincipal, Document* aDoc,
   return NS_OK;
 }
 
+class StoreOptimizedEncodingRunnable final : public Runnable {
+  nsMainThreadPtrHandle<nsICacheInfoChannel> mCache;
+  JS::UniqueOptimizedEncodingBytes mBytes;
+
+ public:
+  StoreOptimizedEncodingRunnable(
+      nsMainThreadPtrHandle<nsICacheInfoChannel>&& aCache,
+      JS::UniqueOptimizedEncodingBytes&& aBytes)
+      : Runnable("StoreOptimizedEncodingRunnable"),
+        mCache(std::move(aCache)),
+        mBytes(std::move(aBytes)) {}
+
+  NS_IMETHOD Run() override {
+    nsresult rv;
+
+    nsCOMPtr<nsIAsyncOutputStream> stream;
+    rv = mCache->OpenAlternativeOutputStream(
+        FetchUtil::WasmAltDataType, mBytes->length(), getter_AddRefs(stream));
+    if (NS_FAILED(rv)) {
+      return rv;
+    }
+
+    auto closeStream = MakeScopeExit([&]() { stream->CloseWithStatus(rv); });
+
+    uint32_t written;
+    rv = stream->Write((char*)mBytes->begin(), mBytes->length(), &written);
+    if (NS_FAILED(rv)) {
+      return rv;
+    }
+
+    MOZ_RELEASE_ASSERT(mBytes->length() == written);
+    return NS_OK;
+  };
+};
+
 class WindowStreamOwner final : public nsIObserver,
                                 public nsSupportsWeakReference {
   // Read from any thread but only set/cleared on the main thread. The lifecycle
@@ -302,17 +340,25 @@ class WorkerStreamOwner final {
   RefPtr<WeakWorkerRef> mWorkerRef;
 };
 
-class JSStreamConsumer final : public nsIInputStreamCallback {
+class JSStreamConsumer final : public nsIInputStreamCallback,
+                               public JS::OptimizedEncodingListener {
   nsCOMPtr<nsIEventTarget> mOwningEventTarget;
   RefPtr<WindowStreamOwner> mWindowStreamOwner;
   RefPtr<WorkerStreamOwner> mWorkerStreamOwner;
+  nsMainThreadPtrHandle<nsICacheInfoChannel> mCache;
+  const bool mOptimizedEncoding;
+  Vector<uint8_t> mOptimizedEncodingBytes;
   JS::StreamConsumer* mConsumer;
   bool mConsumerAborted;
 
   JSStreamConsumer(already_AddRefed<WindowStreamOwner> aWindowStreamOwner,
-                   nsIGlobalObject* aGlobal, JS::StreamConsumer* aConsumer)
+                   nsIGlobalObject* aGlobal, JS::StreamConsumer* aConsumer,
+                   nsMainThreadPtrHandle<nsICacheInfoChannel>&& aCache,
+                   bool aOptimizedEncoding)
       : mOwningEventTarget(aGlobal->EventTargetFor(TaskCategory::Other)),
         mWindowStreamOwner(aWindowStreamOwner),
+        mCache(std::move(aCache)),
+        mOptimizedEncoding(aOptimizedEncoding),
         mConsumer(aConsumer),
         mConsumerAborted(false) {
     MOZ_DIAGNOSTIC_ASSERT(mWindowStreamOwner);
@@ -320,9 +366,13 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
   }
 
   JSStreamConsumer(RefPtr<WorkerStreamOwner> aWorkerStreamOwner,
-                   nsIGlobalObject* aGlobal, JS::StreamConsumer* aConsumer)
+                   nsIGlobalObject* aGlobal, JS::StreamConsumer* aConsumer,
+                   nsMainThreadPtrHandle<nsICacheInfoChannel>&& aCache,
+                   bool aOptimizedEncoding)
       : mOwningEventTarget(aGlobal->EventTargetFor(TaskCategory::Other)),
         mWorkerStreamOwner(std::move(aWorkerStreamOwner)),
+        mCache(std::move(aCache)),
+        mOptimizedEncoding(aOptimizedEncoding),
         mConsumer(aConsumer),
         mConsumerAborted(false) {
     MOZ_DIAGNOSTIC_ASSERT(mWorkerStreamOwner);
@@ -352,11 +402,19 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
     JSStreamConsumer* self = reinterpret_cast<JSStreamConsumer*>(aClosure);
     MOZ_DIAGNOSTIC_ASSERT(!self->mConsumerAborted);
 
-    // This callback can be called on any thread which is explicitly allowed by
-    // this particular JS API call.
-    if (!self->mConsumer->consumeChunk((const uint8_t*)aFromSegment, aCount)) {
-      self->mConsumerAborted = true;
-      return NS_ERROR_UNEXPECTED;
+    if (self->mOptimizedEncoding) {
+      if (!self->mOptimizedEncodingBytes.append((const uint8_t*)aFromSegment,
+                                                aCount)) {
+        return NS_ERROR_UNEXPECTED;
+      }
+    } else {
+      // This callback can be called on any thread which is explicitly allowed
+      // by this particular JS API call.
+      if (!self->mConsumer->consumeChunk((const uint8_t*)aFromSegment,
+                                         aCount)) {
+        self->mConsumerAborted = true;
+        return NS_ERROR_UNEXPECTED;
+      }
     }
 
     *aWriteCount = aCount;
@@ -366,9 +424,10 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
  public:
   NS_DECL_THREADSAFE_ISUPPORTS
 
-  static bool Start(nsCOMPtr<nsIInputStream>&& aStream,
-                    JS::StreamConsumer* aConsumer, nsIGlobalObject* aGlobal,
-                    WorkerPrivate* aMaybeWorker) {
+  static bool Start(nsCOMPtr<nsIInputStream> aStream, nsIGlobalObject* aGlobal,
+                    WorkerPrivate* aMaybeWorker, JS::StreamConsumer* aConsumer,
+                    nsMainThreadPtrHandle<nsICacheInfoChannel>&& aCache,
+                    bool aOptimizedEncoding) {
     nsCOMPtr<nsIAsyncInputStream> asyncStream;
     nsresult rv = NS_MakeAsyncNonBlockingInputStream(
         aStream.forget(), getter_AddRefs(asyncStream));
@@ -384,7 +443,8 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
         return false;
       }
 
-      consumer = new JSStreamConsumer(std::move(owner), aGlobal, aConsumer);
+      consumer = new JSStreamConsumer(std::move(owner), aGlobal, aConsumer,
+                                      std::move(aCache), aOptimizedEncoding);
     } else {
       RefPtr<WindowStreamOwner> owner =
           WindowStreamOwner::Create(asyncStream, aGlobal);
@@ -392,7 +452,8 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
         return false;
       }
 
-      consumer = new JSStreamConsumer(owner.forget(), aGlobal, aConsumer);
+      consumer = new JSStreamConsumer(owner.forget(), aGlobal, aConsumer,
+                                      std::move(aCache), aOptimizedEncoding);
     }
 
     // This AsyncWait() creates a ref-cycle between asyncStream and consumer:
@@ -421,7 +482,16 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
     }
 
     if (rv == NS_BASE_STREAM_CLOSED) {
-      mConsumer->streamEnd();
+      if (mOptimizedEncoding) {
+        mConsumer->consumeOptimizedEncoding(mOptimizedEncodingBytes.begin(),
+                                            mOptimizedEncodingBytes.length());
+      } else {
+        // If there is cache entry associated with this stream, then listen for
+        // an optimized encoding so we can store it in the alt data. By JS API
+        // contract, the compilation process will hold a refcount to 'this'
+        // until it's done, optionally calling storeOptimizedEncoding().
+        mConsumer->streamEnd(mCache ? this : nullptr);
+      }
       return NS_OK;
     }
 
@@ -450,10 +520,42 @@ class JSStreamConsumer final : public nsIInputStreamCallback {
 
     return NS_OK;
   }
+
+  // JS::OptimizedEncodingListener
+
+  void storeOptimizedEncoding(JS::UniqueOptimizedEncodingBytes bytes) override {
+    MOZ_ASSERT(mCache, "we only listen if there's a cache entry");
+
+    NS_DispatchToMainThread(new StoreOptimizedEncodingRunnable(
+        std::move(mCache), std::move(bytes)));
+  }
 };
 
 NS_IMPL_ISUPPORTS(JSStreamConsumer, nsIInputStreamCallback)
 
+// static
+const nsCString FetchUtil::WasmAltDataType;
+
+// static
+void FetchUtil::InitWasmAltDataType() {
+  nsCString& type = const_cast<nsCString&>(WasmAltDataType);
+  MOZ_ASSERT(type.IsEmpty());
+
+  RunOnShutdown([]() {
+    // Avoid nsStringBuffer leak tests failures.
+    const_cast<nsCString&>(WasmAltDataType).Truncate();
+  });
+
+  type.Append(nsLiteralCString("wasm-"));
+
+  JS::BuildIdCharVector buildId;
+  if (!JS::GetOptimizedEncodingBuildId(&buildId)) {
+    MOZ_CRASH("build id oom");
+  }
+
+  type.Append(buildId.begin(), buildId.length());
+}
+
 static bool ThrowException(JSContext* aCx, unsigned errorNumber) {
   JS_ReportErrorNumberASCII(aCx, js::GetErrorMessage, nullptr, errorNumber);
   return false;
@@ -464,6 +566,7 @@ bool FetchUtil::StreamResponseToJS(JSContext* aCx, JS::HandleObject aObj,
                                    JS::MimeType aMimeType,
                                    JS::StreamConsumer* aConsumer,
                                    WorkerPrivate* aMaybeWorker) {
+  MOZ_ASSERT(!WasmAltDataType.IsEmpty());
   MOZ_ASSERT(!aMaybeWorker == NS_IsMainThread());
 
   RefPtr<Response> response;
@@ -529,13 +632,47 @@ bool FetchUtil::StreamResponseToJS(JSContext* aCx, JS::HandleObject aObj,
     return ThrowException(aCx, JSMSG_OUT_OF_MEMORY);
   }
 
-  nsCOMPtr<nsIInputStream> body;
-  ir->GetUnfilteredBody(getter_AddRefs(body));
-  if (!body) {
-    aConsumer->streamEnd();
-    return true;
+  nsCOMPtr<nsIInputStream> stream;
+
+  nsMainThreadPtrHandle<nsICacheInfoChannel> cache;
+  bool optimizedEncoding = false;
+  if (ir->HasCacheInfoChannel()) {
+    cache = ir->TakeCacheInfoChannel();
+
+    nsAutoCString altDataType;
+    if (NS_SUCCEEDED(cache->GetAlternativeDataType(altDataType)) &&
+        WasmAltDataType.Equals(altDataType)) {
+      optimizedEncoding = true;
+      rv = cache->GetAlternativeDataInputStream(getter_AddRefs(stream));
+      if (NS_WARN_IF(NS_FAILED(rv))) {
+        return ThrowException(aCx, JSMSG_OUT_OF_MEMORY);
+      }
+      if (ir->HasBeenCloned()) {
+        // If `Response` is cloned, clone alternative data stream instance.
+        // The cache entry does not clone automatically, and multiple
+        // JSStreamConsumer instances will collide during read if not cloned.
+        nsCOMPtr<nsICloneableInputStream> original = do_QueryInterface(stream);
+        if (NS_WARN_IF(!original)) {
+          return ThrowException(aCx, JSMSG_OUT_OF_MEMORY);
+        }
+        rv = original->Clone(getter_AddRefs(stream));
+        if (NS_WARN_IF(NS_FAILED(rv))) {
+          return ThrowException(aCx, JSMSG_OUT_OF_MEMORY);
+        }
+      }
+    }
   }
 
+  if (!optimizedEncoding) {
+    ir->GetUnfilteredBody(getter_AddRefs(stream));
+    if (!stream) {
+      aConsumer->streamEnd();
+      return true;
+    }
+  }
+
+  MOZ_ASSERT(stream);
+
   IgnoredErrorResult error;
   response->SetBodyUsed(aCx, error);
   if (NS_WARN_IF(error.Failed())) {
@@ -544,8 +681,8 @@ bool FetchUtil::StreamResponseToJS(JSContext* aCx, JS::HandleObject aObj,
 
   nsIGlobalObject* global = xpc::NativeGlobal(js::UncheckedUnwrap(aObj));
 
-  if (!JSStreamConsumer::Start(std::move(body), aConsumer, global,
-                               aMaybeWorker)) {
+  if (!JSStreamConsumer::Start(stream, global, aMaybeWorker, aConsumer,
+                               std::move(cache), optimizedEncoding)) {
     return ThrowException(aCx, JSMSG_OUT_OF_MEMORY);
   }
 

dom/fetch/FetchUtil.h

@@ -14,7 +14,6 @@
 #include "mozilla/dom/FormData.h"
 
 #define WASM_CONTENT_TYPE "application/wasm"
-#define WASM_ALT_DATA_TYPE_V1 "wasm/machine-code/1"
 
 class nsIPrincipal;
 class nsIHttpChannel;
@@ -52,6 +51,15 @@ class FetchUtil final {
                                      nsIHttpChannel* aChannel,
                                      InternalRequest& aRequest);
 
+  /**
+   * The WebAssembly alt data type includes build-id, cpu-id and other relevant
+   * state that is necessary to ensure the validity of caching machine code and
+   * metadata in alt data. InitWasmAltDataType() must be called during startup
+   * before the first fetch(), ensuring that !WasmAltDataType.IsEmpty().
+   */
+  static const nsCString WasmAltDataType;
+  static void InitWasmAltDataType();
+
   /**
    * Check that the given object is a Response and, if so, stream to the given
    * JS consumer. On any failure, this function will report an error on the

dom/fetch/InternalResponse.cpp

@@ -52,7 +52,8 @@ InternalResponse::InternalResponse(uint16_t aStatus,
       mBodySize(UNKNOWN_BODY_SIZE),
       mPaddingSize(UNKNOWN_PADDING_SIZE),
       mErrorCode(NS_OK),
-      mCredentialsMode(aCredentialsMode) {}
+      mCredentialsMode(aCredentialsMode),
+      mCloned(false) {}
 
 /* static */ RefPtr<InternalResponse> InternalResponse::FromIPC(
     const IPCInternalResponse& aIPCResponse) {
@@ -164,6 +165,7 @@ void InternalResponse::ToIPC(
 already_AddRefed<InternalResponse> InternalResponse::Clone(
     CloneType aCloneType) {
   RefPtr<InternalResponse> clone = CreateIncompleteCopy();
+  clone->mCloned = (mCloned = true);
 
   clone->mHeaders = new InternalHeaders(*mHeaders);
 

dom/fetch/InternalResponse.h

@@ -298,6 +298,8 @@ class InternalResponse final {
     return !!mCacheInfoChannel;
   }
 
+  bool HasBeenCloned() const { return mCloned; }
+
   void InitChannelInfo(nsIChannel* aChannel) {
     mChannelInfo.InitFromChannel(aChannel);
   }
@@ -361,6 +363,7 @@ class InternalResponse final {
   nsCString mAlternativeDataType;
   nsCOMPtr<nsIInputStream> mAlternativeBody;
   nsMainThreadPtrHandle<nsICacheInfoChannel> mCacheInfoChannel;
+  bool mCloned;
 
  public:
   static const int64_t UNKNOWN_BODY_SIZE = -1;

dom/promise/tests/test_webassembly_compile.html

@@ -10,13 +10,17 @@
 </head>
 <body>
 <script>
-const wasmIsSupported = SpecialPowers.Cu.getJSTestingFunctions().wasmIsSupported;
+const testingFunctions = SpecialPowers.Cu.getJSTestingFunctions();
+const wasmIsSupported = SpecialPowers.unwrap(testingFunctions.wasmIsSupported);
+const wasmHasTier2CompilationCompleted = SpecialPowers.unwrap(testingFunctions.wasmHasTier2CompilationCompleted);
+const wasmLoadedFromCache = SpecialPowers.unwrap(testingFunctions.wasmLoadedFromCache);
 
 // The test_webassembly_compile_sample.wasm is a medium-sized module with 100
 // functions that call each other recursively, returning a computed sum.
 // Any other non-trivial module could be generated and used.
 var sampleCode;
 const sampleURL = "test_webassembly_compile_sample.wasm";
+const sampleURLWithRandomQuery = () => sampleURL + "?id=" + String(Math.ceil(Math.random()*100000));
 const sampleExportName = "run";
 const sampleResult = 1275;
 
@@ -219,6 +223,78 @@ function compileStreamingFetch() {
   .catch(err => { ok(false, String(err)); });
 }
 
+function compileCachedBasic() {
+  const url = sampleURLWithRandomQuery();
+  WebAssembly.compileStreaming(fetch(url))
+  .then(module => {
+    checkSampleModule(module);
+    ok(!wasmLoadedFromCache(module), "not cached yet");
+    while(!wasmHasTier2CompilationCompleted(module));
+    return WebAssembly.compileStreaming(fetch(url));
+  })
+  .then(module => {
+    checkSampleModule(module);
+    ok(wasmLoadedFromCache(module), "loaded from cache");
+  })
+  .then(() => runTest())
+  .catch(err => { ok(false, String(err)) });
+}
+
+const Original = "original";
+const Clone = "clone";
+
+function compileCachedBothClonesHitCache(which) {
+  const url = sampleURLWithRandomQuery();
+  WebAssembly.compileStreaming(fetch(url))
+  .then(module => {
+    checkSampleModule(module);
+    ok(!wasmLoadedFromCache(module), "not cached yet");
+    while(!wasmHasTier2CompilationCompleted(module));
+    return fetch(url);
+  })
+  .then(original => {
+    let clone = original.clone();
+    if (which === Clone) [clone, original] = [original, clone];
+    return Promise.all([
+      WebAssembly.compileStreaming(original),
+      WebAssembly.compileStreaming(clone)
+    ]);
+  })
+  .then(([m1, m2]) => {
+    checkSampleModule(m1);
+    ok(wasmLoadedFromCache(m1), "clone loaded from cache");
+    checkSampleModule(m2);
+    ok(wasmLoadedFromCache(m2), "original loaded from cache");
+  })
+  .then(() => runTest())
+  .catch(err => { ok(false, String(err)) });
+}
+
+function compileCachedCacheThroughClone(which) {
+  const url = sampleURLWithRandomQuery();
+  fetch(url)
+  .then(original => {
+    ok(true, "fun time");
+    let clone = original.clone();
+    if (which === Clone) [clone, original] = [original, clone];
+    return Promise.all([
+      WebAssembly.compileStreaming(original),
+      clone.arrayBuffer()
+    ]);
+  })
+  .then(([module, buffer]) => {
+    ok(!wasmLoadedFromCache(module), "not cached yet");
+    ok(buffer instanceof ArrayBuffer);
+    while(!wasmHasTier2CompilationCompleted(module));
+    return WebAssembly.compileStreaming(fetch(url));
+  })
+  .then(m => {
+    ok(wasmLoadedFromCache(m), "cache hit of " + which);
+  })
+  .then(() => runTest())
+  .catch(err => { ok(false, String(err)) });
+}
+
 function instantiateStreamingFetch() {
   WebAssembly.instantiateStreaming(fetch(sampleURL))
   .then(({module, instance}) => { checkSampleModule(module); checkSampleInstance(instance); runTest(); })
@@ -277,6 +353,11 @@ var tests = [ propertiesExist,
               compileStreamingDoubleUseFail,
               compileStreamingNullBody,
               compileStreamingFetch,
+              compileCachedBasic,
+              compileCachedBothClonesHitCache.bind(Original),
+              compileCachedBothClonesHitCache.bind(Clone),
+              compileCachedCacheThroughClone.bind(Original),
+              compileCachedCacheThroughClone.bind(Clone),
               instantiateStreamingFetch,
               compileManyStreamingFetch,
               runWorkerTests,

dom/tests/mochitest/fetch/mochitest.ini

@@ -44,7 +44,6 @@ support-files =
   !/dom/security/test/csp/file_redirects_resource.sjs
   !/dom/base/test/referrer_helper.js
   !/dom/base/test/referrer_testserver.sjs
-  !/dom/promise/tests/test_webassembly_compile_sample.wasm
 prefs =
   javascript.options.streams=true
 [test_headers.html]

js/xpconnect/src/XPCJSRuntime.cpp

@@ -64,6 +64,7 @@
 #include "mozilla/dom/GeneratedAtomList.h"
 #include "mozilla/dom/BindingUtils.h"
 #include "mozilla/dom/Element.h"
+#include "mozilla/dom/FetchUtil.h"
 #include "mozilla/dom/WindowBinding.h"
 #include "mozilla/Atomics.h"
 #include "mozilla/Attributes.h"
@@ -3022,7 +3023,10 @@ void XPCJSRuntime::Initialize(JSContext* cx) {
   JS::SetXrayJitInfo(&gXrayJitInfo);
   JS::SetProcessLargeAllocationFailureCallback(
       OnLargeAllocationFailureCallback);
+
+  // The WasmAltDataType is build by the JS engine from the build id.
   JS::SetProcessBuildIdOp(GetBuildId);
+  FetchUtil::InitWasmAltDataType();
 
   // The JS engine needs to keep the source code around in order to implement
   // Function.prototype.toSource(). It'd be nice to not have to do this for