TaskGraphRunner Group support

Adds a group to TaskGraphRunner, this can be used by an
implementation along with priority in order to decide how to execute
tasks.

This is a prerequisite patch for thread affinity changes which are
coming in a follow-up.

This patch causes a moderate (~15%) regression in cc_perftests for TaskGraphRunner.

The main regressions I'm seeing are:

In ScheduleTasks, we're seeing an additional 3% time spent in operator[] for the two new maps vs ToT.

We're also seeing an additional 5.5% of time spent in test code which builds the task graph.

Other regressions are less clear, but probably have to do with the additional cost of ScheduleTasks in the function body itself (inlined looping).

BUG=
CQ_INCLUDE_TRYBOTS=tryserver.blink:linux_blink_rel

Review URL: https://codereview.chromium.org/1489233003

Cr-Commit-Position: refs/heads/master@{#365140}

10 files changed, 234 insertions, 100 deletions

diff --git a/cc/raster/single_thread_task_graph_runner.cc b/cc/raster/single_thread_task_graph_runner.cc
index 24994c9..80c85db 100644
--- a/cc/raster/single_thread_task_graph_runner.cc
+++ b/cc/raster/single_thread_task_graph_runner.cc
@@ -112,7 +112,7 @@ void SingleThreadTaskGraphRunner::Run() {
   base::AutoLock lock(lock_);
 
   while (true) {
-    if (!work_queue_.HasReadyToRunTasks()) {
+    if (!RunTaskWithLockAcquired()) {
       // Exit when shutdown is set and no more tasks are pending.
       if (shutdown_)
         break;
@@ -121,18 +121,29 @@ void SingleThreadTaskGraphRunner::Run() {
       has_ready_to_run_tasks_cv_.Wait();
       continue;
     }
-
-    RunTaskWithLockAcquired();
   }
 }
 
-void SingleThreadTaskGraphRunner::RunTaskWithLockAcquired() {
+bool SingleThreadTaskGraphRunner::RunTaskWithLockAcquired() {
   TRACE_EVENT0("toplevel",
                "SingleThreadTaskGraphRunner::RunTaskWithLockAcquired");
 
   lock_.AssertAcquired();
 
-  auto prioritized_task = work_queue_.GetNextTaskToRun();
+  // Find the first category with any tasks to run. This task graph runner
+  // treats categories as an additional priority.
+  const auto& ready_to_run_namespaces = work_queue_.ready_to_run_namespaces();
+  auto found = std::find_if(
+      ready_to_run_namespaces.cbegin(), ready_to_run_namespaces.cend(),
+      [](const std::pair<uint16_t, TaskGraphWorkQueue::TaskNamespace::Vector>&
+             pair) { return !pair.second.empty(); });
+
+  if (found == ready_to_run_namespaces.cend()) {
+    return false;
+  }
+
+  const uint16_t category = found->first;
+  auto prioritized_task = work_queue_.GetNextTaskToRun(category);
   Task* task = prioritized_task.task;
 
   // Call WillRun() before releasing |lock_| and running task.
@@ -152,6 +163,8 @@ void SingleThreadTaskGraphRunner::RunTaskWithLockAcquired() {
   if (work_queue_.HasFinishedRunningTasksInNamespace(
           prioritized_task.task_namespace))
     has_namespaces_with_finished_running_tasks_cv_.Signal();
+
+  return true;
 }
 
 }  // namespace cc
diff --git a/cc/raster/single_thread_task_graph_runner.h b/cc/raster/single_thread_task_graph_runner.h
index 120cd7b..3e8dec0 100644
--- a/cc/raster/single_thread_task_graph_runner.h
+++ b/cc/raster/single_thread_task_graph_runner.h
@@ -40,7 +40,8 @@ class CC_EXPORT SingleThreadTaskGraphRunner
   void Shutdown();
 
  private:
-  void RunTaskWithLockAcquired();
+  // Returns true if there was a task to run.
+  bool RunTaskWithLockAcquired();
 
   scoped_ptr<base::SimpleThread> thread_;
 
diff --git a/cc/raster/synchronous_task_graph_runner.cc b/cc/raster/synchronous_task_graph_runner.cc
index ed05cfb..ea7b15b 100644
--- a/cc/raster/synchronous_task_graph_runner.cc
+++ b/cc/raster/synchronous_task_graph_runner.cc
@@ -4,6 +4,9 @@
 
 #include "cc/raster/synchronous_task_graph_runner.h"
 
+#include <algorithm>
+#include <utility>
+
 #include "base/threading/simple_thread.h"
 #include "base/threading/thread_restrictions.h"
 #include "base/trace_event/trace_event.h"
@@ -43,7 +46,7 @@ void SynchronousTaskGraphRunner::WaitForTasksToFinishRunning(
 
   while (
       !TaskGraphWorkQueue::HasFinishedRunningTasksInNamespace(task_namespace)) {
-    RunTask();
+    DCHECK(RunTask());
   }
 }
 
@@ -57,14 +60,27 @@ void SynchronousTaskGraphRunner::CollectCompletedTasks(
 }
 
 void SynchronousTaskGraphRunner::RunUntilIdle() {
-  while (work_queue_.HasReadyToRunTasks())
-    RunTask();
+  while (RunTask()) {
+  }
 }
 
-void SynchronousTaskGraphRunner::RunTask() {
+bool SynchronousTaskGraphRunner::RunTask() {
   TRACE_EVENT0("toplevel", "SynchronousTaskGraphRunner::RunTask");
 
-  auto prioritized_task = work_queue_.GetNextTaskToRun();
+  // Find the first category with any tasks to run. This task graph runner
+  // treats categories as an additional priority.
+  const auto& ready_to_run_namespaces = work_queue_.ready_to_run_namespaces();
+  auto found = std::find_if(
+      ready_to_run_namespaces.cbegin(), ready_to_run_namespaces.cend(),
+      [](const std::pair<uint16_t, TaskGraphWorkQueue::TaskNamespace::Vector>&
+             pair) { return !pair.second.empty(); });
+
+  if (found == ready_to_run_namespaces.cend()) {
+    return false;
+  }
+
+  const uint16_t category = found->first;
+  auto prioritized_task = work_queue_.GetNextTaskToRun(category);
 
   Task* task = prioritized_task.task;
   task->WillRun();
@@ -72,6 +88,8 @@ void SynchronousTaskGraphRunner::RunTask() {
   task->DidRun();
 
   work_queue_.CompleteTask(prioritized_task);
+
+  return true;
 }
 
 }  // namespace cc
diff --git a/cc/raster/synchronous_task_graph_runner.h b/cc/raster/synchronous_task_graph_runner.h
index 89ecca2..65a7716 100644
--- a/cc/raster/synchronous_task_graph_runner.h
+++ b/cc/raster/synchronous_task_graph_runner.h
@@ -28,7 +28,8 @@ class CC_EXPORT SynchronousTaskGraphRunner : public TaskGraphRunner {
   void RunUntilIdle();
 
  private:
-  void RunTask();
+  // Returns true if there was a task to run.
+  bool RunTask();
 
   // Stores the actual tasks to be run, sorted by priority.
   TaskGraphWorkQueue work_queue_;
diff --git a/cc/raster/task_graph_runner.h b/cc/raster/task_graph_runner.h
index b5ffade..181a455 100644
--- a/cc/raster/task_graph_runner.h
+++ b/cc/raster/task_graph_runner.h
@@ -46,19 +46,29 @@ class CC_EXPORT Task : public base::RefCountedThreadSafe<Task> {
 
 // A task dependency graph describes the order in which to execute a set
 // of tasks. Dependencies are represented as edges. Each node is assigned
-// a priority and a run count that matches the number of dependencies.
-// Priority range from 0 (most favorable scheduling) to UINT_MAX
-// (least favorable).
+// a category, a priority and a run count that matches the number of
+// dependencies. Priority range from 0 (most favorable scheduling) to UINT16_MAX
+// (least favorable). Categories range from 0 to UINT16_MAX. It is up to the
+// implementation and its consumer to determine the meaning (if any) of a
+// category. A TaskGraphRunner implementation may chose to prioritize certain
+// categories over others, regardless of the individual priorities of tasks.
 struct CC_EXPORT TaskGraph {
   struct Node {
     typedef std::vector<Node> Vector;
 
-    Node(Task* task, size_t priority, size_t dependencies)
-        : task(task), priority(priority), dependencies(dependencies) {}
+    Node(Task* task,
+         uint16_t category,
+         uint16_t priority,
+         uint32_t dependencies)
+        : task(task),
+          category(category),
+          priority(priority),
+          dependencies(dependencies) {}
 
     Task* task;
-    size_t priority;
-    size_t dependencies;
+    uint16_t category;
+    uint16_t priority;
+    uint32_t dependencies;
   };
 
   struct Edge {
diff --git a/cc/raster/task_graph_runner_perftest.cc b/cc/raster/task_graph_runner_perftest.cc
index f571d33..0df0757 100644
--- a/cc/raster/task_graph_runner_perftest.cc
+++ b/cc/raster/task_graph_runner_perftest.cc
@@ -230,13 +230,13 @@ class TaskGraphRunnerPerfTest : public testing::Test {
     for (PerfTaskImpl::Vector::const_iterator it = leaf_tasks.begin();
          it != leaf_tasks.end();
          ++it) {
-      graph->nodes.push_back(TaskGraph::Node(it->get(), 0u, 0u));
+      graph->nodes.push_back(TaskGraph::Node(it->get(), 0u, 0u, 0u));
     }
 
     for (PerfTaskImpl::Vector::const_iterator it = tasks.begin();
-         it != tasks.end();
-         ++it) {
-      graph->nodes.push_back(TaskGraph::Node(it->get(), 0u, leaf_tasks.size()));
+         it != tasks.end(); ++it) {
+      graph->nodes.push_back(
+          TaskGraph::Node(it->get(), 0u, 0u, leaf_tasks.size()));
 
       for (PerfTaskImpl::Vector::const_iterator leaf_it = leaf_tasks.begin();
            leaf_it != leaf_tasks.end();
@@ -255,7 +255,7 @@ class TaskGraphRunnerPerfTest : public testing::Test {
     for (PerfTaskImpl::Vector::const_iterator it = top_level_tasks.begin();
          it != top_level_tasks.end();
          ++it) {
-      graph->nodes.push_back(TaskGraph::Node(it->get(), 0u, tasks.size()));
+      graph->nodes.push_back(TaskGraph::Node(it->get(), 0u, 0u, tasks.size()));
     }
   }
 
diff --git a/cc/raster/task_graph_work_queue.cc b/cc/raster/task_graph_work_queue.cc
index 655a005..95faedc 100644
--- a/cc/raster/task_graph_work_queue.cc
+++ b/cc/raster/task_graph_work_queue.cc
@@ -5,11 +5,41 @@
 #include "cc/raster/task_graph_work_queue.h"
 
 #include <algorithm>
+#include <map>
 #include <utility>
 
 #include "base/trace_event/trace_event.h"
 
 namespace cc {
+namespace {
+
+bool CompareTaskPriority(const TaskGraphWorkQueue::PrioritizedTask& a,
+                         const TaskGraphWorkQueue::PrioritizedTask& b) {
+  // In this system, numerically lower priority is run first.
+  return a.priority > b.priority;
+}
+
+class CompareTaskNamespacePriority {
+ public:
+  explicit CompareTaskNamespacePriority(uint16_t category)
+      : category_(category) {}
+
+  bool operator()(const TaskGraphWorkQueue::TaskNamespace* a,
+                  const TaskGraphWorkQueue::TaskNamespace* b) {
+    DCHECK(!a->ready_to_run_tasks.at(category_).empty());
+    DCHECK(!b->ready_to_run_tasks.at(category_).empty());
+
+    // Compare based on task priority of the ready_to_run_tasks heap .front()
+    // will hold the max element of the heap, except after pop_heap, when max
+    // element is moved to .back().
+    return CompareTaskPriority(a->ready_to_run_tasks.at(category_).front(),
+                               b->ready_to_run_tasks.at(category_).front());
+  }
+
+ private:
+  uint16_t category_;
+};
+}  // namespace
 
 TaskGraphWorkQueue::TaskNamespace::TaskNamespace() {}
 
@@ -37,7 +67,9 @@ void TaskGraphWorkQueue::ScheduleTasks(NamespaceToken token, TaskGraph* graph) {
   }
 
   // Build new "ready to run" queue and remove nodes from old graph.
-  task_namespace.ready_to_run_tasks.clear();
+  for (auto& ready_to_run_tasks_it : task_namespace.ready_to_run_tasks) {
+    ready_to_run_tasks_it.second.clear();
+  }
   for (const TaskGraph::Node& node : graph->nodes) {
     // Remove any old nodes that are associated with this task. The result is
     // that the old graph is left with all nodes not present in this graph,
@@ -66,14 +98,17 @@ void TaskGraphWorkQueue::ScheduleTasks(NamespaceToken token, TaskGraph* graph) {
                   node.task) != task_namespace.running_tasks.end())
       continue;
 
-    task_namespace.ready_to_run_tasks.push_back(
-        PrioritizedTask(node.task, &task_namespace, node.priority));
+    task_namespace.ready_to_run_tasks[node.category].push_back(PrioritizedTask(
+        node.task, &task_namespace, node.category, node.priority));
   }
 
-  // Rearrange the elements in |ready_to_run_tasks| in such a way that they
-  // form a heap.
-  std::make_heap(task_namespace.ready_to_run_tasks.begin(),
-                 task_namespace.ready_to_run_tasks.end(), CompareTaskPriority);
+  // Rearrange the elements in each vector within |ready_to_run_tasks| in such a
+  // way that they form a heap.
+  for (auto& it : task_namespace.ready_to_run_tasks) {
+    auto& ready_to_run_tasks = it.second;
+    std::make_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
+                   CompareTaskPriority);
+  }
 
   // Swap task graph.
   task_namespace.graph.Swap(graph);
@@ -100,41 +135,59 @@ void TaskGraphWorkQueue::ScheduleTasks(NamespaceToken token, TaskGraph* graph) {
   }
 
   // Build new "ready to run" task namespaces queue.
-  ready_to_run_namespaces_.clear();
-  for (auto& it : namespaces_) {
-    if (!it.second.ready_to_run_tasks.empty())
-      ready_to_run_namespaces_.push_back(&it.second);
+  for (auto& ready_to_run_namespaces_it : ready_to_run_namespaces_) {
+    ready_to_run_namespaces_it.second.clear();
+  }
+  for (auto& namespace_it : namespaces_) {
+    auto& task_namespace = namespace_it.second;
+    for (auto& ready_to_run_tasks_it : task_namespace.ready_to_run_tasks) {
+      auto& ready_to_run_tasks = ready_to_run_tasks_it.second;
+      uint16_t category = ready_to_run_tasks_it.first;
+      if (!ready_to_run_tasks.empty()) {
+        ready_to_run_namespaces_[category].push_back(&task_namespace);
+      }
+    }
   }
 
   // Rearrange the task namespaces in |ready_to_run_namespaces| in such a
   // way that they form a heap.
-  std::make_heap(ready_to_run_namespaces_.begin(),
-                 ready_to_run_namespaces_.end(), CompareTaskNamespacePriority);
+  for (auto& it : ready_to_run_namespaces_) {
+    uint16_t category = it.first;
+    auto& task_namespace = it.second;
+    std::make_heap(task_namespace.begin(), task_namespace.end(),
+                   CompareTaskNamespacePriority(category));
+  }
 }
 
-TaskGraphWorkQueue::PrioritizedTask TaskGraphWorkQueue::GetNextTaskToRun() {
-  DCHECK(!ready_to_run_namespaces_.empty());
+TaskGraphWorkQueue::PrioritizedTask TaskGraphWorkQueue::GetNextTaskToRun(
+    uint16_t category) {
+  TaskNamespace::Vector& ready_to_run_namespaces =
+      ready_to_run_namespaces_[category];
+  DCHECK(!ready_to_run_namespaces.empty());
 
-  // Take top priority TaskNamespace from |ready_to_run_namespaces_|.
-  std::pop_heap(ready_to_run_namespaces_.begin(),
-                ready_to_run_namespaces_.end(), CompareTaskNamespacePriority);
-  TaskNamespace* task_namespace = ready_to_run_namespaces_.back();
-  ready_to_run_namespaces_.pop_back();
-  DCHECK(!task_namespace->ready_to_run_tasks.empty());
+  // Take top priority TaskNamespace from |ready_to_run_namespaces|.
+  std::pop_heap(ready_to_run_namespaces.begin(), ready_to_run_namespaces.end(),
+                CompareTaskNamespacePriority(category));
+  TaskNamespace* task_namespace = ready_to_run_namespaces.back();
+  ready_to_run_namespaces.pop_back();
+
+  PrioritizedTask::Vector& ready_to_run_tasks =
+      task_namespace->ready_to_run_tasks[category];
+  DCHECK(!ready_to_run_tasks.empty());
 
   // Take top priority task from |ready_to_run_tasks|.
-  std::pop_heap(task_namespace->ready_to_run_tasks.begin(),
-                task_namespace->ready_to_run_tasks.end(), CompareTaskPriority);
-  PrioritizedTask task = task_namespace->ready_to_run_tasks.back();
-  task_namespace->ready_to_run_tasks.pop_back();
+  std::pop_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
+                CompareTaskPriority);
+  PrioritizedTask task = ready_to_run_tasks.back();
+  ready_to_run_tasks.pop_back();
 
-  // Add task namespace back to |ready_to_run_namespaces_| if not empty after
+  // Add task namespace back to |ready_to_run_namespaces| if not empty after
   // taking top priority task.
-  if (!task_namespace->ready_to_run_tasks.empty()) {
-    ready_to_run_namespaces_.push_back(task_namespace);
-    std::push_heap(ready_to_run_namespaces_.begin(),
-                   ready_to_run_namespaces_.end(),
-                   CompareTaskNamespacePriority);
+  if (!ready_to_run_tasks.empty()) {
+    ready_to_run_namespaces.push_back(task_namespace);
+    std::push_heap(ready_to_run_namespaces.begin(),
+                   ready_to_run_namespaces.end(),
+                   CompareTaskNamespacePriority(category));
   }
 
   // Add task to |running_tasks|.
@@ -164,19 +217,26 @@ void TaskGraphWorkQueue::CompleteTask(const PrioritizedTask& completed_task) {
     dependent_node.dependencies--;
     // Task is ready if it has no dependencies. Add it to |ready_to_run_tasks_|.
     if (!dependent_node.dependencies) {
-      bool was_empty = task_namespace->ready_to_run_tasks.empty();
-      task_namespace->ready_to_run_tasks.push_back(PrioritizedTask(
-          dependent_node.task, task_namespace, dependent_node.priority));
-      std::push_heap(task_namespace->ready_to_run_tasks.begin(),
-                     task_namespace->ready_to_run_tasks.end(),
+      PrioritizedTask::Vector& ready_to_run_tasks =
+          task_namespace->ready_to_run_tasks[dependent_node.category];
+
+      bool was_empty = ready_to_run_tasks.empty();
+      ready_to_run_tasks.push_back(
+          PrioritizedTask(dependent_node.task, task_namespace,
+                          dependent_node.category, dependent_node.priority));
+      std::push_heap(ready_to_run_tasks.begin(), ready_to_run_tasks.end(),
                      CompareTaskPriority);
+
       // Task namespace is ready if it has at least one ready to run task. Add
       // it to |ready_to_run_namespaces_| if it just become ready.
       if (was_empty) {
-        DCHECK(std::find(ready_to_run_namespaces_.begin(),
-                         ready_to_run_namespaces_.end(),
-                         task_namespace) == ready_to_run_namespaces_.end());
-        ready_to_run_namespaces_.push_back(task_namespace);
+        TaskNamespace::Vector& ready_to_run_namespaces =
+            ready_to_run_namespaces_[dependent_node.category];
+
+        DCHECK(std::find(ready_to_run_namespaces.begin(),
+                         ready_to_run_namespaces.end(),
+                         task_namespace) == ready_to_run_namespaces.end());
+        ready_to_run_namespaces.push_back(task_namespace);
       }
       ready_to_run_namespaces_has_heap_properties = false;
     }
@@ -185,9 +245,13 @@ void TaskGraphWorkQueue::CompleteTask(const PrioritizedTask& completed_task) {
   // Rearrange the task namespaces in |ready_to_run_namespaces_| in such a way
   // that they yet again form a heap.
   if (!ready_to_run_namespaces_has_heap_properties) {
-    std::make_heap(ready_to_run_namespaces_.begin(),
-                   ready_to_run_namespaces_.end(),
-                   CompareTaskNamespacePriority);
+    for (auto& it : ready_to_run_namespaces_) {
+      uint16_t category = it.first;
+      auto& ready_to_run_namespaces = it.second;
+      std::make_heap(ready_to_run_namespaces.begin(),
+                     ready_to_run_namespaces.end(),
+                     CompareTaskNamespacePriority(category));
+    }
   }
 
   // Finally add task to |completed_tasks_|.
@@ -209,7 +273,7 @@ void TaskGraphWorkQueue::CollectCompletedTasks(NamespaceToken token,
 
   // Remove namespace if finished running tasks.
   DCHECK_EQ(0u, task_namespace.completed_tasks.size());
-  DCHECK_EQ(0u, task_namespace.ready_to_run_tasks.size());
+  DCHECK(!HasReadyToRunTasksInNamespace(&task_namespace));
   DCHECK_EQ(0u, task_namespace.running_tasks.size());
   namespaces_.erase(it);
 }
diff --git a/cc/raster/task_graph_work_queue.h b/cc/raster/task_graph_work_queue.h
index 8a46123..7b6840c 100644
--- a/cc/raster/task_graph_work_queue.h
+++ b/cc/raster/task_graph_work_queue.h
@@ -5,6 +5,7 @@
 #ifndef CC_RASTER_TASK_GRAPH_WORK_QUEUE_H_
 #define CC_RASTER_TASK_GRAPH_WORK_QUEUE_H_
 
+#include <algorithm>
 #include <map>
 #include <vector>
 
@@ -16,6 +17,11 @@ namespace cc {
 // Implements a queue of incoming TaskGraph work. Designed for use by
 // implementations of TaskGraphRunner. Not thread safe, so the caller is
 // responsible for all necessary locking.
+//
+// Tasks in the queue are divided into categories. Tasks from a single graph may
+// be put into different categories, each of which is prioritized independently
+// from the others. It is up to the implementation of TaskGraphRunner to
+// define the meaning of the categories and handle them appropriately.
 class CC_EXPORT TaskGraphWorkQueue {
  public:
   struct TaskNamespace;
@@ -23,12 +29,19 @@ class CC_EXPORT TaskGraphWorkQueue {
   struct PrioritizedTask {
     typedef std::vector<PrioritizedTask> Vector;
 
-    PrioritizedTask(Task* task, TaskNamespace* task_namespace, size_t priority)
-        : task(task), task_namespace(task_namespace), priority(priority) {}
+    PrioritizedTask(Task* task,
+                    TaskNamespace* task_namespace,
+                    uint16_t category,
+                    uint16_t priority)
+        : task(task),
+          task_namespace(task_namespace),
+          category(category),
+          priority(priority) {}
 
     Task* task;
     TaskNamespace* task_namespace;
-    size_t priority;
+    uint16_t category;
+    uint16_t priority;
   };
 
   // Helper classes and static methods used by dependent classes.
@@ -41,8 +54,9 @@ class CC_EXPORT TaskGraphWorkQueue {
     // Current task graph.
     TaskGraph graph;
 
-    // Ordered set of tasks that are ready to run.
-    PrioritizedTask::Vector ready_to_run_tasks;
+    // Map from category to a vector of tasks that are ready to run for that
+    // category.
+    std::map<uint16_t, PrioritizedTask::Vector> ready_to_run_tasks;
 
     // Completed tasks not yet collected by origin thread.
     Task::Vector completed_tasks;
@@ -62,8 +76,8 @@ class CC_EXPORT TaskGraphWorkQueue {
   // previous tasks in the graph being replaced.
   void ScheduleTasks(NamespaceToken token, TaskGraph* graph);
 
-  // Returns the next task to run paired with its namespace.
-  PrioritizedTask GetNextTaskToRun();
+  // Returns the next task to run for the given category.
+  PrioritizedTask GetNextTaskToRun(uint16_t category);
 
   // Marks a task as completed, adding it to its namespace's list of completed
   // tasks and updating the list of |ready_to_run_namespaces|.
@@ -84,13 +98,35 @@ class CC_EXPORT TaskGraphWorkQueue {
     return &it->second;
   }
 
+  static bool HasReadyToRunTasksInNamespace(
+      const TaskNamespace* task_namespace) {
+    return std::find_if(task_namespace->ready_to_run_tasks.begin(),
+                        task_namespace->ready_to_run_tasks.end(),
+                        [](const std::pair<uint16_t, PrioritizedTask::Vector>&
+                               ready_to_run_tasks) {
+                          return !ready_to_run_tasks.second.empty();
+                        }) != task_namespace->ready_to_run_tasks.end();
+  }
+
   static bool HasFinishedRunningTasksInNamespace(
       const TaskNamespace* task_namespace) {
     return task_namespace->running_tasks.empty() &&
-           task_namespace->ready_to_run_tasks.empty();
+           !HasReadyToRunTasksInNamespace(task_namespace);
   }
 
-  bool HasReadyToRunTasks() const { return !ready_to_run_namespaces_.empty(); }
+  bool HasReadyToRunTasks() const {
+    return std::find_if(ready_to_run_namespaces_.begin(),
+                        ready_to_run_namespaces_.end(),
+                        [](const std::pair<uint16_t, TaskNamespace::Vector>&
+                               ready_to_run_namespaces) {
+                          return !ready_to_run_namespaces.second.empty();
+                        }) != ready_to_run_namespaces_.end();
+  }
+
+  bool HasReadyToRunTasksForCategory(uint16_t category) const {
+    auto found = ready_to_run_namespaces_.find(category);
+    return found != ready_to_run_namespaces_.end() && !found->second.empty();
+  }
 
   bool HasAnyNamespaces() const { return !namespaces_.empty(); }
 
@@ -102,6 +138,11 @@ class CC_EXPORT TaskGraphWorkQueue {
                }) == namespaces_.end();
   }
 
+  const std::map<uint16_t, TaskNamespace::Vector>& ready_to_run_namespaces()
+      const {
+    return ready_to_run_namespaces_;
+  }
+
   // Helper function which ensures that graph dependencies were correctly
   // configured.
   static bool DependencyMismatch(const TaskGraph* graph);
@@ -116,29 +157,13 @@ class CC_EXPORT TaskGraphWorkQueue {
     }
   };
 
-  static bool CompareTaskPriority(const PrioritizedTask& a,
-                                  const PrioritizedTask& b) {
-    // In this system, numerically lower priority is run first.
-    return a.priority > b.priority;
-  }
-
-  static bool CompareTaskNamespacePriority(const TaskNamespace* a,
-                                           const TaskNamespace* b) {
-    DCHECK(!a->ready_to_run_tasks.empty());
-    DCHECK(!b->ready_to_run_tasks.empty());
-
-    // Compare based on task priority of the ready_to_run_tasks heap .front()
-    // will hold the max element of the heap, except after pop_heap, when max
-    // element is moved to .back().
-    return CompareTaskPriority(a->ready_to_run_tasks.front(),
-                               b->ready_to_run_tasks.front());
-  }
-
   using TaskNamespaceMap =
       std::map<NamespaceToken, TaskNamespace, CompareToken>;
 
   TaskNamespaceMap namespaces_;
-  TaskNamespace::Vector ready_to_run_namespaces_;
+
+  // Map from category to a vector of ready to run namespaces for that category.
+  std::map<uint16_t, TaskNamespace::Vector> ready_to_run_namespaces_;
 
   // Provides a unique id to each NamespaceToken.
   int next_namespace_id_;
diff --git a/cc/raster/tile_task_worker_pool_perftest.cc b/cc/raster/tile_task_worker_pool_perftest.cc
index 1e1ff1c..7035441 100644
--- a/cc/raster/tile_task_worker_pool_perftest.cc
+++ b/cc/raster/tile_task_worker_pool_perftest.cc
@@ -223,13 +223,14 @@ class TileTaskWorkerPoolPerfTestBase {
 
       for (auto& decode_task : raster_task->dependencies()) {
         graph->nodes.push_back(
-            TaskGraph::Node(decode_task.get(), priority, 0u));
+            TaskGraph::Node(decode_task.get(), 0u /* group */, priority, 0u));
         graph->edges.push_back(
             TaskGraph::Edge(raster_task.get(), decode_task.get()));
       }
 
-      graph->nodes.push_back(TaskGraph::Node(
-          raster_task.get(), priority, raster_task->dependencies().size()));
+      graph->nodes.push_back(
+          TaskGraph::Node(raster_task.get(), 0u /* group */, priority,
+                          raster_task->dependencies().size()));
     }
   }
 
diff --git a/cc/raster/tile_task_worker_pool_unittest.cc b/cc/raster/tile_task_worker_pool_unittest.cc
index c5d3a41..941cd47 100644
--- a/cc/raster/tile_task_worker_pool_unittest.cc
+++ b/cc/raster/tile_task_worker_pool_unittest.cc
@@ -195,7 +195,8 @@ class TileTaskWorkerPoolTest
 
     for (RasterTaskVector::const_iterator it = tasks_.begin();
          it != tasks_.end(); ++it) {
-      graph_.nodes.emplace_back(it->get(), priority++, 0 /* dependencies */);
+      graph_.nodes.emplace_back(it->get(), 0 /* group */, priority++,
+                                0 /* dependencies */);
     }
 
     tile_task_worker_pool_->AsTileTaskRunner()->ScheduleTasks(&graph_);