Bug 966384 - Promises on workers use correct busy count. r=khuey

dom/promise/Promise.cpp

@@ -57,11 +57,11 @@ private:
   nsRefPtr<Promise> mPromise;
 };
 
-class WorkerPromiseTask MOZ_FINAL : public WorkerRunnable
+class WorkerPromiseTask MOZ_FINAL : public WorkerSameThreadRunnable
 {
 public:
   WorkerPromiseTask(WorkerPrivate* aWorkerPrivate, Promise* aPromise)
-    : WorkerRunnable(aWorkerPrivate, WorkerThreadUnchangedBusyCount)
+    : WorkerSameThreadRunnable(aWorkerPrivate)
     , mPromise(aPromise)
   {
     MOZ_ASSERT(aPromise);
@@ -161,7 +161,7 @@ public:
   }
 };
 
-class WorkerPromiseResolverTask MOZ_FINAL : public WorkerRunnable,
+class WorkerPromiseResolverTask MOZ_FINAL : public WorkerSameThreadRunnable,
                                             public PromiseResolverMixin
 {
 public:
@@ -169,7 +169,7 @@ public:
                             Promise* aPromise,
                             JS::Handle<JS::Value> aValue,
                             Promise::PromiseState aState)
-    : WorkerRunnable(aWorkerPrivate, WorkerThreadUnchangedBusyCount),
+    : WorkerSameThreadRunnable(aWorkerPrivate),
       PromiseResolverMixin(aPromise, aValue, aState)
   {}
 
@@ -832,7 +832,7 @@ Promise::AppendCallbacks(PromiseCallback* aResolveCallback,
       WorkerPrivate* worker = GetCurrentThreadWorkerPrivate();
       MOZ_ASSERT(worker);
       nsRefPtr<WorkerPromiseTask> task = new WorkerPromiseTask(worker, this);
-      worker->Dispatch(task);
+      task->Dispatch(worker->GetJSContext());
     }
     mTaskPending = true;
   }
@@ -1045,7 +1045,7 @@ Promise::RunResolveTask(JS::Handle<JS::Value> aValue,
       MOZ_ASSERT(worker);
       nsRefPtr<WorkerPromiseResolverTask> task =
         new WorkerPromiseResolverTask(worker, this, aValue, aState);
-      worker->Dispatch(task);
+      task->Dispatch(worker->GetJSContext());
     }
     return;
   }

dom/workers/WorkerRunnable.cpp

@@ -7,10 +7,12 @@
 
 #include "nsIEventTarget.h"
 #include "nsIRunnable.h"
+#include "nsThreadUtils.h"
+
+#include "mozilla/DebugOnly.h"
 
 #include "js/RootingAPI.h"
 #include "js/Value.h"
-#include "nsThreadUtils.h"
 
 #include "WorkerPrivate.h"
 
@@ -478,3 +480,43 @@ MainThreadWorkerControlRunnable::PostDispatch(JSContext* aCx,
 }
 
 NS_IMPL_ISUPPORTS_INHERITED0(WorkerControlRunnable, WorkerRunnable)
+
+bool
+WorkerSameThreadRunnable::PreDispatch(JSContext* aCx,
+                                      WorkerPrivate* aWorkerPrivate)
+{
+  aWorkerPrivate->AssertIsOnWorkerThread();
+  return true;
+}
+
+void
+WorkerSameThreadRunnable::PostDispatch(JSContext* aCx,
+                                       WorkerPrivate* aWorkerPrivate,
+                                       bool aDispatchResult)
+{
+  aWorkerPrivate->AssertIsOnWorkerThread();
+  if (aDispatchResult) {
+    DebugOnly<bool> willIncrement = aWorkerPrivate->ModifyBusyCountFromWorker(aCx, true);
+    // Should never fail since if this thread is still running, so should the
+    // parent and it should be able to process a control runnable.
+    MOZ_ASSERT(willIncrement);
+  }
+}
+
+void
+WorkerSameThreadRunnable::PostRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate,
+                                  bool aRunResult)
+{
+  MOZ_ASSERT(aCx);
+  MOZ_ASSERT(aWorkerPrivate);
+
+  aWorkerPrivate->AssertIsOnWorkerThread();
+
+  DebugOnly<bool> willDecrement = aWorkerPrivate->ModifyBusyCountFromWorker(aCx, false);
+  MOZ_ASSERT(willDecrement);
+
+  if (!aRunResult) {
+    JS_ReportPendingException(aCx);
+  }
+}
+

dom/workers/WorkerRunnable.h

@@ -314,6 +314,34 @@ protected:
                bool aDispatchResult) MOZ_OVERRIDE;
 };
 
+// A WorkerRunnable that should be dispatched from the worker to itself for
+// async tasks. This will increment the busy count PostDispatch() (only if
+// dispatch was successful) and decrement it in PostRun().
+//
+// Async tasks will almost always want to use this since
+// a WorkerSameThreadRunnable keeps the Worker from being GCed.
+class WorkerSameThreadRunnable : public WorkerRunnable
+{
+protected:
+  WorkerSameThreadRunnable(WorkerPrivate* aWorkerPrivate)
+  : WorkerRunnable(aWorkerPrivate, WorkerThreadModifyBusyCount)
+  { }
+
+  virtual ~WorkerSameThreadRunnable()
+  { }
+
+  virtual bool
+  PreDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate) MOZ_OVERRIDE;
+
+  virtual void
+  PostDispatch(JSContext* aCx, WorkerPrivate* aWorkerPrivate,
+               bool aDispatchResult) MOZ_OVERRIDE;
+
+  virtual void
+  PostRun(JSContext* aCx, WorkerPrivate* aWorkerPrivate,
+          bool aRunResult) MOZ_OVERRIDE;
+};
+
 END_WORKERS_NAMESPACE
 
 #endif // mozilla_dom_workers_workerrunnable_h__