Bug 1606706 - Part 1: Add new TaskController code to the tree. r=smaug,froydnj
authorBas Schouten <bschouten@mozilla.com>
Wed, 03 Jun 2020 23:39:58 +0000
changeset 533823 5f8501eecd1bb5d46b1060d50ffaabe051c55f7e
parent 533822 b57a601bf74caded488ebd5e6222dd6a7bf15df1
child 533824 5b8af5e91633a3e446b9e2894a89310161f9933e
push id37478
push userabutkovits@mozilla.com
push dateThu, 04 Jun 2020 09:29:07 +0000
treeherdermozilla-central@e87e4800d332 [default view] [failures only]
perfherder[talos] [build metrics] [platform microbench] (compared to previous push)
reviewerssmaug, froydnj
bugs1606706
milestone79.0a1
first release with
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
last release without
nightly linux32
nightly linux64
nightly mac
nightly win32
nightly win64
Bug 1606706 - Part 1: Add new TaskController code to the tree. r=smaug,froydnj Differential Revision: https://phabricator.services.mozilla.com/D74671
xpcom/threads/TaskController.cpp
xpcom/threads/TaskController.h
xpcom/threads/moz.build
xpcom/threads/nsThread.cpp
xpcom/threads/nsThread.h
new file mode 100644
--- /dev/null
+++ b/xpcom/threads/TaskController.cpp
@@ -0,0 +1,585 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#include "TaskController.h"
+#include "nsIIdleRunnable.h"
+#include "nsIRunnable.h"
+#include "nsThreadUtils.h"
+#include <algorithm>
+#include <initializer_list>
+#include "mozilla/AbstractEventQueue.h"
+#include "mozilla/StaticMutex.h"
+#include "mozilla/ScopeExit.h"
+#include "mozilla/Unused.h"
+#include "nsIThreadInternal.h"
+#include "nsThread.h"
+
+namespace mozilla {
+
+std::unique_ptr<TaskController> TaskController::sSingleton;
+uint64_t Task::sCurrentTaskSeqNo = 0;
+
+bool TaskManager::
+    UpdateCachesForCurrentIterationAndReportPriorityModifierChanged(
+        const MutexAutoLock& aProofOfLock, IterationType aIterationType) {
+  mCurrentSuspended = IsSuspended(aProofOfLock);
+
+  if (aIterationType == IterationType::EVENT_LOOP_TURN) {
+    int32_t oldModifier = mCurrentPriorityModifier;
+    mCurrentPriorityModifier =
+        GetPriorityModifierForEventLoopTurn(aProofOfLock);
+
+    if (mCurrentPriorityModifier != oldModifier) {
+      return true;
+    }
+  }
+  return false;
+}
+
+Task* Task::GetHighestPriorityDependency() {
+  Task* currentTask = this;
+
+  while (!currentTask->mDependencies.empty()) {
+    auto iter = currentTask->mDependencies.begin();
+
+    while (iter != currentTask->mDependencies.end()) {
+      if ((*iter)->mCompleted) {
+        auto oldIter = iter;
+        iter++;
+        // Completed tasks are removed here to prevent needlessly keeping them
+        // alive or iterating over them in the future.
+        currentTask->mDependencies.erase(oldIter);
+        continue;
+      }
+
+      currentTask = iter->get();
+      break;
+    }
+  }
+
+  return currentTask == this ? nullptr : currentTask;
+}
+
+TaskController* TaskController::Get() {
+  MOZ_ASSERT(sSingleton.get());
+  return sSingleton.get();
+}
+
+bool TaskController::Initialize() {
+  MOZ_ASSERT(!sSingleton);
+  sSingleton = std::make_unique<TaskController>();
+  return sSingleton->InitializeInternal();
+}
+
+bool TaskController::InitializeInternal() {
+  mMTProcessingRunnable = NS_NewRunnableFunction(
+      "TaskController::ExecutePendingMTTasks()",
+      []() { TaskController::Get()->ProcessPendingMTTask(); });
+
+  return true;
+}
+
+void TaskController::SetPerformanceCounterState(
+    PerformanceCounterState* aPerformanceCounterState) {
+  mPerformanceCounterState = aPerformanceCounterState;
+}
+
+/* static */
+void TaskController::Shutdown() {
+  if (sSingleton) {
+    sSingleton->ShutdownInternal();
+  }
+  MOZ_ASSERT(!sSingleton);
+}
+
+void TaskController::ShutdownInternal() { sSingleton = nullptr; }
+
+void TaskController::AddTask(already_AddRefed<Task>&& aTask) {
+  MutexAutoLock lock(mGraphMutex);
+
+  RefPtr<Task> task(aTask);
+
+  if (TaskManager* manager = task->GetManager()) {
+    if (manager->mTaskCount == 0) {
+      mTaskManagers.insert(manager);
+    }
+    manager->DidQueueTask();
+
+    // Set this here since if this manager's priority modifier doesn't change
+    // we will not reprioritize when iterating over the queue.
+    task->mPriorityModifier = manager->mCurrentPriorityModifier;
+  }
+
+#ifdef DEBUG
+  task->mIsInGraph = true;
+
+  for (const RefPtr<Task>& otherTask : task->mDependencies) {
+    MOZ_ASSERT(!otherTask->mTaskManager ||
+               otherTask->mTaskManager == task->mTaskManager);
+  }
+#endif
+
+  auto insertion = mMainThreadTasks.insert(std::move(task));
+  MOZ_ASSERT(insertion.second);
+  (*insertion.first)->mIterator = insertion.first;
+
+  MaybeInterruptTask(*insertion.first);
+}
+
+void TaskController::WaitForTaskOrMessage() {
+  MutexAutoLock lock(mGraphMutex);
+  while (!mMayHaveMainThreadTask) {
+    mMainThreadCV.Wait();
+  }
+}
+
+void TaskController::ExecuteNextTaskOnlyMainThread() {
+  MOZ_ASSERT(NS_IsMainThread());
+  MutexAutoLock lock(mGraphMutex);
+  ExecuteNextTaskOnlyMainThreadInternal(lock);
+}
+
+void TaskController::ProcessPendingMTTask() {
+  MOZ_ASSERT(NS_IsMainThread());
+  MutexAutoLock lock(mGraphMutex);
+  mHasScheduledMTRunnable = false;
+
+  ExecuteNextTaskOnlyMainThreadInternal(lock);
+
+  if (mMayHaveMainThreadTask) {
+    EnsureMainThreadTasksScheduled();
+  }
+}
+
+void TaskController::ReprioritizeTask(Task* aTask, uint32_t aPriority) {
+  MutexAutoLock lock(mGraphMutex);
+  std::set<RefPtr<Task>, Task::PriorityCompare>* queue = &mMainThreadTasks;
+
+  MOZ_ASSERT(aTask->mIterator != queue->end());
+  queue->erase(aTask->mIterator);
+
+  aTask->mPriority = aPriority;
+
+  auto insertion = queue->insert(aTask);
+  MOZ_ASSERT(insertion.second);
+  aTask->mIterator = insertion.first;
+
+  MaybeInterruptTask(aTask);
+}
+
+// Code supporting runnable compatibility.
+// Task that wraps a runnable.
+class RunnableTask : public Task {
+ public:
+  RunnableTask(already_AddRefed<nsIRunnable>&& aRunnable, int32_t aPriority,
+               bool aMainThread = true)
+      : Task(aMainThread, aPriority), mRunnable(aRunnable) {}
+
+  virtual bool Run() override {
+    mRunnable->Run();
+    mRunnable = nullptr;
+    return true;
+  }
+
+  void SetIdleDeadline(TimeStamp aDeadline) override {
+    nsCOMPtr<nsIIdleRunnable> idleRunnable = do_QueryInterface(mRunnable);
+    if (idleRunnable) {
+      idleRunnable->SetDeadline(aDeadline);
+    }
+  }
+
+  PerformanceCounter* GetPerformanceCounter() const override {
+    return nsThread::GetPerformanceCounterBase(mRunnable);
+  }
+
+ private:
+  RefPtr<nsIRunnable> mRunnable;
+};
+
+void TaskController::DispatchRunnable(already_AddRefed<nsIRunnable>&& aRunnable,
+                                      uint32_t aPriority,
+                                      TaskManager* aManager) {
+  RefPtr<RunnableTask> task = new RunnableTask(std::move(aRunnable), aPriority);
+
+  task->SetManager(aManager);
+  TaskController::Get()->AddTask(task.forget());
+}
+
+nsIRunnable* TaskController::GetRunnableForMTTask(bool aReallyWait) {
+  MutexAutoLock lock(mGraphMutex);
+
+  while (mMainThreadTasks.empty()) {
+    if (!aReallyWait) {
+      return nullptr;
+    }
+
+    AUTO_PROFILER_LABEL("TaskController::GetRunnableForMTTask::Wait", IDLE);
+    mMainThreadCV.Wait();
+  }
+
+  return mMTProcessingRunnable;
+}
+
+bool TaskController::HasMainThreadPendingTasks() {
+  auto resetIdleState = MakeScopeExit([&idleManager = mIdleTaskManager] {
+    if (idleManager) {
+      idleManager->State().ClearCachedIdleDeadline();
+    }
+  });
+
+  for (bool considerIdle : {false, true}) {
+    if (considerIdle && !mIdleTaskManager) {
+      continue;
+    }
+
+    MutexAutoLock lock(mGraphMutex);
+
+    if (considerIdle) {
+      mIdleTaskManager->State().ForgetPendingTaskGuarantee();
+      // Temporarily unlock so we can peek our idle deadline.
+      // XXX We could do this _before_ we take the lock if the API would let us.
+      // We do want to do this before looking at mMainThreadTasks, in case
+      // someone adds one while we're unlocked.
+      {
+        MutexAutoUnlock unlock(mGraphMutex);
+        mIdleTaskManager->State().CachePeekedIdleDeadline(unlock);
+      }
+    }
+
+    // Return early if there's no tasks at all.
+    if (mMainThreadTasks.empty()) {
+      return false;
+    }
+
+    // We can cheaply count how many tasks are suspended.
+    uint64_t totalSuspended = 0;
+    for (TaskManager* manager : mTaskManagers) {
+      DebugOnly<bool> modifierChanged =
+          manager
+              ->UpdateCachesForCurrentIterationAndReportPriorityModifierChanged(
+                  lock, TaskManager::IterationType::NOT_EVENT_LOOP_TURN);
+      MOZ_ASSERT(!modifierChanged);
+
+      // The idle manager should be suspended unless we're doing the idle pass.
+      MOZ_ASSERT(manager != mIdleTaskManager || manager->mCurrentSuspended ||
+                     considerIdle,
+                 "Why are idle tasks not suspended here?");
+
+      if (manager->mCurrentSuspended) {
+        // XXX - If managers manage off-main-thread tasks this breaks! This
+        // scenario is explicitly not supported.
+        //
+        // This is only incremented inside the lock -or- decremented on the main
+        // thread so this is safe.
+        totalSuspended += manager->mTaskCount;
+      }
+    }
+
+    // Thi would break down if we have a non-suspended task depending on a
+    // suspended task. This is why for the moment we do not allow tasks
+    // to be dependent on tasks managed by another taskmanager.
+    if (mMainThreadTasks.size() > totalSuspended) {
+      // If mIdleTaskManager->mTaskCount is 0, we never updated the suspended
+      // state of mIdleTaskManager above, hence shouldn't even check it here.
+      // But in that case idle tasks are not contributing to our suspended task
+      // count anyway.
+      if (mIdleTaskManager && mIdleTaskManager->mTaskCount &&
+          !mIdleTaskManager->mCurrentSuspended) {
+        MOZ_ASSERT(considerIdle, "Why is mIdleTaskManager not suspended?");
+        // Check whether the idle tasks were really needed to make our "we have
+        // an unsuspended task" decision.  If they were, we need to force-enable
+        // idle tasks until we run our next task.
+        if (mMainThreadTasks.size() - mIdleTaskManager->mTaskCount <=
+            totalSuspended) {
+          mIdleTaskManager->State().EnforcePendingTaskGuarantee();
+        }
+      }
+      return true;
+    }
+  }
+  return false;
+}
+
+void TaskController::ExecuteNextTaskOnlyMainThreadInternal(
+    const MutexAutoLock& aProofOfLock) {
+  // Block to make it easier to jump to our cleanup.
+  do {
+    bool taskRan = DoExecuteNextTaskOnlyMainThreadInternal(aProofOfLock);
+    if (taskRan) {
+      break;
+    }
+
+    if (!mIdleTaskManager) {
+      break;
+    }
+
+    if (mIdleTaskManager->mTaskCount) {
+      // We have idle tasks that we may not have gotten above because
+      // our idle state is not up to date.  We need to update the idle state
+      // and try again.  We need to temporarily release the lock while we do
+      // that.
+      MutexAutoUnlock unlock(mGraphMutex);
+      mIdleTaskManager->State().UpdateCachedIdleDeadline(unlock);
+    } else {
+      MutexAutoUnlock unlock(mGraphMutex);
+      mIdleTaskManager->State().RanOutOfTasks(unlock);
+    }
+
+    // When we unlocked, someone may have queued a new task on us.  So try to
+    // see whether we can run things again.
+    Unused << DoExecuteNextTaskOnlyMainThreadInternal(aProofOfLock);
+  } while (false);
+
+  if (mIdleTaskManager) {
+    // The pending task guarantee is not needed anymore, since we just tried
+    // running a task
+    mIdleTaskManager->State().ForgetPendingTaskGuarantee();
+
+    if (mMainThreadTasks.empty()) {
+      // XXX the IdlePeriodState API demands we have a MutexAutoUnlock for it.
+      // Otherwise we could perhaps just do this after we exit the locked block,
+      // by pushing the lock down into this method.  Though it's not clear that
+      // we could check mMainThreadTasks.size() once we unlock, and whether we
+      // could maybe substitute mMayHaveMainThreadTask for that check.
+      MutexAutoUnlock unlock(mGraphMutex);
+      mIdleTaskManager->State().RanOutOfTasks(unlock);
+    }
+  }
+}
+
+bool TaskController::DoExecuteNextTaskOnlyMainThreadInternal(
+    const MutexAutoLock& aProofOfLock) {
+  uint32_t totalSuspended = 0;
+  for (TaskManager* manager : mTaskManagers) {
+    bool modifierChanged =
+        manager
+            ->UpdateCachesForCurrentIterationAndReportPriorityModifierChanged(
+                aProofOfLock, TaskManager::IterationType::EVENT_LOOP_TURN);
+    if (modifierChanged) {
+      ProcessUpdatedPriorityModifier(manager);
+    }
+    if (manager->mCurrentSuspended) {
+      totalSuspended += manager->mTaskCount;
+    }
+  }
+
+  MOZ_ASSERT(mMainThreadTasks.size() >= totalSuspended);
+
+  // This would break down if we have a non-suspended task depending on a
+  // suspended task. This is why for the moment we do not allow tasks
+  // to be dependent on tasks managed by another taskmanager.
+  if (mMainThreadTasks.size() > totalSuspended) {
+    for (auto iter = mMainThreadTasks.begin(); iter != mMainThreadTasks.end();
+         iter++) {
+      Task* task = iter->get();
+
+      if (task->mTaskManager && task->mTaskManager->mCurrentSuspended) {
+        // Even though we may want to run some dependencies of this task, we
+        // will run them at their own priority level and not the priority
+        // level of their dependents.
+        continue;
+      }
+
+      task = GetFinalDependency(task);
+
+      if (!task->IsMainThreadOnly() || task->mInProgress ||
+          (task->mTaskManager && task->mTaskManager->mCurrentSuspended)) {
+        continue;
+      }
+
+      mCurrentTasksMT.push(task);
+      mMainThreadTasks.erase(task->mIterator);
+      task->mIterator = mMainThreadTasks.end();
+      task->mInProgress = true;
+      TaskManager* manager = task->GetManager();
+      bool result = false;
+
+      {
+        MutexAutoUnlock unlock(mGraphMutex);
+        if (manager) {
+          manager->WillRunTask();
+          if (manager != mIdleTaskManager) {
+            // Notify the idle period state that we're running a non-idle task.
+            // This needs to happen while our mutex is not locked!
+            mIdleTaskManager->State().FlagNotIdle();
+          } else {
+            TimeStamp idleDeadline =
+                mIdleTaskManager->State().GetCachedIdleDeadline();
+            MOZ_ASSERT(
+                idleDeadline,
+                "How can we not have a deadline if our manager is enabled?");
+            task->SetIdleDeadline(idleDeadline);
+          }
+        }
+        if (mIdleTaskManager) {
+          // We found a task to run; we can clear the idle deadline on our idle
+          // task manager.  This _must_ be done before we actually run the task,
+          // because running the task could reenter via spinning the event loop
+          // and we want to make sure there's no cached idle deadline at that
+          // point.  But we have to make sure we do it after out SetIdleDeadline
+          // call above, in the case when the task is actually an idle task.
+          mIdleTaskManager->State().ClearCachedIdleDeadline();
+        }
+
+        PerformanceCounterState::Snapshot snapshot =
+            mPerformanceCounterState->RunnableWillRun(
+                task->GetPerformanceCounter(), TimeStamp::Now(),
+                manager == mIdleTaskManager);
+
+        result = task->Run();
+        // Task itself should keep manager alive.
+        if (manager) {
+          manager->DidRunTask();
+        }
+
+        mPerformanceCounterState->RunnableDidRun(std::move(snapshot));
+      }
+
+      // Task itself should keep manager alive.
+      if (manager && result && manager->mTaskCount == 0) {
+        mTaskManagers.erase(manager);
+      }
+
+      task->mInProgress = false;
+
+      if (!result) {
+        // Presumably this task was interrupted, leave its dependencies
+        // unresolved and reinsert into the queue.
+        auto insertion =
+            mMainThreadTasks.insert(std::move(mCurrentTasksMT.top()));
+        MOZ_ASSERT(insertion.second);
+        task->mIterator = insertion.first;
+        manager->WillRunTask();
+      } else {
+        task->mCompleted = true;
+#ifdef DEBUG
+        task->mIsInGraph = false;
+#endif
+        // Clear dependencies to release references.
+        task->mDependencies.clear();
+      }
+
+      mCurrentTasksMT.pop();
+      return true;
+    }
+  }
+
+  mMayHaveMainThreadTask = false;
+  if (mIdleTaskManager) {
+    // We did not find a task to run.  We still need to clear the cached idle
+    // deadline on our idle state, because that deadline was only relevant to
+    // the execution of this function.  Had we found a task, we would have
+    // cleared the deadline before running that task.
+    mIdleTaskManager->State().ClearCachedIdleDeadline();
+  }
+  return false;
+}
+
+Task* TaskController::GetFinalDependency(Task* aTask) {
+  Task* nextTask;
+
+  while ((nextTask = aTask->GetHighestPriorityDependency())) {
+    aTask = nextTask;
+  }
+
+  return aTask;
+}
+
+void TaskController::MaybeInterruptTask(Task* aTask) {
+  mGraphMutex.AssertCurrentThreadOwns();
+
+  if (!aTask) {
+    return;
+  }
+
+  // This optimization prevents many slow lookups in long chains of similar
+  // priority.
+  if (!aTask->mDependencies.empty()) {
+    Task* firstDependency = aTask->mDependencies.begin()->get();
+    if (aTask->GetPriority() <= firstDependency->GetPriority() &&
+        !firstDependency->mCompleted &&
+        aTask->IsMainThreadOnly() == firstDependency->IsMainThreadOnly()) {
+      // This task has the same or a higher priority as one of its dependencies,
+      // never any need to interrupt.
+      return;
+    }
+  }
+
+  Task* finalDependency = GetFinalDependency(aTask);
+
+  if (finalDependency->mInProgress) {
+    // No need to wake anything, we can't schedule this task right now anyway.
+    return;
+  }
+
+  EnsureMainThreadTasksScheduled();
+
+  mMayHaveMainThreadTask = true;
+
+  if (mCurrentTasksMT.empty()) {
+    return;
+  }
+
+  // We could go through the steps above here and interrupt an off main
+  // thread task in case it has a lower priority.
+  if (!finalDependency->IsMainThreadOnly()) {
+    return;
+  }
+
+  if (mCurrentTasksMT.top()->GetPriority() < aTask->GetPriority()) {
+    mCurrentTasksMT.top()->RequestInterrupt(aTask->GetPriority());
+  }
+}
+
+Task* TaskController::GetHighestPriorityMTTask() {
+  mGraphMutex.AssertCurrentThreadOwns();
+
+  if (!mMainThreadTasks.empty()) {
+    return mMainThreadTasks.begin()->get();
+  }
+  return nullptr;
+}
+
+void TaskController::EnsureMainThreadTasksScheduled() {
+  if (mObserver) {
+    mObserver->OnDispatchedEvent();
+  }
+  if (mExternalCondVar) {
+    mExternalCondVar->Notify();
+  }
+  mMainThreadCV.Notify();
+}
+
+void TaskController::ProcessUpdatedPriorityModifier(TaskManager* aManager) {
+  mGraphMutex.AssertCurrentThreadOwns();
+
+  MOZ_ASSERT(NS_IsMainThread());
+
+  int32_t modifier = aManager->mCurrentPriorityModifier;
+
+  std::vector<RefPtr<Task>> storedTasks;
+  // Find all relevant tasks.
+  for (auto iter = mMainThreadTasks.begin(); iter != mMainThreadTasks.end();) {
+    if ((*iter)->mTaskManager == aManager) {
+      storedTasks.push_back(*iter);
+      iter = mMainThreadTasks.erase(iter);
+    } else {
+      iter++;
+    }
+  }
+
+  // Reinsert found tasks with their new priorities.
+  for (RefPtr<Task>& ref : storedTasks) {
+    // Kept alive at first by the vector and then by mMainThreadTasks.
+    Task* task = ref;
+    task->mPriorityModifier = modifier;
+    auto insertion = mMainThreadTasks.insert(std::move(ref));
+    MOZ_ASSERT(insertion.second);
+    task->mIterator = insertion.first;
+  }
+}
+
+}  // namespace mozilla
new file mode 100644
--- /dev/null
+++ b/xpcom/threads/TaskController.h
@@ -0,0 +1,360 @@
+/* -*- Mode: C++; tab-width: 8; indent-tabs-mode: nil; c-basic-offset: 2 -*- */
+/* vim: set ts=8 sts=2 et sw=2 tw=80: */
+/* This Source Code Form is subject to the terms of the Mozilla Public
+ * License, v. 2.0. If a copy of the MPL was not distributed with this
+ * file, You can obtain one at http://mozilla.org/MPL/2.0/. */
+
+#ifndef mozilla_TaskController_h
+#define mozilla_TaskController_h
+
+#include "mozilla/CondVar.h"
+#include "mozilla/IdlePeriodState.h"
+#include "mozilla/RefPtr.h"
+#include "mozilla/Mutex.h"
+#include "mozilla/StaticMutex.h"
+#include "mozilla/TimeStamp.h"
+#include "mozilla/EventQueue.h"
+#include "nsISupportsImpl.h"
+#include "nsIEventTarget.h"
+
+#include <thread>
+#include <memory>
+#include <vector>
+#include <set>
+#include <list>
+#include <stack>
+
+class nsIRunnable;
+class nsIThreadObserver;
+
+namespace mozilla {
+
+class Task;
+class TaskController;
+class PerformanceCounter;
+class PerformanceCounterState;
+
+const uint32_t kDefaultPriorityValue = uint32_t(EventQueuePriority::Normal);
+
+// This file contains the core classes to access the Gecko scheduler. The
+// scheduler forms a graph of prioritize tasks, and is responsible for ensuring
+// the execution of tasks or their dependencies in order of inherited priority.
+//
+// The core class is the 'Task' class. The task class describes a single unit of
+// work. Users scheduling work implement this class and are required to
+// reimplement the 'Run' function in order to do work.
+//
+// The TaskManager class is reimplemented by users that require
+// the ability to reprioritize or suspend tasks.
+//
+// The TaskController is responsible for scheduling the work itself. The AddTask
+// function is used to schedule work. The ReprioritizeTask function may be used
+// to change the priority of a task already in the task graph, without
+// unscheduling it.
+
+// The TaskManager is the baseclass used to atomically manage a large set of
+// tasks. API users reimplementing TaskManager may reimplement a number of
+// functions that they may use to indicate to the scheduler changes in the state
+// for any tasks they manage. They may be used to reprioritize or suspend tasks
+// under their control, and will also be notified before and after tasks under
+// their control are executed. Their methods will only be called once per event
+// loop turn, however they may still incur some performance overhead. In
+// addition to this frequent reprioritizations may incur a significant
+// performance overhead and are discouraged. A TaskManager may currently only be
+// used to manage tasks that are bound to the Gecko Main Thread.
+class TaskManager {
+ public:
+  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(TaskManager)
+
+  TaskManager() : mTaskCount(0) {}
+
+  // Subclasses implementing task manager will have this function called to
+  // determine whether their associated tasks are currently suspended. This
+  // will only be called once per iteration of the task queue, this means that
+  // suspension of tasks managed by a single TaskManager may be assumed to
+  // occur atomically.
+  virtual bool IsSuspended(const MutexAutoLock& aProofOfLock) { return false; }
+
+  // Subclasses may implement this in order to supply a priority adjustment
+  // to their managed tasks. This is called once per iteration of the task
+  // queue, and may be assumed to occur atomically for all managed tasks.
+  virtual int32_t GetPriorityModifierForEventLoopTurn(
+      const MutexAutoLock& aProofOfLock) {
+    return 0;
+  }
+
+  void DidQueueTask() { ++mTaskCount; }
+  // This is called when a managed task is about to be executed by the
+  // scheduler. Anyone reimplementing this should ensure to call the parent or
+  // decrement mTaskCount.
+  virtual void WillRunTask() { --mTaskCount; }
+  // This is called when a managed task has finished being executed by the
+  // scheduler.
+  virtual void DidRunTask() {}
+  uint32_t PendingTaskCount() { return mTaskCount; }
+
+ protected:
+  virtual ~TaskManager() {}
+
+ private:
+  friend class TaskController;
+
+  enum class IterationType { NOT_EVENT_LOOP_TURN, EVENT_LOOP_TURN };
+  bool UpdateCachesForCurrentIterationAndReportPriorityModifierChanged(
+      const MutexAutoLock& aProofOfLock, IterationType aIterationType);
+
+  bool mCurrentSuspended = false;
+  int32_t mCurrentPriorityModifier = 0;
+
+  Atomic<uint32_t> mTaskCount;
+};
+
+// A Task is the the base class for any unit of work that may be scheduled.
+// Subclasses may specify their priority and whether they should be bound to
+// the Gecko Main thread. When not bound to the main thread tasks may be
+// executed on any available thread (including the main thread), but they may
+// also be executed in parallel to any other task they do not have a dependency
+// relationship with. Tasks will be run in order of object creation.
+class Task {
+ public:
+  NS_INLINE_DECL_THREADSAFE_REFCOUNTING(Task)
+
+  bool IsMainThreadOnly() { return mMainThreadOnly; }
+
+  // This returns the current task priority with its modifier applied.
+  uint32_t GetPriority() { return mPriority + mPriorityModifier; }
+  uint64_t GetSeqNo() { return mSeqNo; }
+
+  // Callee needs to assume this may be called on any thread.
+  // aInterruptPriority passes the priority of the higher priority task that
+  // is ready to be executed. The task may safely ignore this function, or
+  // interrupt any work being done. It may return 'false' from its run function
+  // in order to be run automatically in the future, or true if it will
+  // reschedule incomplete work manually.
+  virtual void RequestInterrupt(uint32_t aInterruptPriority) {}
+
+  // At the moment this -must- be called before the task is added to the
+  // controller. Calling this after tasks have been added to the controller
+  // results in undefined behavior!
+  // At submission, tasks must depend only on tasks managed by the same, or
+  // no idle manager.
+  void AddDependency(Task* aTask) {
+    MOZ_ASSERT(aTask);
+    MOZ_ASSERT(!mIsInGraph);
+    mDependencies.insert(aTask);
+  }
+
+  // This sets the TaskManager for the current task. Calling this after the
+  // task has been added to the TaskController results in undefined behavior.
+  void SetManager(TaskManager* aManager) {
+    MOZ_ASSERT(mMainThreadOnly);
+    MOZ_ASSERT(!mIsInGraph);
+    mTaskManager = aManager;
+  }
+  TaskManager* GetManager() { return mTaskManager; }
+
+  struct PriorityCompare {
+    bool operator()(const RefPtr<Task>& aTaskA,
+                    const RefPtr<Task>& aTaskB) const {
+      uint32_t prioA = aTaskA->GetPriority();
+      uint32_t prioB = aTaskB->GetPriority();
+      return (prioA > prioB) ||
+             (prioA == prioB && (aTaskA->GetSeqNo() < aTaskB->GetSeqNo()));
+    }
+  };
+
+  // Tell the task about its idle deadline.  Will only be called for
+  // tasks managed by an IdleTaskManager, right before the task runs.
+  virtual void SetIdleDeadline(TimeStamp aDeadline) {}
+
+  virtual PerformanceCounter* GetPerformanceCounter() const { return nullptr; }
+
+ protected:
+  Task(bool aMainThreadOnly, uint32_t aPriority = kDefaultPriorityValue)
+      : mMainThreadOnly(aMainThreadOnly),
+        mSeqNo(sCurrentTaskSeqNo++),
+        mPriority(aPriority) {}
+
+  virtual ~Task() {}
+
+  friend class TaskController;
+
+  // When this returns false, the task is considered incomplete and will be
+  // rescheduled at the current 'mPriority' level.
+  virtual bool Run() = 0;
+
+ private:
+  Task* GetHighestPriorityDependency();
+
+  // Iterator pointing to this task's position in
+  // mThreadableTasks/mMainThreadTasks if, and only if this task is currently
+  // scheduled to be executed. This allows fast access to the task's position
+  // in the set, allowing for fast removal.
+  // This is safe, and remains valid unless the task is removed from the set.
+  // See also iterator invalidation in:
+  // https://en.cppreference.com/w/cpp/container
+  //
+  // Or the spec:
+  // "All Associative Containers: The insert and emplace members shall not
+  // affect the validity of iterators and references to the container
+  // [26.2.6/9]" "All Associative Containers: The erase members shall invalidate
+  // only iterators and references to the erased elements [26.2.6/9]"
+  std::set<RefPtr<Task>, PriorityCompare>::iterator mIterator;
+  std::set<RefPtr<Task>, PriorityCompare> mDependencies;
+
+  RefPtr<TaskManager> mTaskManager;
+
+  // Access to these variables is protected by the GraphMutex.
+  bool mMainThreadOnly;
+  bool mCompleted = false;
+  bool mInProgress = false;
+#ifdef DEBUG
+  bool mIsInGraph = false;
+#endif
+
+  static uint64_t sCurrentTaskSeqNo;
+  int64_t mSeqNo;
+  uint32_t mPriority;
+  // Modifier currently being applied to this task by its taskmanager.
+  int32_t mPriorityModifier = 0;
+};
+
+// A task manager implementation for priority levels that should only
+// run during idle periods.
+class IdleTaskManager : public TaskManager {
+ public:
+  explicit IdleTaskManager(already_AddRefed<nsIIdlePeriod>&& aIdlePeriod)
+      : mIdlePeriodState(std::move(aIdlePeriod)) {}
+
+  IdlePeriodState& State() { return mIdlePeriodState; }
+
+  bool IsSuspended(const MutexAutoLock& aProofOfLock) override {
+    TimeStamp idleDeadline = State().GetCachedIdleDeadline();
+    return !idleDeadline;
+  }
+
+ private:
+  // Tracking of our idle state of various sorts.
+  IdlePeriodState mIdlePeriodState;
+};
+
+// The TaskController is the core class of the scheduler. It is used to
+// schedule tasks to be executed, as well as to reprioritize tasks that have
+// already been scheduled. The core functions to do this are AddTask and
+// ReprioritizeTask.
+class TaskController {
+ public:
+  TaskController()
+      : mGraphMutex("TaskController::mGraphMutex"),
+        mMainThreadCV(mGraphMutex, "TaskController::mMainThreadCV") {}
+
+  static TaskController* Get();
+
+  static bool Initialize();
+
+  void SetThreadObserver(nsIThreadObserver* aObserver) {
+    mObserver = aObserver;
+  }
+  void SetConditionVariable(CondVar* aExternalCondVar) {
+    mExternalCondVar = aExternalCondVar;
+  }
+
+  void SetIdleTaskManager(IdleTaskManager* aIdleTaskManager) {
+    mIdleTaskManager = aIdleTaskManager;
+  }
+
+  // Initialization and shutdown code.
+  bool InitializeInternal();
+  void SetPerformanceCounterState(
+      PerformanceCounterState* aPerformanceCounterState);
+
+  static void Shutdown();
+
+  // This adds a task to the TaskController graph.
+  // This may be called on any thread.
+  void AddTask(already_AddRefed<Task>&& aTask);
+
+  // This wait function is the theoretical function you would need if our main
+  // thread needs to also process OS messages or something along those lines.
+  void WaitForTaskOrMessage();
+
+  // This gets the next (highest priority) task that is only allowed to execute
+  // on the main thread.
+  void ExecuteNextTaskOnlyMainThread();
+
+  // Process all pending main thread tasks.
+  void ProcessPendingMTTask();
+
+  // This allows reprioritization of a task already in the task graph.
+  // This may be called on any thread.
+  void ReprioritizeTask(Task* aTask, uint32_t aPriority);
+
+  void DispatchRunnable(already_AddRefed<nsIRunnable>&& aRunnable,
+                        uint32_t aPriority, TaskManager* aManager = nullptr);
+
+  nsIRunnable* GetRunnableForMTTask(bool aReallyWait);
+
+  bool HasMainThreadPendingTasks();
+
+ private:
+  // This gets the next (highest priority) task that is only allowed to execute
+  // on the main thread, if any, and executes it.
+  void ExecuteNextTaskOnlyMainThreadInternal(const MutexAutoLock& aProofOfLock);
+
+  // The guts of ExecuteNextTaskOnlyMainThreadInternal, which get idle handling
+  // wrapped around them.  Returns whether a task actually ran.
+  bool DoExecuteNextTaskOnlyMainThreadInternal(
+      const MutexAutoLock& aProofOfLock);
+
+  Task* GetFinalDependency(Task* aTask);
+  void MaybeInterruptTask(Task* aTask);
+  Task* GetHighestPriorityMTTask();
+
+  void EnsureMainThreadTasksScheduled();
+
+  void ProcessUpdatedPriorityModifier(TaskManager* aManager);
+
+  void ShutdownInternal();
+
+  static std::unique_ptr<TaskController> sSingleton;
+  static StaticMutex sSingletonMutex;
+
+  // This protects access to the task graph.
+  Mutex mGraphMutex;
+
+  CondVar mMainThreadCV;
+
+  // Variables below are protected by mGraphMutex.
+
+  std::stack<RefPtr<Task>> mCurrentTasksMT;
+
+  // A list of all tasks ordered by priority.
+  std::set<RefPtr<Task>, Task::PriorityCompare> mMainThreadTasks;
+
+  // TaskManagers currently active.
+  // We can use a raw pointer since tasks always hold on to their TaskManager.
+  std::set<TaskManager*> mTaskManagers;
+
+  // This ensures we keep running the main thread if we processed a task there.
+  bool mMayHaveMainThreadTask = true;
+
+  // Whether we have scheduled a runnable on the main thread event loop.
+  // This is used for nsIRunnable compatibility.
+  bool mHasScheduledMTRunnable = false;
+  RefPtr<nsIRunnable> mMTProcessingRunnable;
+
+  // XXX - Thread observer to notify when a new event has been dispatched
+  nsIThreadObserver* mObserver = nullptr;
+  // XXX - External condvar to notify when we have received an event
+  CondVar* mExternalCondVar = nullptr;
+  // Idle task manager so we can properly do idle state stuff.
+  RefPtr<IdleTaskManager> mIdleTaskManager;
+
+  // Our tracking of our performance counter and long task state,
+  // shared with nsThread.
+  PerformanceCounterState* mPerformanceCounterState = nullptr;
+};
+
+}  // namespace mozilla
+
+#endif  // mozilla_TaskController_h
--- a/xpcom/threads/moz.build
+++ b/xpcom/threads/moz.build
@@ -64,16 +64,17 @@ EXPORTS.mozilla += [
     'RWLock.h',
     'SchedulerGroup.h',
     'SharedThreadPool.h',
     'StateMirroring.h',
     'StateWatching.h',
     'SynchronizedEventQueue.h',
     'SyncRunnable.h',
     'TaskCategory.h',
+    'TaskController.h',
     'TaskDispatcher.h',
     'TaskQueue.h',
     'ThreadBound.h',
     'ThreadEventQueue.h',
     'ThrottledEventQueue.h',
 ]
 
 SOURCES += [
@@ -101,16 +102,17 @@ UNIFIED_SOURCES += [
     'nsTimerImpl.cpp',
     'PerformanceCounter.cpp',
     'PrioritizedEventQueue.cpp',
     'RecursiveMutex.cpp',
     'RWLock.cpp',
     'SchedulerGroup.cpp',
     'SharedThreadPool.cpp',
     'SynchronizedEventQueue.cpp',
+    'TaskController.cpp',
     'TaskQueue.cpp',
     'ThreadEventQueue.cpp',
     'ThreadEventTarget.cpp',
     'ThreadLocalVariables.cpp',
     'ThrottledEventQueue.cpp',
     'TimerThread.cpp',
 ]
 
--- a/xpcom/threads/nsThread.cpp
+++ b/xpcom/threads/nsThread.cpp
@@ -1003,16 +1003,22 @@ static bool GetLabeledRunnableName(nsIRu
   }
 
   return labeled;
 }
 #endif
 
 mozilla::PerformanceCounter* nsThread::GetPerformanceCounter(
     nsIRunnable* aEvent) const {
+  return GetPerformanceCounterBase(aEvent);
+}
+
+// static
+mozilla::PerformanceCounter* nsThread::GetPerformanceCounterBase(
+    nsIRunnable* aEvent) {
   RefPtr<SchedulerGroup::Runnable> docRunnable = do_QueryObject(aEvent);
   if (docRunnable) {
     mozilla::dom::DocGroup* docGroup = docRunnable->DocGroup();
     if (docGroup) {
       return docGroup->GetPerformanceCounter();
     }
   }
   return nullptr;
--- a/xpcom/threads/nsThread.h
+++ b/xpcom/threads/nsThread.h
@@ -225,16 +225,19 @@ class nsThread : public nsIThreadInterna
 
   mozilla::SynchronizedEventQueue* EventQueue() { return mEvents.get(); }
 
   bool ShuttingDown() const { return mShutdownContext != nullptr; }
 
   virtual mozilla::PerformanceCounter* GetPerformanceCounter(
       nsIRunnable* aEvent) const;
 
+  static mozilla::PerformanceCounter* GetPerformanceCounterBase(
+      nsIRunnable* aEvent);
+
   size_t SizeOfIncludingThis(mozilla::MallocSizeOf aMallocSizeOf) const;
 
   // Returns the size of this object, its PRThread, and its shutdown contexts,
   // but excluding its event queues.
   size_t ShallowSizeOfIncludingThis(mozilla::MallocSizeOf aMallocSizeOf) const;
 
   size_t SizeOfEventQueues(mozilla::MallocSizeOf aMallocSizeOf) const;