[Jobs API]: Back pressure on many workers.

queue_time for Jobs is adjusted based on number of workers to put
back pressure on Jobs that use many workers. This prioritizes newer work
that uses less workers (or simple sequences).
To achieve this, GetSortKey is extracted from Transaction and uses
atomic load. This is safe since the result of GetSortKey and the queue
itself is only modified under lock.

Follow-up: consider integrating full TaskSourceSortKey in ShouldYield

Change-Id: If26e00b40502b0a4e09f10e3627de8b7e4a52488
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2377893
Commit-Queue: Etienne Pierre-Doray <etiennep@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Cr-Commit-Position: refs/heads/master@{#807090}

base/task/thread_pool/job_task_source.cc

@@ -95,7 +95,7 @@ JobTaskSource::JobTaskSource(const Location& from_here,
             self->worker_task_.Run(&job_delegate);
           },
           base::Unretained(this))),
-      queue_time_(TimeTicks::Now()),
+      ready_time_(TimeTicks::Now()),
       delegate_(delegate) {
   DCHECK(delegate_);
 }
@@ -348,7 +348,8 @@ bool JobTaskSource::DidProcessTask(TaskSource::Transaction* /*transaction*/) {
 }
 
 TaskSourceSortKey JobTaskSource::GetSortKey() const {
-  return TaskSourceSortKey(traits_.priority(), queue_time_);
+  return TaskSourceSortKey(priority_racy(), ready_time_,
+                           TS_UNCHECKED_READ(state_).Load().worker_count());
 }
 
 Task JobTaskSource::Clear(TaskSource::Transaction* transaction) {

base/task/thread_pool/job_task_source.h

@@ -68,6 +68,7 @@ class BASE_EXPORT JobTaskSource : public TaskSource {
   // TaskSource:
   ExecutionEnvironment GetExecutionEnvironment() override;
   size_t GetRemainingConcurrency() const override;
+  TaskSourceSortKey GetSortKey() const override;
 
   bool IsCompleted() const;
   size_t GetWorkerCount() const;
@@ -182,7 +183,6 @@ class BASE_EXPORT JobTaskSource : public TaskSource {
   Task TakeTask(TaskSource::Transaction* transaction) override;
   Task Clear(TaskSource::Transaction* transaction) override;
   bool DidProcessTask(TaskSource::Transaction* transaction) override;
-  TaskSourceSortKey GetSortKey() const override;
 
   // Synchronizes access to workers state.
   mutable CheckedLock worker_lock_{UniversalSuccessor()};
@@ -207,7 +207,7 @@ class BASE_EXPORT JobTaskSource : public TaskSource {
   // Task returned from TakeTask(), that calls |worker_task_| internally.
   RepeatingClosure primary_task_;
 
-  const TimeTicks queue_time_;
+  const TimeTicks ready_time_;
   PooledTaskRunnerDelegate* delegate_;
 
   DISALLOW_COPY_AND_ASSIGN(JobTaskSource);

base/task/thread_pool/job_task_source_unittest.cc

@@ -203,6 +203,7 @@ TEST_F(ThreadPoolJobTaskSourceTest, RunTasksInParallel) {
   EXPECT_EQ(registered_task_source_a.WillRunTask(),
             TaskSource::RunStatus::kAllowedNotSaturated);
   EXPECT_EQ(1U, task_source->GetWorkerCount());
+  EXPECT_EQ(1U, task_source->GetSortKey().worker_count());
   auto task_a = registered_task_source_a.TakeTask();
 
   auto registered_task_source_b =
@@ -210,6 +211,7 @@ TEST_F(ThreadPoolJobTaskSourceTest, RunTasksInParallel) {
   EXPECT_EQ(registered_task_source_b.WillRunTask(),
             TaskSource::RunStatus::kAllowedSaturated);
   EXPECT_EQ(2U, task_source->GetWorkerCount());
+  EXPECT_EQ(2U, task_source->GetSortKey().worker_count());
   auto task_b = registered_task_source_b.TakeTask();
 
   // WillRunTask() should return a null RunStatus once the max concurrency is
@@ -222,9 +224,11 @@ TEST_F(ThreadPoolJobTaskSourceTest, RunTasksInParallel) {
   // source to re-enqueue.
   job_task->SetNumTasksToRun(2);
   EXPECT_TRUE(registered_task_source_a.DidProcessTask());
+  EXPECT_EQ(1U, task_source->GetSortKey().worker_count());
 
   std::move(task_b.task).Run();
   EXPECT_TRUE(registered_task_source_b.DidProcessTask());
+  EXPECT_EQ(0U, task_source->GetSortKey().worker_count());
 
   EXPECT_EQ(0U, task_source->GetWorkerCount());
 

base/task/thread_pool/priority_queue.cc

@@ -100,7 +100,7 @@ PriorityQueue& PriorityQueue::operator=(PriorityQueue&& other) = default;
 void PriorityQueue::Push(
     TransactionWithRegisteredTaskSource transaction_with_task_source) {
   auto task_source_sort_key =
-      transaction_with_task_source.transaction.GetSortKey();
+      transaction_with_task_source.task_source->GetSortKey();
   container_.insert(
       TaskSourceAndSortKey(std::move(transaction_with_task_source.task_source),
                            task_source_sort_key));
@@ -160,29 +160,27 @@ RegisteredTaskSource PriorityQueue::RemoveTaskSource(
   return registered_task_source;
 }
 
-void PriorityQueue::UpdateSortKey(TaskSource::Transaction transaction) {
-  DCHECK(transaction);
-
+void PriorityQueue::UpdateSortKey(const TaskSource& task_source,
+                                  TaskSourceSortKey sort_key) {
   if (IsEmpty())
     return;
 
-  const HeapHandle heap_handle = transaction.task_source()->heap_handle();
+  const HeapHandle heap_handle = task_source.heap_handle();
   if (!heap_handle.IsValid())
     return;
 
   auto old_sort_key = container_.at(heap_handle).sort_key();
-  auto new_sort_key = transaction.GetSortKey();
   auto registered_task_source =
       const_cast<PriorityQueue::TaskSourceAndSortKey&>(
           container_.at(heap_handle))
           .take_task_source();
 
   DecrementNumTaskSourcesForPriority(old_sort_key.priority());
-  IncrementNumTaskSourcesForPriority(new_sort_key.priority());
+  IncrementNumTaskSourcesForPriority(sort_key.priority());
 
   container_.ChangeKey(
       heap_handle,
-      TaskSourceAndSortKey(std::move(registered_task_source), new_sort_key));
+      TaskSourceAndSortKey(std::move(registered_task_source), sort_key));
 }
 
 bool PriorityQueue::IsEmpty() const {

base/task/thread_pool/priority_queue.h

@@ -51,10 +51,10 @@ class BASE_EXPORT PriorityQueue {
   // empty.
   RegisteredTaskSource RemoveTaskSource(const TaskSource& task_source);
 
-  // Updates the sort key of the TaskSource in |transaction| to
-  // match its current traits. No-ops if the TaskSource is not in the
-  // PriorityQueue or the PriorityQueue is empty.
-  void UpdateSortKey(TaskSource::Transaction transaction);
+  // Updates the sort key of |task_source| to |sort_key|, reordering
+  // |task_source| in the queue if necessary. No-ops if the TaskSource is not in
+  // the PriorityQueue or the PriorityQueue is empty.
+  void UpdateSortKey(const TaskSource& task_source, TaskSourceSortKey sort_key);
 
   // Returns true if the PriorityQueue is empty.
   bool IsEmpty() const;

base/task/thread_pool/priority_queue_unittest.cc

@@ -59,19 +59,19 @@ class PriorityQueueWithSequencesTest : public testing::Test {
 
   scoped_refptr<TaskSource> sequence_a =
       MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_VISIBLE));
-  TaskSourceSortKey sort_key_a = sequence_a->BeginTransaction().GetSortKey();
+  TaskSourceSortKey sort_key_a = sequence_a->GetSortKey();
 
   scoped_refptr<TaskSource> sequence_b =
       MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_BLOCKING));
-  TaskSourceSortKey sort_key_b = sequence_b->BeginTransaction().GetSortKey();
+  TaskSourceSortKey sort_key_b = sequence_b->GetSortKey();
 
   scoped_refptr<TaskSource> sequence_c =
       MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::USER_BLOCKING));
-  TaskSourceSortKey sort_key_c = sequence_c->BeginTransaction().GetSortKey();
+  TaskSourceSortKey sort_key_c = sequence_c->GetSortKey();
 
   scoped_refptr<TaskSource> sequence_d =
       MakeSequenceWithTraitsAndTask(TaskTraits(TaskPriority::BEST_EFFORT));
-  TaskSourceSortKey sort_key_d = sequence_d->BeginTransaction().GetSortKey();
+  TaskSourceSortKey sort_key_d = sequence_d->GetSortKey();
 
   PriorityQueue pq;
 };
@@ -193,7 +193,7 @@ TEST_F(PriorityQueueWithSequencesTest, UpdateSortKey) {
     auto sequence_b_transaction = sequence_b->BeginTransaction();
     sequence_b_transaction.UpdatePriority(TaskPriority::BEST_EFFORT);
 
-    pq.UpdateSortKey(std::move(sequence_b_transaction));
+    pq.UpdateSortKey(*sequence_b, sequence_b->GetSortKey());
     EXPECT_EQ(sort_key_c, pq.PeekSortKey());
     ExpectNumSequences(2U, 1U, 1U);
   }
@@ -205,7 +205,7 @@ TEST_F(PriorityQueueWithSequencesTest, UpdateSortKey) {
     auto sequence_c_transaction = sequence_c->BeginTransaction();
     sequence_c_transaction.UpdatePriority(TaskPriority::USER_BLOCKING);
 
-    pq.UpdateSortKey(std::move(sequence_c_transaction));
+    pq.UpdateSortKey(*sequence_c, sequence_c->GetSortKey());
     ExpectNumSequences(2U, 1U, 1U);
 
     // Note: |sequence_c| is popped for comparison as |sort_key_c| becomes
@@ -222,7 +222,7 @@ TEST_F(PriorityQueueWithSequencesTest, UpdateSortKey) {
     auto sequence_d_and_transaction = sequence_d->BeginTransaction();
     sequence_d_and_transaction.UpdatePriority(TaskPriority::USER_BLOCKING);
 
-    pq.UpdateSortKey(std::move(sequence_d_and_transaction));
+    pq.UpdateSortKey(*sequence_d, sequence_d->GetSortKey());
     ExpectNumSequences(1U, 1U, 1U);
 
     // Note: |sequence_d| is popped for comparison as |sort_key_d| becomes
@@ -235,7 +235,7 @@ TEST_F(PriorityQueueWithSequencesTest, UpdateSortKey) {
   }
 
   {
-    pq.UpdateSortKey(sequence_d->BeginTransaction());
+    pq.UpdateSortKey(*sequence_d, sequence_d->GetSortKey());
     ExpectNumSequences(1U, 1U, 0U);
     EXPECT_EQ(sequence_a, pq.PopTaskSource().Unregister());
     ExpectNumSequences(1U, 0U, 0U);
@@ -245,7 +245,7 @@ TEST_F(PriorityQueueWithSequencesTest, UpdateSortKey) {
 
   {
     // No-op if UpdateSortKey() is called on an empty PriorityQueue.
-    pq.UpdateSortKey(sequence_b->BeginTransaction());
+    pq.UpdateSortKey(*sequence_b, sequence_b->GetSortKey());
     EXPECT_TRUE(pq.IsEmpty());
     ExpectNumSequences(0U, 0U, 0U);
   }

base/task/thread_pool/sequence.cc

@@ -49,6 +49,8 @@ void Sequence::Transaction::PushTask(Task task) {
                                         std::move(task.task))
                   : std::move(task.task);
 
+  if (sequence()->queue_.empty())
+    sequence()->ready_time_.store(task.queue_time, std::memory_order_relaxed);
   sequence()->queue_.push(std::move(task));
 
   // AddRef() matched by manual Release() when the sequence has no more tasks
@@ -83,6 +85,9 @@ Task Sequence::TakeTask(TaskSource::Transaction* transaction) {
 
   auto next_task = std::move(queue_.front());
   queue_.pop();
+  if (!queue_.empty()) {
+    ready_time_.store(queue_.front().queue_time, std::memory_order_relaxed);
+  }
   return next_task;
 }
 
@@ -104,8 +109,8 @@ bool Sequence::DidProcessTask(TaskSource::Transaction* transaction) {
 }
 
 TaskSourceSortKey Sequence::GetSortKey() const {
-  DCHECK(!queue_.empty());
-  return TaskSourceSortKey(traits_.priority(), queue_.front().queue_time);
+  return TaskSourceSortKey(priority_racy(),
+                           ready_time_.load(std::memory_order_relaxed));
 }
 
 Task Sequence::Clear(TaskSource::Transaction* transaction) {

base/task/thread_pool/sequence.h

@@ -86,6 +86,7 @@ class BASE_EXPORT Sequence : public TaskSource {
   // TaskSource:
   ExecutionEnvironment GetExecutionEnvironment() override;
   size_t GetRemainingConcurrency() const override;
+  TaskSourceSortKey GetSortKey() const override;
 
   // Returns a token that uniquely identifies this Sequence.
   const SequenceToken& token() const { return token_; }
@@ -102,7 +103,6 @@ class BASE_EXPORT Sequence : public TaskSource {
   Task TakeTask(TaskSource::Transaction* transaction) override;
   Task Clear(TaskSource::Transaction* transaction) override;
   bool DidProcessTask(TaskSource::Transaction* transaction) override;
-  TaskSourceSortKey GetSortKey() const override;
 
   // Releases reference to TaskRunner.
   void ReleaseTaskRunner();
@@ -112,6 +112,8 @@ class BASE_EXPORT Sequence : public TaskSource {
   // Queue of tasks to execute.
   base::queue<Task> queue_;
 
+  std::atomic<TimeTicks> ready_time_{TimeTicks()};
+
   // True if a worker is currently associated with a Task from this Sequence.
   bool has_worker_ = false;
 

base/task/thread_pool/sequence_unittest.cc

@@ -128,7 +128,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyBestEffort) {
 
   // Get the sort key.
   const TaskSourceSortKey best_effort_sort_key =
-      best_effort_sequence_transaction.GetSortKey();
+      best_effort_sequence->GetSortKey();
 
   // Take the task from the sequence, so that its sequenced time is available
   // for the check below.
@@ -141,7 +141,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyBestEffort) {
   // Verify the sort key.
   EXPECT_EQ(TaskPriority::BEST_EFFORT, best_effort_sort_key.priority());
   EXPECT_EQ(take_best_effort_task.queue_time,
-            best_effort_sort_key.next_task_sequenced_time());
+            best_effort_sort_key.ready_time());
 
   // DidProcessTask for correctness.
   best_effort_registered_task_source.DidProcessTask(
@@ -162,7 +162,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyForeground) {
 
   // Get the sort key.
   const TaskSourceSortKey foreground_sort_key =
-      foreground_sequence_transaction.GetSortKey();
+      foreground_sequence->GetSortKey();
 
   // Take the task from the sequence, so that its sequenced time is available
   // for the check below.
@@ -174,8 +174,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyForeground) {
 
   // Verify the sort key.
   EXPECT_EQ(TaskPriority::USER_VISIBLE, foreground_sort_key.priority());
-  EXPECT_EQ(take_foreground_task.queue_time,
-            foreground_sort_key.next_task_sequenced_time());
+  EXPECT_EQ(take_foreground_task.queue_time, foreground_sort_key.ready_time());
 
   // DidProcessTask for correctness.
   foreground_registered_task_source.DidProcessTask(

base/task/thread_pool/task_source.cc

@@ -32,10 +32,6 @@ TaskSource::Transaction::~Transaction() {
   }
 }
 
-TaskSourceSortKey TaskSource::Transaction::GetSortKey() const {
-  return task_source_->GetSortKey();
-}
-
 void TaskSource::Transaction::UpdatePriority(TaskPriority priority) {
   if (FeatureList::IsEnabled(kAllTasksUserBlocking))
     return;

base/task/thread_pool/task_source.h

@@ -105,10 +105,6 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
 
     operator bool() const { return !!task_source_; }
 
-    // Returns a TaskSourceSortKey representing the priority of the TaskSource.
-    // Cannot be called on an empty TaskSource.
-    TaskSourceSortKey GetSortKey() const;
-
     // Sets TaskSource priority to |priority|.
     void UpdatePriority(TaskPriority priority);
 
@@ -148,6 +144,9 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
   // are needed. This may be called on an empty task source.
   virtual size_t GetRemainingConcurrency() const = 0;
 
+  // Returns a TaskSourceSortKey representing the priority of the TaskSource.
+  virtual TaskSourceSortKey GetSortKey() const = 0;
+
   // Support for IntrusiveHeap.
   void SetHeapHandle(const HeapHandle& handle);
   void ClearHeapHandle();
@@ -193,8 +192,6 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
   // are concurrently ready.
   virtual Task Clear(TaskSource::Transaction* transaction) = 0;
 
-  virtual TaskSourceSortKey GetSortKey() const = 0;
-
   // Sets TaskSource priority to |priority|.
   void UpdatePriority(TaskPriority priority);
 

base/task/thread_pool/task_source_sort_key.cc

@@ -7,14 +7,20 @@
 namespace base {
 namespace internal {
 
+static_assert(sizeof(TaskSourceSortKey) <= 2 * sizeof(uint64_t),
+              "Members in TaskSourceSortKey should be ordered to be compact.");
+
 TaskSourceSortKey::TaskSourceSortKey(TaskPriority priority,
-                                     TimeTicks next_task_sequenced_time)
+                                     TimeTicks ready_time,
+                                     uint8_t worker_count)
     : priority_(priority),
-      next_task_sequenced_time_(next_task_sequenced_time) {}
+      worker_count_(worker_count),
+      ready_time_(ready_time) {}
 
 bool TaskSourceSortKey::operator<=(const TaskSourceSortKey& other) const {
   // This TaskSourceSortKey is considered more important than |other| if it has
-  // a higher priority or if it has the same priority but its next task was
+  // a higher priority or if it has the same priority but fewer workers, or if
+  // it has the same priority and same worker count but its next task was
   // posted sooner than |other|'s.
   const int priority_diff =
       static_cast<int>(priority_) - static_cast<int>(other.priority_);
@@ -22,7 +28,11 @@ bool TaskSourceSortKey::operator<=(const TaskSourceSortKey& other) const {
     return true;
   if (priority_diff < 0)
     return false;
-  return next_task_sequenced_time_ <= other.next_task_sequenced_time_;
+  if (worker_count_ < other.worker_count_)
+    return true;
+  if (worker_count_ > other.worker_count_)
+    return false;
+  return ready_time_ <= other.ready_time_;
 }
 
 }  // namespace internal

base/task/thread_pool/task_source_sort_key.h

@@ -16,19 +16,21 @@ namespace internal {
 class BASE_EXPORT TaskSourceSortKey final {
  public:
   TaskSourceSortKey() = default;
-  TaskSourceSortKey(TaskPriority priority, TimeTicks next_task_sequenced_time);
+  TaskSourceSortKey(TaskPriority priority,
+                    TimeTicks ready_time,
+                    uint8_t worker_count = 0);
 
   TaskPriority priority() const { return priority_; }
-  TimeTicks next_task_sequenced_time() const {
-    return next_task_sequenced_time_;
-  }
+  uint8_t worker_count() const { return worker_count_; }
+  TimeTicks ready_time() const { return ready_time_; }
 
   // Lower sort key means more important.
   bool operator<=(const TaskSourceSortKey& other) const;
 
   bool operator==(const TaskSourceSortKey& other) const {
     return priority_ == other.priority_ &&
-           next_task_sequenced_time_ == other.next_task_sequenced_time_;
+           worker_count_ == other.worker_count_ &&
+           ready_time_ == other.ready_time_;
   }
   bool operator!=(const TaskSourceSortKey& other) const {
     return !(other == *this);
@@ -42,9 +44,13 @@ class BASE_EXPORT TaskSourceSortKey final {
   // created.
   TaskPriority priority_;
 
-  // Sequenced time of the next task to run in the sequence at the time this
-  // sort key was created.
-  TimeTicks next_task_sequenced_time_;
+  // Number of workers running the task source, used as secondary sort key
+  // prioritizing task sources with fewer workers.
+  uint8_t worker_count_;
+
+  // Time since the task source has been ready to run upcoming work, used as
+  // secondary sort key after |worker_count| prioritizing older task sources.
+  TimeTicks ready_time_;
 };
 
 }  // namespace internal

base/task/thread_pool/task_source_sort_key_unittest.cc

@@ -20,9 +20,13 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
                           TimeTicks::FromInternalValue(1000));
   TaskSourceSortKey key_d(TaskPriority::USER_VISIBLE,
                           TimeTicks::FromInternalValue(2000));
-  TaskSourceSortKey key_e(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_e(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(1000), 1);
+  TaskSourceSortKey key_f(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(2000), 1);
+  TaskSourceSortKey key_g(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(1000));
-  TaskSourceSortKey key_f(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_h(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(2000));
 
   EXPECT_LE(key_a, key_a);
@@ -31,6 +35,8 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_FALSE(key_d <= key_a);
   EXPECT_FALSE(key_e <= key_a);
   EXPECT_FALSE(key_f <= key_a);
+  EXPECT_FALSE(key_g <= key_a);
+  EXPECT_FALSE(key_h <= key_a);
 
   EXPECT_LE(key_a, key_b);
   EXPECT_LE(key_b, key_b);
@@ -38,6 +44,8 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_FALSE(key_d <= key_b);
   EXPECT_FALSE(key_e <= key_b);
   EXPECT_FALSE(key_f <= key_b);
+  EXPECT_FALSE(key_g <= key_b);
+  EXPECT_FALSE(key_h <= key_b);
 
   EXPECT_LE(key_a, key_c);
   EXPECT_LE(key_b, key_c);
@@ -45,6 +53,8 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_FALSE(key_d <= key_c);
   EXPECT_FALSE(key_e <= key_c);
   EXPECT_FALSE(key_f <= key_c);
+  EXPECT_FALSE(key_g <= key_c);
+  EXPECT_FALSE(key_h <= key_c);
 
   EXPECT_LE(key_a, key_d);
   EXPECT_LE(key_b, key_d);
@@ -52,6 +62,8 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_LE(key_d, key_d);
   EXPECT_FALSE(key_e <= key_d);
   EXPECT_FALSE(key_f <= key_d);
+  EXPECT_FALSE(key_g <= key_d);
+  EXPECT_FALSE(key_h <= key_d);
 
   EXPECT_LE(key_a, key_e);
   EXPECT_LE(key_b, key_e);
@@ -59,6 +71,8 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_LE(key_d, key_e);
   EXPECT_LE(key_e, key_e);
   EXPECT_FALSE(key_f <= key_e);
+  EXPECT_FALSE(key_g <= key_e);
+  EXPECT_FALSE(key_h <= key_e);
 
   EXPECT_LE(key_a, key_f);
   EXPECT_LE(key_b, key_f);
@@ -66,6 +80,26 @@ TEST(TaskSourceSortKeyTest, OperatorLessThanOrEqual) {
   EXPECT_LE(key_d, key_f);
   EXPECT_LE(key_e, key_f);
   EXPECT_LE(key_f, key_f);
+  EXPECT_FALSE(key_g <= key_f);
+  EXPECT_FALSE(key_h <= key_f);
+
+  EXPECT_LE(key_a, key_g);
+  EXPECT_LE(key_b, key_g);
+  EXPECT_LE(key_c, key_g);
+  EXPECT_LE(key_d, key_g);
+  EXPECT_LE(key_e, key_g);
+  EXPECT_LE(key_f, key_g);
+  EXPECT_LE(key_g, key_g);
+  EXPECT_FALSE(key_h <= key_g);
+
+  EXPECT_LE(key_a, key_h);
+  EXPECT_LE(key_b, key_h);
+  EXPECT_LE(key_c, key_h);
+  EXPECT_LE(key_d, key_h);
+  EXPECT_LE(key_e, key_h);
+  EXPECT_LE(key_f, key_h);
+  EXPECT_LE(key_g, key_h);
+  EXPECT_LE(key_h, key_h);
 }
 
 TEST(TaskSourceSortKeyTest, OperatorEqual) {
@@ -77,9 +111,13 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
                           TimeTicks::FromInternalValue(1000));
   TaskSourceSortKey key_d(TaskPriority::USER_VISIBLE,
                           TimeTicks::FromInternalValue(2000));
-  TaskSourceSortKey key_e(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_e(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(1000), 1);
+  TaskSourceSortKey key_f(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(2000), 1);
+  TaskSourceSortKey key_g(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(1000));
-  TaskSourceSortKey key_f(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_h(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(2000));
 
   EXPECT_EQ(key_a, key_a);
@@ -88,6 +126,8 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_FALSE(key_d == key_a);
   EXPECT_FALSE(key_e == key_a);
   EXPECT_FALSE(key_f == key_a);
+  EXPECT_FALSE(key_g == key_a);
+  EXPECT_FALSE(key_h == key_a);
 
   EXPECT_FALSE(key_a == key_b);
   EXPECT_EQ(key_b, key_b);
@@ -95,6 +135,8 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_FALSE(key_d == key_b);
   EXPECT_FALSE(key_e == key_b);
   EXPECT_FALSE(key_f == key_b);
+  EXPECT_FALSE(key_g == key_b);
+  EXPECT_FALSE(key_h == key_b);
 
   EXPECT_FALSE(key_a == key_c);
   EXPECT_FALSE(key_b == key_c);
@@ -102,6 +144,8 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_FALSE(key_d == key_c);
   EXPECT_FALSE(key_e == key_c);
   EXPECT_FALSE(key_f == key_c);
+  EXPECT_FALSE(key_g == key_c);
+  EXPECT_FALSE(key_h == key_c);
 
   EXPECT_FALSE(key_a == key_d);
   EXPECT_FALSE(key_b == key_d);
@@ -109,6 +153,8 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_EQ(key_d, key_d);
   EXPECT_FALSE(key_e == key_d);
   EXPECT_FALSE(key_f == key_d);
+  EXPECT_FALSE(key_g == key_d);
+  EXPECT_FALSE(key_h == key_d);
 
   EXPECT_FALSE(key_a == key_e);
   EXPECT_FALSE(key_b == key_e);
@@ -116,6 +162,8 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_FALSE(key_d == key_e);
   EXPECT_EQ(key_e, key_e);
   EXPECT_FALSE(key_f == key_e);
+  EXPECT_FALSE(key_g == key_e);
+  EXPECT_FALSE(key_h == key_e);
 
   EXPECT_FALSE(key_a == key_f);
   EXPECT_FALSE(key_b == key_f);
@@ -123,6 +171,26 @@ TEST(TaskSourceSortKeyTest, OperatorEqual) {
   EXPECT_FALSE(key_d == key_f);
   EXPECT_FALSE(key_e == key_f);
   EXPECT_EQ(key_f, key_f);
+  EXPECT_FALSE(key_g == key_f);
+  EXPECT_FALSE(key_h == key_f);
+
+  EXPECT_FALSE(key_a == key_g);
+  EXPECT_FALSE(key_b == key_g);
+  EXPECT_FALSE(key_c == key_g);
+  EXPECT_FALSE(key_d == key_g);
+  EXPECT_FALSE(key_e == key_g);
+  EXPECT_FALSE(key_f == key_g);
+  EXPECT_EQ(key_g, key_g);
+  EXPECT_FALSE(key_h == key_g);
+
+  EXPECT_FALSE(key_a == key_h);
+  EXPECT_FALSE(key_b == key_h);
+  EXPECT_FALSE(key_c == key_h);
+  EXPECT_FALSE(key_d == key_h);
+  EXPECT_FALSE(key_e == key_h);
+  EXPECT_FALSE(key_f == key_h);
+  EXPECT_FALSE(key_g == key_h);
+  EXPECT_EQ(key_h, key_h);
 }
 
 TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
@@ -134,9 +202,13 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
                           TimeTicks::FromInternalValue(1000));
   TaskSourceSortKey key_d(TaskPriority::USER_VISIBLE,
                           TimeTicks::FromInternalValue(2000));
-  TaskSourceSortKey key_e(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_e(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(1000), 1);
+  TaskSourceSortKey key_f(TaskPriority::USER_VISIBLE,
+                          TimeTicks::FromInternalValue(2000), 1);
+  TaskSourceSortKey key_g(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(1000));
-  TaskSourceSortKey key_f(TaskPriority::BEST_EFFORT,
+  TaskSourceSortKey key_h(TaskPriority::BEST_EFFORT,
                           TimeTicks::FromInternalValue(2000));
 
   EXPECT_FALSE(key_a != key_a);
@@ -145,6 +217,8 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_NE(key_d, key_a);
   EXPECT_NE(key_e, key_a);
   EXPECT_NE(key_f, key_a);
+  EXPECT_NE(key_g, key_a);
+  EXPECT_NE(key_h, key_a);
 
   EXPECT_NE(key_a, key_b);
   EXPECT_FALSE(key_b != key_b);
@@ -152,6 +226,8 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_NE(key_d, key_b);
   EXPECT_NE(key_e, key_b);
   EXPECT_NE(key_f, key_b);
+  EXPECT_NE(key_g, key_b);
+  EXPECT_NE(key_h, key_b);
 
   EXPECT_NE(key_a, key_c);
   EXPECT_NE(key_b, key_c);
@@ -159,6 +235,8 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_NE(key_d, key_c);
   EXPECT_NE(key_e, key_c);
   EXPECT_NE(key_f, key_c);
+  EXPECT_NE(key_g, key_c);
+  EXPECT_NE(key_h, key_c);
 
   EXPECT_NE(key_a, key_d);
   EXPECT_NE(key_b, key_d);
@@ -166,6 +244,8 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_FALSE(key_d != key_d);
   EXPECT_NE(key_e, key_d);
   EXPECT_NE(key_f, key_d);
+  EXPECT_NE(key_g, key_d);
+  EXPECT_NE(key_h, key_d);
 
   EXPECT_NE(key_a, key_e);
   EXPECT_NE(key_b, key_e);
@@ -173,6 +253,8 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_NE(key_d, key_e);
   EXPECT_FALSE(key_e != key_e);
   EXPECT_NE(key_f, key_e);
+  EXPECT_NE(key_g, key_e);
+  EXPECT_NE(key_h, key_e);
 
   EXPECT_NE(key_a, key_f);
   EXPECT_NE(key_b, key_f);
@@ -180,6 +262,26 @@ TEST(TaskSourceSortKeyTest, OperatorNotEqual) {
   EXPECT_NE(key_d, key_f);
   EXPECT_NE(key_e, key_f);
   EXPECT_FALSE(key_f != key_f);
+  EXPECT_NE(key_g, key_f);
+  EXPECT_NE(key_h, key_f);
+
+  EXPECT_NE(key_a, key_g);
+  EXPECT_NE(key_b, key_g);
+  EXPECT_NE(key_c, key_g);
+  EXPECT_NE(key_d, key_g);
+  EXPECT_NE(key_e, key_g);
+  EXPECT_NE(key_f, key_g);
+  EXPECT_FALSE(key_g != key_g);
+  EXPECT_NE(key_h, key_g);
+
+  EXPECT_NE(key_a, key_h);
+  EXPECT_NE(key_b, key_h);
+  EXPECT_NE(key_c, key_h);
+  EXPECT_NE(key_d, key_h);
+  EXPECT_NE(key_e, key_h);
+  EXPECT_NE(key_f, key_h);
+  EXPECT_NE(key_g, key_h);
+  EXPECT_FALSE(key_h != key_h);
 }
 
 }  // namespace internal

base/task/thread_pool/thread_group.cc

@@ -194,7 +194,7 @@ RegisteredTaskSource ThreadGroup::TakeRegisteredTaskSource(
   // If the TaskSource isn't saturated, check whether TaskTracker allows it to
   // remain in the PriorityQueue.
   // The canonical way of doing this is to pop the task source to return, call
-  // WillQueueTaskSource() to get an additional RegisteredTaskSource, and
+  // RegisterTaskSource() to get an additional RegisteredTaskSource, and
   // reenqueue that task source if valid. Instead, it is cheaper and equivalent
   // to peek the task source, call RegisterTaskSource() to get an additional
   // RegisteredTaskSource to replace if valid, and only pop |priority_queue_|
@@ -203,14 +203,17 @@ RegisteredTaskSource ThreadGroup::TakeRegisteredTaskSource(
       task_tracker_->RegisterTaskSource(priority_queue_.PeekTaskSource().get());
   if (!task_source)
     return priority_queue_.PopTaskSource();
-  return std::exchange(priority_queue_.PeekTaskSource(),
-                       std::move(task_source));
+  // Replace the top task_source and then update the queue.
+  std::swap(priority_queue_.PeekTaskSource(), task_source);
+  priority_queue_.UpdateSortKey(*task_source.get(), task_source->GetSortKey());
+  return task_source;
 }
 
 void ThreadGroup::UpdateSortKeyImpl(BaseScopedCommandsExecutor* executor,
                                     TaskSource::Transaction transaction) {
   CheckedAutoLock auto_lock(lock_);
-  priority_queue_.UpdateSortKey(std::move(transaction));
+  priority_queue_.UpdateSortKey(*transaction.task_source(),
+                                transaction.task_source()->GetSortKey());
   EnsureEnoughWorkersLockRequired(executor);
 }