Bug 1652126: Obtain an OffThreadToken immediately so parse tasks can be canceled anytime, and clean up dangling Runnables during cancellation. r=smaug

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

dom/script/ScriptLoadRequest.cpp

@@ -58,6 +58,9 @@ NS_IMPL_CYCLE_COLLECTION_CLASS(ScriptLoadRequest)
 NS_IMPL_CYCLE_COLLECTION_UNLINK_BEGIN(ScriptLoadRequest)
   NS_IMPL_CYCLE_COLLECTION_UNLINK(mFetchOptions, mCacheInfo)
   tmp->mScript = nullptr;
+  if (Runnable* runnable = tmp->mRunnable.exchange(nullptr)) {
+    runnable->Release();
+  }
   tmp->DropBytecodeCacheReferences();
   tmp->MaybeUnblockOnload();
 NS_IMPL_CYCLE_COLLECTION_UNLINK_END
@@ -90,6 +93,7 @@ ScriptLoadRequest::ScriptLoadRequest(ScriptKind aKind, nsIURI* aURI,
       mIsTracking(false),
       mFetchOptions(aFetchOptions),
       mOffThreadToken(nullptr),
+      mRunnable(nullptr),
       mScriptTextLength(0),
       mScriptBytecode(),
       mBytecodeOffset(0),
@@ -147,6 +151,14 @@ void ScriptLoadRequest::MaybeCancelOffThreadScript() {
     MOZ_ASSERT(IsBytecode());
     JS::CancelOffThreadScriptDecoder(cx, mOffThreadToken);
   }
+
+  // Cancellation request above should guarantee removal of the parse task, so
+  // releasing the runnable should be safe to do here.
+  if (Runnable* runnable = mRunnable.exchange(nullptr)) {
+    runnable->Release();
+  }
+
+  MaybeUnblockOnload();
   mOffThreadToken = nullptr;
 }
 

dom/script/ScriptLoadRequest.h

@@ -7,6 +7,7 @@
 #ifndef mozilla_dom_ScriptLoadRequest_h
 #define mozilla_dom_ScriptLoadRequest_h
 
+#include "mozilla/Atomics.h"
 #include "mozilla/Assertions.h"
 #include "mozilla/CORSMode.h"
 #include "mozilla/dom/Element.h"
@@ -324,6 +325,10 @@ class ScriptLoadRequest
   JS::OffThreadToken* mOffThreadToken;  // Off-thread parsing token.
   Maybe<nsString> mSourceMapURL;  // Holds source map url for loaded scripts
 
+  Atomic<Runnable*> mRunnable;  // Runnable created when dispatching off thread
+                                // compile. Tracked here so that it can be
+                                // properly released during cancellation.
+
   // Holds the top-level JSScript that corresponds to the current source, once
   // it is parsed, and planned to be saved in the bytecode cache.
   JS::Heap<JSScript*> mScript;

dom/script/ScriptLoader.cpp

@@ -230,9 +230,7 @@ ScriptLoader::~ScriptLoader() {
     mPendingChildLoaders[j]->RemoveParserBlockingScriptExecutionBlocker();
   }
 
-  // Cancel any unused preload requests
   for (size_t i = 0; i < mPreloads.Length(); i++) {
-    mPreloads[i].mRequest->Cancel();
     AccumulateCategorical(LABELS_DOM_SCRIPT_PRELOAD_RESULT::NotUsed);
   }
 }
@@ -2200,6 +2198,9 @@ NotifyOffThreadScriptLoadCompletedRunnable::Run() {
   // function.
   RefPtr<ScriptLoadRequest> request = std::move(mRequest);
 
+  // Runnable pointer should have been cleared in the offthread callback.
+  MOZ_ASSERT(!request->mRunnable);
+
 #ifdef MOZ_GECKO_PROFILER
   if (profiler_is_active()) {
     const char* scriptSourceString;
@@ -2223,16 +2224,19 @@ NotifyOffThreadScriptLoadCompletedRunnable::Run() {
 
   RefPtr<ScriptLoader> loader = std::move(mLoader);
 
-  request->mOffThreadToken = mToken;
-  nsresult rv = loader->ProcessOffThreadRequest(request);
+  // Request was already cancelled at some earlier point.
+  if (!request->mOffThreadToken) {
+    return NS_OK;
+  }
 
-  return rv;
+  return loader->ProcessOffThreadRequest(request);
 }
 
 static void OffThreadScriptLoaderCallback(JS::OffThreadToken* aToken,
                                           void* aCallbackData) {
   RefPtr<NotifyOffThreadScriptLoadCompletedRunnable> aRunnable = dont_AddRef(
       static_cast<NotifyOffThreadScriptLoadCompletedRunnable*>(aCallbackData));
+  MOZ_ASSERT(aRunnable.get() == aRunnable->GetScriptLoadRequest()->mRunnable);
 
 #ifdef MOZ_GECKO_PROFILER
   aRunnable->GetScriptLoadRequest()->mOffThreadParseStopTime =
@@ -2242,6 +2246,12 @@ static void OffThreadScriptLoaderCallback(JS::OffThreadToken* aToken,
   LogRunnable::Run run(aRunnable);
 
   aRunnable->SetToken(aToken);
+
+  // If mRunnable was cleared then request was canceled so do nothing.
+  if (!aRunnable->GetScriptLoadRequest()->mRunnable.exchange(nullptr)) {
+    return;
+  }
+
   NotifyOffThreadScriptLoadCompletedRunnable::Dispatch(aRunnable.forget());
 }
 
@@ -2302,6 +2312,9 @@ nsresult ScriptLoader::AttemptAsyncScriptCompile(ScriptLoadRequest* aRequest,
   aRequest->mOffThreadParseStartTime = TimeStamp::NowUnfuzzed();
 #endif
 
+  // Save the runnable so it can be properly cleared during cancellation.
+  aRequest->mRunnable = runnable.get();
+
   if (aRequest->IsModuleRequest()) {
     MOZ_ASSERT(aRequest->IsTextSource());
     MaybeSourceText maybeSource;
@@ -2311,17 +2324,19 @@ nsresult ScriptLoader::AttemptAsyncScriptCompile(ScriptLoadRequest* aRequest,
     if (maybeSource.constructed<SourceText<char16_t>>()
             ? !JS::CompileOffThreadModule(
                   cx, options, maybeSource.ref<SourceText<char16_t>>(),
-                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable))
+                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable),
+                  &aRequest->mOffThreadToken)
             : !JS::CompileOffThreadModule(
                   cx, options, maybeSource.ref<SourceText<Utf8Unit>>(),
-                  OffThreadScriptLoaderCallback,
-                  static_cast<void*>(runnable))) {
+                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable),
+                  &aRequest->mOffThreadToken)) {
       return NS_ERROR_OUT_OF_MEMORY;
     }
   } else if (aRequest->IsBytecode()) {
     if (!JS::DecodeOffThreadScript(
             cx, options, aRequest->mScriptBytecode, aRequest->mBytecodeOffset,
-            OffThreadScriptLoaderCallback, static_cast<void*>(runnable))) {
+            OffThreadScriptLoaderCallback, static_cast<void*>(runnable),
+            &aRequest->mOffThreadToken)) {
       return NS_ERROR_OUT_OF_MEMORY;
     }
   } else {
@@ -2333,11 +2348,12 @@ nsresult ScriptLoader::AttemptAsyncScriptCompile(ScriptLoadRequest* aRequest,
     if (maybeSource.constructed<SourceText<char16_t>>()
             ? !JS::CompileOffThread(
                   cx, options, maybeSource.ref<SourceText<char16_t>>(),
-                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable))
-            : !JS::CompileOffThread(cx, options,
-                                    maybeSource.ref<SourceText<Utf8Unit>>(),
-                                    OffThreadScriptLoaderCallback,
-                                    static_cast<void*>(runnable))) {
+                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable),
+                  &aRequest->mOffThreadToken)
+            : !JS::CompileOffThread(
+                  cx, options, maybeSource.ref<SourceText<Utf8Unit>>(),
+                  OffThreadScriptLoaderCallback, static_cast<void*>(runnable),
+                  &aRequest->mOffThreadToken)) {
       return NS_ERROR_OUT_OF_MEMORY;
     }
   }
@@ -2358,7 +2374,7 @@ nsresult ScriptLoader::CompileOffThreadOrProcessRequest(
   NS_ASSERTION(nsContentUtils::IsSafeToRunScript(),
                "Processing requests when running scripts is unsafe.");
 
-  if (!aRequest->mOffThreadToken && !aRequest->mWasCompiledOMT) {
+  if (!aRequest->mOffThreadToken && !aRequest->InCompilingStage()) {
     bool couldCompile = false;
     nsresult rv = AttemptAsyncScriptCompile(aRequest, &couldCompile);
     if (NS_FAILED(rv)) {
@@ -3937,6 +3953,11 @@ void ScriptLoader::ParsingComplete(bool aTerminated) {
     mParserBlockingRequest = nullptr;
   }
 
+  // Cancel any unused scripts that were compiled speculatively
+  for (size_t i = 0; i < mPreloads.Length(); i++) {
+    mPreloads[i].mRequest->MaybeCancelOffThreadScript();
+  }
+
   // Have to call this even if aTerminated so we'll correctly unblock
   // onload and all.
   DeferCheckpointReached();