Bug 1405821 - Move microtask handling to CycleCollectedJSContext

Tag UXP Issue #1344

1 files changed, 51 insertions, 0 deletions

diff --git a/xpcom/base/CycleCollectedJSContext.h b/xpcom/base/CycleCollectedJSContext.h
index ac4cf4361..2197eae92 100644
--- a/xpcom/base/CycleCollectedJSContext.h
+++ b/xpcom/base/CycleCollectedJSContext.h
@@ -402,6 +402,37 @@ public:
   // Queue an async microtask to the current main or worker thread.
   virtual void DispatchToMicroTask(already_AddRefed<nsIRunnable> aRunnable);
 
+  // Call EnterMicroTask when you're entering JS execution.
+  // Usually the best way to do this is to use nsAutoMicroTask.
+  void EnterMicroTask()
+  {
+    ++mMicroTaskLevel;
+  }
+
+  void LeaveMicroTask()
+  {
+    if (--mMicroTaskLevel == 0) {
+      PerformMainThreadMicroTaskCheckpoint();
+    }
+  }
+
+  bool IsInMicroTask()
+  {
+    return mMicroTaskLevel != 0;
+  }
+
+  uint32_t MicroTaskLevel()
+  {
+    return mMicroTaskLevel;
+  }
+
+  void SetMicroTaskLevel(uint32_t aLevel)
+  {
+    mMicroTaskLevel = aLevel;
+  }
+
+  void PerformMainThreadMicroTaskCheckpoint();
+
   // Storage for watching rejected promises waiting for some client to
   // consume their rejection.
 
@@ -452,6 +483,7 @@ private:
 
   bool mDisableMicroTaskCheckpoint;
 
+  uint32_t mMicroTaskLevel;
   OOMState mOutOfMemoryState;
   OOMState mLargeAllocationFailureState;
 
@@ -470,6 +502,25 @@ private:
   EnvironmentPreparer mEnvironmentPreparer;
 };
 
+class MOZ_STACK_CLASS nsAutoMicroTask
+{
+public:
+  nsAutoMicroTask()
+  {
+    CycleCollectedJSContext* ccjs = CycleCollectedJSContext::Get();
+    if (ccjs) {
+      ccjs->EnterMicroTask();
+    }
+  }
+  ~nsAutoMicroTask()
+  {
+    CycleCollectedJSContext* ccjs = CycleCollectedJSContext::Get();
+    if (ccjs) {
+      ccjs->LeaveMicroTask();
+    }
+  }
+};
+
 void TraceScriptHolder(nsISupports* aHolder, JSTracer* aTracer);
 
 // Returns true if the JS::TraceKind is one the cycle collector cares about.