TaskScheduler: Track number of Sequences per priority.

With this CL, PriorityQueue keeps track of the number of Sequences
of each TaskPriority that it contains. This will be used in an upcoming
CL to ensure that the appropriate number of workers are woken up when
new Sequences are scheduled or when the maxmium number of tasks that
can run concurrently is updated.

Bug: 889029
Change-Id: I644dc3fb5d59f9f408b99d35af8a8187b03cf313
Reviewed-on: https://chromium-review.googlesource.com/c/1452484
Commit-Queue: François Doray <fdoray@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#629606}

base/task/task_scheduler/priority_queue.cc

@@ -8,6 +8,7 @@
 
 #include "base/logging.h"
 #include "base/memory/ptr_util.h"
+#include "base/stl_util.h"
 
 namespace base {
 namespace internal {
@@ -92,6 +93,7 @@ PriorityQueue::~PriorityQueue() {
 void PriorityQueue::Push(scoped_refptr<Sequence> sequence,
                          const SequenceSortKey& sequence_sort_key) {
   container_.insert(SequenceAndSortKey(std::move(sequence), sequence_sort_key));
+  IncrementNumSequencesForPriority(sequence_sort_key.priority());
 }
 
 const SequenceSortKey& PriorityQueue::PeekSortKey() const {
@@ -102,12 +104,13 @@ const SequenceSortKey& PriorityQueue::PeekSortKey() const {
 scoped_refptr<Sequence> PriorityQueue::PopSequence() {
   DCHECK(!IsEmpty());
 
-  // The const_cast on top() is okay since the SequenceAndSortKey is
+  // The const_cast on Min() is okay since the SequenceAndSortKey is
   // transactionally being popped from |container_| right after and taking its
   // Sequence does not alter its sort order.
-  scoped_refptr<Sequence> sequence =
-      const_cast<PriorityQueue::SequenceAndSortKey&>(container_.Min())
-          .take_sequence();
+  auto& sequence_and_sort_key =
+      const_cast<PriorityQueue::SequenceAndSortKey&>(container_.Min());
+  DecrementNumSequencesForPriority(sequence_and_sort_key.sort_key().priority());
+  scoped_refptr<Sequence> sequence = sequence_and_sort_key.take_sequence();
   container_.Pop();
   return sequence;
 }
@@ -122,7 +125,10 @@ bool PriorityQueue::RemoveSequence(scoped_refptr<Sequence> sequence) {
   if (!heap_handle.IsValid())
     return false;
 
-  DCHECK_EQ(container_.at(heap_handle).sequence(), sequence.get());
+  const SequenceAndSortKey& sequence_and_sort_key = container_.at(heap_handle);
+  DCHECK_EQ(sequence_and_sort_key.sequence(), sequence.get());
+
+  DecrementNumSequencesForPriority(sequence_and_sort_key.sort_key().priority());
   container_.erase(heap_handle);
   return true;
 }
@@ -139,10 +145,16 @@ void PriorityQueue::UpdateSortKey(
   if (!heap_handle.IsValid())
     return;
 
-  auto sort_key = sequence_and_transaction.transaction.GetSortKey();
+  auto old_sort_key = container_.at(heap_handle).sort_key();
+  auto new_sort_key = sequence_and_transaction.transaction.GetSortKey();
+
+  DecrementNumSequencesForPriority(old_sort_key.priority());
+  IncrementNumSequencesForPriority(new_sort_key.priority());
+
   container_.ChangeKey(
-      heap_handle, SequenceAndSortKey(
-                       std::move(sequence_and_transaction.sequence), sort_key));
+      heap_handle,
+      SequenceAndSortKey(std::move(sequence_and_transaction.sequence),
+                         new_sort_key));
 }
 
 bool PriorityQueue::IsEmpty() const {
@@ -158,5 +170,14 @@ void PriorityQueue::EnableFlushSequencesOnDestroyForTesting() {
   is_flush_sequences_on_destroy_enabled_ = true;
 }
 
+void PriorityQueue::DecrementNumSequencesForPriority(TaskPriority priority) {
+  DCHECK_GT(num_sequences_per_priority_[static_cast<int>(priority)], 0U);
+  --num_sequences_per_priority_[static_cast<int>(priority)];
+}
+
+void PriorityQueue::IncrementNumSequencesForPriority(TaskPriority priority) {
+  ++num_sequences_per_priority_[static_cast<int>(priority)];
+}
+
 }  // namespace internal
 }  // namespace base

base/task/task_scheduler/priority_queue.h

@@ -58,6 +58,11 @@ class BASE_EXPORT PriorityQueue {
   // Returns the number of Sequences in the PriorityQueue.
   size_t Size() const;
 
+  // Returns the number of Sequences with |priority|.
+  size_t GetNumSequencesWithPriority(TaskPriority priority) const {
+    return num_sequences_per_priority_[static_cast<int>(priority)];
+  }
+
   // Set the PriorityQueue to empty all its Sequences of Tasks when it is
   // destroyed; needed to prevent memory leaks caused by a reference cycle
   // (Sequence -> Task -> TaskRunner -> Sequence...) during test teardown.
@@ -70,8 +75,14 @@ class BASE_EXPORT PriorityQueue {
 
   using ContainerType = IntrusiveHeap<SequenceAndSortKey>;
 
+  void DecrementNumSequencesForPriority(TaskPriority priority);
+  void IncrementNumSequencesForPriority(TaskPriority priority);
+
   ContainerType container_;
 
+  size_t num_sequences_per_priority_[static_cast<int>(TaskPriority::HIGHEST) +
+                                     1] = {};
+
   // Should only be enabled by EnableFlushSequencesOnDestroyForTesting().
   bool is_flush_sequences_on_destroy_enabled_ = false;
 

base/task/task_scheduler/priority_queue_unittest.cc

@@ -33,6 +33,17 @@ scoped_refptr<Sequence> MakeSequenceWithTraitsAndTask(
 
 class TaskSchedulerPriorityQueueWithSequencesTest : public testing::Test {
  protected:
+  void ExpectNumSequences(size_t num_best_effort,
+                          size_t num_user_visible,
+                          size_t num_user_blocking) {
+    EXPECT_EQ(pq.GetNumSequencesWithPriority(TaskPriority::BEST_EFFORT),
+              num_best_effort);
+    EXPECT_EQ(pq.GetNumSequencesWithPriority(TaskPriority::USER_VISIBLE),
+              num_user_visible);
+    EXPECT_EQ(pq.GetNumSequencesWithPriority(TaskPriority::USER_BLOCKING),
+              num_user_blocking);
+  }
+
   scoped_refptr<Sequence> sequence_a =
       MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_VISIBLE));
   SequenceSortKey sort_key_a = sequence_a->BeginTransaction().GetSortKey();
@@ -56,45 +67,54 @@ class TaskSchedulerPriorityQueueWithSequencesTest : public testing::Test {
 
 TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, PushPopPeek) {
   EXPECT_TRUE(pq.IsEmpty());
+  ExpectNumSequences(0U, 0U, 0U);
 
   // Push |sequence_a| in the PriorityQueue. It becomes the sequence with the
   // highest priority.
   pq.Push(sequence_a, sort_key_a);
   EXPECT_EQ(sort_key_a, pq.PeekSortKey());
+  ExpectNumSequences(0U, 1U, 0U);
 
   // Push |sequence_b| in the PriorityQueue. It becomes the sequence with the
   // highest priority.
   pq.Push(sequence_b, sort_key_b);
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(0U, 1U, 1U);
 
   // Push |sequence_c| in the PriorityQueue. |sequence_b| is still the sequence
   // with the highest priority.
   pq.Push(sequence_c, sort_key_c);
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(0U, 1U, 2U);
 
   // Push |sequence_d| in the PriorityQueue. |sequence_b| is still the sequence
   // with the highest priority.
   pq.Push(sequence_d, sort_key_d);
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(1U, 1U, 2U);
 
   // Pop |sequence_b| from the PriorityQueue. |sequence_c| becomes the sequence
   // with the highest priority.
   EXPECT_EQ(sequence_b, pq.PopSequence());
   EXPECT_EQ(sort_key_c, pq.PeekSortKey());
+  ExpectNumSequences(1U, 1U, 1U);
 
   // Pop |sequence_c| from the PriorityQueue. |sequence_a| becomes the sequence
   // with the highest priority.
   EXPECT_EQ(sequence_c, pq.PopSequence());
   EXPECT_EQ(sort_key_a, pq.PeekSortKey());
+  ExpectNumSequences(1U, 1U, 0U);
 
   // Pop |sequence_a| from the PriorityQueue. |sequence_d| becomes the sequence
   // with the highest priority.
   EXPECT_EQ(sequence_a, pq.PopSequence());
   EXPECT_EQ(sort_key_d, pq.PeekSortKey());
+  ExpectNumSequences(1U, 0U, 0U);
 
   // Pop |sequence_d| from the PriorityQueue. It is now empty.
   EXPECT_EQ(sequence_d, pq.PopSequence());
   EXPECT_TRUE(pq.IsEmpty());
+  ExpectNumSequences(0U, 0U, 0U);
 }
 
 TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, RemoveSequence) {
@@ -107,32 +127,39 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, RemoveSequence) {
   pq.Push(sequence_c, sort_key_c);
   pq.Push(sequence_d, sort_key_d);
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(1U, 1U, 2U);
 
   // Remove |sequence_a| from the PriorityQueue. |sequence_b| is still the
   // sequence with the highest priority.
   EXPECT_TRUE(pq.RemoveSequence(sequence_a));
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(1U, 0U, 2U);
 
   // RemoveSequence() should return false if called on a sequence not in the
   // PriorityQueue.
   EXPECT_FALSE(pq.RemoveSequence(sequence_a));
+  ExpectNumSequences(1U, 0U, 2U);
 
   // Remove |sequence_b| from the PriorityQueue. |sequence_c| becomes the
   // sequence with the highest priority.
   EXPECT_TRUE(pq.RemoveSequence(sequence_b));
   EXPECT_EQ(sort_key_c, pq.PeekSortKey());
+  ExpectNumSequences(1U, 0U, 1U);
 
   // Remove |sequence_d| from the PriorityQueue. |sequence_c| is still the
   // sequence with the highest priority.
   EXPECT_TRUE(pq.RemoveSequence(sequence_d));
   EXPECT_EQ(sort_key_c, pq.PeekSortKey());
+  ExpectNumSequences(0U, 0U, 1U);
 
   // Remove |sequence_c| from the PriorityQueue, making it empty.
   EXPECT_TRUE(pq.RemoveSequence(sequence_c));
   EXPECT_TRUE(pq.IsEmpty());
+  ExpectNumSequences(0U, 0U, 0U);
 
   // Return false if RemoveSequence() is called on an empty PriorityQueue.
   EXPECT_FALSE(pq.RemoveSequence(sequence_c));
+  ExpectNumSequences(0U, 0U, 0U);
 }
 
 TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
@@ -145,6 +172,7 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
   pq.Push(sequence_c, sort_key_c);
   pq.Push(sequence_d, sort_key_d);
   EXPECT_EQ(sort_key_b, pq.PeekSortKey());
+  ExpectNumSequences(1U, 1U, 2U);
 
   {
     // Downgrade |sequence_b| from USER_BLOCKING to BEST_EFFORT. |sequence_c|
@@ -156,6 +184,7 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
 
     pq.UpdateSortKey(std::move(sequence_b_and_transaction));
     EXPECT_EQ(sort_key_c, pq.PeekSortKey());
+    ExpectNumSequences(2U, 1U, 1U);
   }
 
   {
@@ -168,12 +197,14 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
         TaskPriority::USER_BLOCKING);
 
     pq.UpdateSortKey(std::move(sequence_c_and_transaction));
+    ExpectNumSequences(2U, 1U, 1U);
 
     // Note: |sequence_c| is popped for comparison as |sort_key_c| becomes
     // obsolete. |sequence_a| (USER_VISIBLE priority) becomes the sequence with
     // the highest priority.
     EXPECT_EQ(sequence_c, pq.PopSequence());
     EXPECT_EQ(sort_key_a, pq.PeekSortKey());
+    ExpectNumSequences(2U, 1U, 0U);
   }
 
   {
@@ -185,6 +216,7 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
         TaskPriority::USER_BLOCKING);
 
     pq.UpdateSortKey(std::move(sequence_d_and_transaction));
+    ExpectNumSequences(1U, 1U, 1U);
 
     // Note: |sequence_d| is popped for comparison as |sort_key_d| becomes
     // obsolete.
@@ -192,6 +224,7 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
     // No-op if UpdateSortKey() is called on a Sequence not in the
     // PriorityQueue.
     EXPECT_EQ(sort_key_a, pq.PeekSortKey());
+    ExpectNumSequences(1U, 1U, 0U);
   }
 
   {
@@ -199,8 +232,11 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
         SequenceAndTransaction::FromSequence(sequence_d);
 
     pq.UpdateSortKey(std::move(sequence_d_and_transaction));
+    ExpectNumSequences(1U, 1U, 0U);
     EXPECT_EQ(sequence_a, pq.PopSequence());
+    ExpectNumSequences(1U, 0U, 0U);
     EXPECT_EQ(sequence_b, pq.PopSequence());
+    ExpectNumSequences(0U, 0U, 0U);
   }
 
   {
@@ -209,6 +245,7 @@ TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
         SequenceAndTransaction::FromSequence(sequence_b);
     pq.UpdateSortKey(std::move(sequence_b_and_transaction));
     EXPECT_TRUE(pq.IsEmpty());
+    ExpectNumSequences(0U, 0U, 0U);
   }
 }