Bug 1894231 - Remove CheckedUnsafePtr<WorkerPrivate> of WebTaskSchedulerWorker. r=dom-worker-reviewers,smaug

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

dom/webscheduling/WebTaskScheduler.cpp

@@ -147,7 +147,7 @@ already_AddRefed<WebTaskSchedulerWorker> WebTaskScheduler::CreateForWorker(
     WorkerPrivate* aWorkerPrivate) {
   aWorkerPrivate->AssertIsOnWorkerThread();
   RefPtr<WebTaskSchedulerWorker> scheduler =
-      new WebTaskSchedulerWorker(aWorkerPrivate);
+      WebTaskSchedulerWorker::Create(aWorkerPrivate);
   return scheduler.forget();
 }
 

dom/webscheduling/WebTaskSchedulerWorker.cpp

@@ -17,9 +17,29 @@ WebTaskWorkerRunnable::WebTaskWorkerRunnable(
   MOZ_ASSERT(mSchedulerWorker);
 }
 
+RefPtr<WebTaskSchedulerWorker> WebTaskSchedulerWorker::Create(
+    WorkerPrivate* aWorkerPrivate) {
+  MOZ_ASSERT(aWorkerPrivate);
+  aWorkerPrivate->AssertIsOnWorkerThread();
+
+  RefPtr<WebTaskSchedulerWorker> scheduler =
+      MakeRefPtr<WebTaskSchedulerWorker>(aWorkerPrivate);
+
+  scheduler->mWorkerRef = StrongWorkerRef::Create(
+      aWorkerPrivate, "WebTaskSchedulerWorker", [scheduler]() {
+        // Set mWorkerIsShuttingDown as true here to avoid dispatching tasks
+        // to worker thread.
+        scheduler->mWorkerIsShuttingDown = true;
+      });
+  if (!scheduler->mWorkerRef) {
+    NS_WARNING("Create WebTaskScheduler when Worker is shutting down");
+    scheduler->mWorkerIsShuttingDown = true;
+  }
+  return scheduler;
+}
+
 WebTaskSchedulerWorker::WebTaskSchedulerWorker(WorkerPrivate* aWorkerPrivate)
-    : WebTaskScheduler(aWorkerPrivate->GlobalScope()),
-      mWorkerPrivate(aWorkerPrivate) {}
+    : WebTaskScheduler(aWorkerPrivate->GlobalScope()) {}
 
 bool WebTaskWorkerRunnable::WorkerRun(JSContext* aCx,
                                       WorkerPrivate* aWorkerPrivate) {
@@ -36,9 +56,18 @@ bool WebTaskWorkerRunnable::WorkerRun(JSContext* aCx,
 
 nsresult WebTaskSchedulerWorker::SetTimeoutForDelayedTask(WebTask* aTask,
                                                           uint64_t aDelay) {
-  if (!mWorkerPrivate) {
+  if (mWorkerIsShuttingDown) {
+    return NS_ERROR_ABORT;
+  }
+
+  if (!mWorkerRef) {
     return NS_ERROR_UNEXPECTED;
   }
+
+  WorkerPrivate* workerPrivate = mWorkerRef->Private();
+  MOZ_ASSERT(workerPrivate);
+  workerPrivate->AssertIsOnWorkerThread();
+
   JSContext* cx = nsContentUtils::GetCurrentJSContext();
   if (!cx) {
     return NS_ERROR_UNEXPECTED;
@@ -48,24 +77,31 @@ nsresult WebTaskSchedulerWorker::SetTimeoutForDelayedTask(WebTask* aTask,
   ErrorResult rv;
 
   int32_t delay = aDelay > INT32_MAX ? INT32_MAX : (int32_t)aDelay;
-  mWorkerPrivate->SetTimeout(cx, handler, delay,
+  workerPrivate->SetTimeout(cx, handler, delay,
                             /* aIsInterval */ false,
                             Timeout::Reason::eDelayedWebTaskTimeout, rv);
   return rv.StealNSResult();
 }
 
 bool WebTaskSchedulerWorker::DispatchEventLoopRunnable() {
-  if (!mWorkerPrivate) {
+  if (mWorkerIsShuttingDown) {
     return false;
   }
+
+  if (!mWorkerRef) {
+    return false;
+  }
+  MOZ_ASSERT(mWorkerRef->Private());
+  mWorkerRef->Private()->AssertIsOnWorkerThread();
+
   RefPtr<WebTaskWorkerRunnable> runnable =
-      new WebTaskWorkerRunnable(mWorkerPrivate, this);
-  return runnable->Dispatch(mWorkerPrivate);
+      new WebTaskWorkerRunnable(mWorkerRef->Private(), this);
+  return runnable->Dispatch(mWorkerRef->Private());
 }
 
 void WebTaskSchedulerWorker::Disconnect() {
-  if (mWorkerPrivate) {
-    mWorkerPrivate = nullptr;
+  if (mWorkerRef) {
+    mWorkerRef = nullptr;
   }
   WebTaskScheduler::Disconnect();
 }

dom/webscheduling/WebTaskSchedulerWorker.h

@@ -13,6 +13,7 @@
 #include "mozilla/LinkedList.h"
 #include "mozilla/dom/WorkerRunnable.h"
 #include "mozilla/dom/WorkerPrivate.h"
+#include "mozilla/dom/WorkerRef.h"
 #include "mozilla/dom/WebTaskSchedulingBinding.h"
 
 namespace mozilla::dom {
@@ -32,6 +33,8 @@ class WebTaskWorkerRunnable final : public WorkerSameThreadRunnable {
 
 class WebTaskSchedulerWorker final : public WebTaskScheduler {
  public:
+  static RefPtr<WebTaskSchedulerWorker> Create(WorkerPrivate* aWorkerPrivate);
+
   explicit WebTaskSchedulerWorker(WorkerPrivate* aWorkerPrivate);
 
   void Disconnect() override;
@@ -42,7 +45,8 @@ class WebTaskSchedulerWorker final : public WebTaskScheduler {
   nsresult SetTimeoutForDelayedTask(WebTask* aTask, uint64_t aDelay) override;
   bool DispatchEventLoopRunnable() override;
 
-  CheckedUnsafePtr<WorkerPrivate> mWorkerPrivate;
+  RefPtr<StrongWorkerRef> mWorkerRef;
+  bool mWorkerIsShuttingDown{false};
 };
 }  // namespace mozilla::dom
 #endif