TaskScheduler: Add methods RemoveSequence() and UpdateSortKey() to PriorityQueue

Bug: 889029
Change-Id: I4fb7cf0011d067c24f26914009681563baf51d11
Reviewed-on: https://chromium-review.googlesource.com/c/1331949
Commit-Queue: Jesse McKenna <jessemckenna@google.com>
Reviewed-by: François Doray <fdoray@chromium.org>
Cr-Commit-Position: refs/heads/master@{#607720}

base/task/task_scheduler/priority_queue.cc

@@ -37,6 +37,7 @@ class PriorityQueue::SequenceAndSortKey {
   // call.
   scoped_refptr<Sequence> take_sequence() {
     DCHECK(sequence_);
+    sequence_->ClearHeapHandle();
     return std::move(sequence_);
   }
 
@@ -47,10 +48,21 @@ class PriorityQueue::SequenceAndSortKey {
   }
 
   // Required by IntrusiveHeap.
-  void SetHeapHandle(const HeapHandle& handle) {}
+  void SetHeapHandle(const HeapHandle& handle) {
+    DCHECK(sequence_);
+    sequence_->SetHeapHandle(handle);
+  }
 
   // Required by IntrusiveHeap.
-  void ClearHeapHandle() {}
+  void ClearHeapHandle() {
+    // Ensure |sequence_| is not nullptr, which may be the case if
+    // take_sequence() was called before this.
+    if (sequence_) {
+      sequence_->ClearHeapHandle();
+    }
+  }
+
+  const Sequence* sequence() const { return sequence_.get(); }
 
   const SequenceSortKey& sort_key() const { return sort_key_; }
 
@@ -92,6 +104,39 @@ scoped_refptr<Sequence> PriorityQueue::Transaction::PopSequence() {
   return sequence;
 }
 
+bool PriorityQueue::Transaction::RemoveSequence(
+    scoped_refptr<Sequence> sequence) {
+  DCHECK(sequence);
+
+  if (IsEmpty())
+    return false;
+
+  const HeapHandle heap_handle = sequence->heap_handle();
+  if (!heap_handle.IsValid())
+    return false;
+
+  DCHECK_EQ(outer_queue_->container_.at(heap_handle).sequence(),
+            sequence.get());
+  outer_queue_->container_.erase(heap_handle);
+  return true;
+}
+
+void PriorityQueue::Transaction::UpdateSortKey(
+    scoped_refptr<Sequence> sequence,
+    const SequenceSortKey& sequence_sort_key) {
+  DCHECK(sequence);
+
+  if (IsEmpty())
+    return;
+
+  const HeapHandle heap_handle = sequence->heap_handle();
+  if (!heap_handle.IsValid())
+    return;
+
+  outer_queue_->container_.ChangeKey(
+      heap_handle, SequenceAndSortKey(std::move(sequence), sequence_sort_key));
+}
+
 bool PriorityQueue::Transaction::IsEmpty() const {
   return outer_queue_->container_.empty();
 }

base/task/task_scheduler/priority_queue.h

@@ -51,6 +51,16 @@ class BASE_EXPORT PriorityQueue {
     // Cannot be called on an empty PriorityQueue.
     scoped_refptr<Sequence> PopSequence();
 
+    // Removes |sequence| from the PriorityQueue. Returns true if successful,
+    // or false if |sequence| is not currently in the PriorityQueue or the
+    // PriorityQueue is empty.
+    bool RemoveSequence(scoped_refptr<Sequence> sequence);
+
+    // Updates the sort key of |sequence| to |sequence_sort_key|. No-ops if
+    // |sequence| is not in the PriorityQueue or the PriorityQueue is empty.
+    void UpdateSortKey(scoped_refptr<Sequence> sequence,
+                       const SequenceSortKey& sequence_sort_key);
+
     // Returns true if the PriorityQueue is empty.
     bool IsEmpty() const;
 

base/task/task_scheduler/priority_queue_unittest.cc

@@ -53,37 +53,43 @@ class ThreadBeginningTransaction : public SimpleThread {
   DISALLOW_COPY_AND_ASSIGN(ThreadBeginningTransaction);
 };
 
-}  // namespace
+scoped_refptr<Sequence> MakeSequenceWithTraitsAndTask(
+    const TaskTraits& traits) {
+  scoped_refptr<Sequence> sequence = MakeRefCounted<Sequence>(traits);
+  sequence->BeginTransaction()->PushTask(
+      Task(FROM_HERE, DoNothing(), TimeDelta()));
+  return sequence;
+}
 
-TEST(TaskSchedulerPriorityQueueTest, PushPopPeek) {
-  // Create test sequences.
+class TaskSchedulerPriorityQueueWithSequencesTest : public testing::Test {
+ protected:
   scoped_refptr<Sequence> sequence_a =
-      MakeRefCounted<Sequence>(TaskTraits(TaskPriority::USER_VISIBLE));
-  sequence_a->BeginTransaction()->PushTask(
-      Task(FROM_HERE, DoNothing(), TimeDelta()));
+      MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_VISIBLE));
   SequenceSortKey sort_key_a = sequence_a->BeginTransaction()->GetSortKey();
 
   scoped_refptr<Sequence> sequence_b =
-      MakeRefCounted<Sequence>(TaskTraits(TaskPriority::USER_BLOCKING));
-  sequence_b->BeginTransaction()->PushTask(
-      Task(FROM_HERE, DoNothing(), TimeDelta()));
+      MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_BLOCKING));
   SequenceSortKey sort_key_b = sequence_b->BeginTransaction()->GetSortKey();
 
   scoped_refptr<Sequence> sequence_c =
-      MakeRefCounted<Sequence>(TaskTraits(TaskPriority::USER_BLOCKING));
-  sequence_c->BeginTransaction()->PushTask(
-      Task(FROM_HERE, DoNothing(), TimeDelta()));
+      MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_BLOCKING));
   SequenceSortKey sort_key_c = sequence_c->BeginTransaction()->GetSortKey();
 
   scoped_refptr<Sequence> sequence_d =
-      MakeRefCounted<Sequence>(TaskTraits(TaskPriority::BEST_EFFORT));
-  sequence_d->BeginTransaction()->PushTask(
-      Task(FROM_HERE, DoNothing(), TimeDelta()));
+      MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::BEST_EFFORT));
   SequenceSortKey sort_key_d = sequence_d->BeginTransaction()->GetSortKey();
 
-  // Create a PriorityQueue and a Transaction.
+ private:
   PriorityQueue pq;
-  auto transaction(pq.BeginTransaction());
+
+ protected:
+  std::unique_ptr<PriorityQueue::Transaction> transaction =
+      pq.BeginTransaction();
+};
+
+}  // namespace
+
+TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, PushPopPeek) {
   EXPECT_TRUE(transaction->IsEmpty());
 
   // Push |sequence_a| in the PriorityQueue. It becomes the sequence with the
@@ -126,6 +132,91 @@ TEST(TaskSchedulerPriorityQueueTest, PushPopPeek) {
   EXPECT_TRUE(transaction->IsEmpty());
 }
 
+TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, RemoveSequence) {
+  EXPECT_TRUE(transaction->IsEmpty());
+
+  // Push all test Sequences into the PriorityQueue. |sequence_b|
+  // will be the sequence with the highest priority.
+  transaction->Push(sequence_a, sort_key_a);
+  transaction->Push(sequence_b, sort_key_b);
+  transaction->Push(sequence_c, sort_key_c);
+  transaction->Push(sequence_d, sort_key_d);
+  EXPECT_EQ(sort_key_b, transaction->PeekSortKey());
+
+  // Remove |sequence_a| from the PriorityQueue. |sequence_b| is still the
+  // sequence with the highest priority.
+  EXPECT_TRUE(transaction->RemoveSequence(sequence_a));
+  EXPECT_EQ(sort_key_b, transaction->PeekSortKey());
+
+  // RemoveSequence() should return false if called on a sequence not in the
+  // PriorityQueue.
+  EXPECT_FALSE(transaction->RemoveSequence(sequence_a));
+
+  // Remove |sequence_b| from the PriorityQueue. |sequence_c| becomes the
+  // sequence with the highest priority.
+  EXPECT_TRUE(transaction->RemoveSequence(sequence_b));
+  EXPECT_EQ(sort_key_c, transaction->PeekSortKey());
+
+  // Remove |sequence_d| from the PriorityQueue. |sequence_c| is still the
+  // sequence with the highest priority.
+  EXPECT_TRUE(transaction->RemoveSequence(sequence_d));
+  EXPECT_EQ(sort_key_c, transaction->PeekSortKey());
+
+  // Remove |sequence_c| from the PriorityQueue, making it empty.
+  EXPECT_TRUE(transaction->RemoveSequence(sequence_c));
+  EXPECT_TRUE(transaction->IsEmpty());
+
+  // Return false if RemoveSequence() is called on an empty PriorityQueue.
+  EXPECT_FALSE(transaction->RemoveSequence(sequence_c));
+}
+
+TEST_F(TaskSchedulerPriorityQueueWithSequencesTest, UpdateSortKey) {
+  EXPECT_TRUE(transaction->IsEmpty());
+
+  // Push all test Sequences into the PriorityQueue. |sequence_b|
+  // becomes the sequence with the highest priority.
+  transaction->Push(sequence_a, sort_key_a);
+  transaction->Push(sequence_b, sort_key_b);
+  transaction->Push(sequence_c, sort_key_c);
+  transaction->Push(sequence_d, sort_key_d);
+  EXPECT_EQ(sort_key_b, transaction->PeekSortKey());
+
+  // Downgrade |sequence_b| from USER_BLOCKING to BEST_EFFORT. |sequence_c|
+  // (USER_BLOCKING priority) becomes the sequence with the highest priority.
+  const SequenceSortKey best_effort_sort_key =
+      SequenceSortKey(TaskPriority::BEST_EFFORT, TimeTicks::Now());
+  transaction->UpdateSortKey(sequence_b, best_effort_sort_key);
+  EXPECT_EQ(sort_key_c, transaction->PeekSortKey());
+
+  // Update |sequence_c|'s sort key to one with the same priority. |sequence_c|
+  // (USER_BLOCKING priority) is still the sequence with the highest priority.
+  const SequenceSortKey user_blocking_sort_key =
+      SequenceSortKey(TaskPriority::USER_BLOCKING, TimeTicks::Now());
+  transaction->UpdateSortKey(sequence_c, user_blocking_sort_key);
+  // 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, transaction->PopSequence());
+  EXPECT_EQ(sort_key_a, transaction->PeekSortKey());
+
+  // Upgrade |sequence_d| from BEST_EFFORT to USER_BLOCKING. |sequence_d|
+  // becomes the sequence with the highest priority.
+  transaction->UpdateSortKey(sequence_d, user_blocking_sort_key);
+  // Note: |sequence_d| is popped for comparison as |sort_key_d| becomes
+  // obsolete.
+  EXPECT_EQ(sequence_d, transaction->PopSequence());
+
+  // No-op if UpdateSortKey() is called on a Sequence not in the PriorityQueue.
+  EXPECT_EQ(sort_key_a, transaction->PeekSortKey());
+  transaction->UpdateSortKey(sequence_d, user_blocking_sort_key);
+  EXPECT_EQ(sequence_a, transaction->PopSequence());
+  EXPECT_EQ(sequence_b, transaction->PopSequence());
+
+  // No-op if UpdateSortKey() is called on an empty PriorityQueue.
+  transaction->UpdateSortKey(sequence_b, sort_key_b);
+  EXPECT_TRUE(transaction->IsEmpty());
+}
+
 // Check that creating Transactions on the same thread for 2 unrelated
 // PriorityQueues causes a crash.
 TEST(TaskSchedulerPriorityQueueTest, IllegalTwoTransactionsSameThread) {

base/task/task_scheduler/sequence.cc

@@ -67,6 +67,14 @@ SequenceSortKey Sequence::Transaction::GetSortKey() const {
                          next_task_sequenced_time);
 }
 
+void Sequence::SetHeapHandle(const HeapHandle& handle) {
+  heap_handle_ = handle;
+}
+
+void Sequence::ClearHeapHandle() {
+  heap_handle_ = HeapHandle();
+}
+
 Sequence::Sequence(
     const TaskTraits& traits,
     scoped_refptr<SchedulerParallelTaskRunner> scheduler_parallel_task_runner)

base/task/task_scheduler/sequence.h

@@ -13,6 +13,7 @@
 #include "base/memory/ref_counted.h"
 #include "base/optional.h"
 #include "base/sequence_token.h"
+#include "base/task/common/intrusive_heap.h"
 #include "base/task/task_scheduler/scheduler_lock.h"
 #include "base/task/task_scheduler/scheduler_parallel_task_runner.h"
 #include "base/task/task_scheduler/sequence_sort_key.h"
@@ -75,6 +76,8 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
     // Cannot be called on an empty Sequence.
     SequenceSortKey GetSortKey() const;
 
+    scoped_refptr<Sequence> sequence() { return sequence_; }
+
    private:
     friend class Sequence;
 
@@ -96,6 +99,11 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
   // active Sequence::Transaction.
   std::unique_ptr<Transaction> BeginTransaction();
 
+  // Support for IntrusiveHeap.
+  void SetHeapHandle(const HeapHandle& handle);
+  void ClearHeapHandle();
+  HeapHandle heap_handle() const { return heap_handle_; }
+
   // Returns a token that uniquely identifies this Sequence.
   const SequenceToken& token() const { return token_; }
 
@@ -130,6 +138,10 @@ class BASE_EXPORT Sequence : public RefCountedThreadSafe<Sequence> {
   const scoped_refptr<SchedulerParallelTaskRunner>
       scheduler_parallel_task_runner_;
 
+  // The Sequence's position in its current PriorityQueue. Access is protected
+  // by the PriorityQueue's lock.
+  HeapHandle heap_handle_;
+
   DISALLOW_COPY_AND_ASSIGN(Sequence);
 };