[ThreadPool]: Implement standalone JobTaskSource.

This CL Creates a new derived JobTaskSource that tracks the number
of concurrent worker.
It's a precursor of https://chromium-review.googlesource.com/c/chromium/src/+/1582427
which integrates JobTaskSource in ThreadGroups.

Bug: 839091
Change-Id: Ie407823bdd6d4c84fd8b3db7d6d6a58a33846045
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1693182
Commit-Queue: François Doray <fdoray@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Reviewed-by: François Doray <fdoray@chromium.org>
Cr-Commit-Position: refs/heads/master@{#677058}

base/BUILD.gn

@@ -829,6 +829,8 @@ jumbo_component("base") {
     "task/thread_pool/environment_config.h",
     "task/thread_pool/initialization_util.cc",
     "task/thread_pool/initialization_util.h",
+    "task/thread_pool/job_task_source.cc",
+    "task/thread_pool/job_task_source.h",
     "task/thread_pool/pooled_parallel_task_runner.cc",
     "task/thread_pool/pooled_parallel_task_runner.h",
     "task/thread_pool/pooled_sequenced_task_runner.cc",
@@ -2686,6 +2688,7 @@ test("base_unittests") {
     "task/thread_pool/can_run_policy_test.h",
     "task/thread_pool/delayed_task_manager_unittest.cc",
     "task/thread_pool/environment_config_unittest.cc",
+    "task/thread_pool/job_task_source_unittest.cc",
     "task/thread_pool/pooled_single_thread_task_runner_manager_unittest.cc",
     "task/thread_pool/priority_queue_unittest.cc",
     "task/thread_pool/sequence_sort_key_unittest.cc",

base/task/thread_pool/job_task_source.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "base/task/thread_pool/job_task_source.h"
+
+#include <utility>
+
+#include "base/logging.h"
+#include "base/memory/ptr_util.h"
+#include "base/task/task_features.h"
+#include "base/task/thread_pool/thread_pool_clock.h"
+#include "base/time/time.h"
+
+namespace base {
+namespace internal {
+
+JobTaskSource::JobTaskSource(const Location& from_here,
+                             base::RepeatingClosure worker_task,
+                             const TaskTraits& traits)
+    : TaskSource(traits, nullptr, TaskSourceExecutionMode::kJob),
+      from_here_(from_here),
+      worker_task_(std::move(worker_task)),
+      queue_time_(ThreadPoolClock::Now()) {}
+
+JobTaskSource::~JobTaskSource() = default;
+
+ExecutionEnvironment JobTaskSource::GetExecutionEnvironment() {
+  return {SequenceToken::Create(), nullptr};
+}
+
+TaskSource::RunIntent JobTaskSource::WillRunTask() {
+  const size_t max_concurrency = GetMaxConcurrency();
+  const size_t worker_count_initial =
+      worker_count_.load(std::memory_order_relaxed);
+  // Don't allow this worker to run the task if either:
+  //   A) |worker_count_| is already at |max_concurrency|.
+  //   B) |max_concurrency| was lowered below or to |worker_count_|.
+  if (worker_count_initial >= max_concurrency) {
+    // The caller receives an invalid RunIntent and should skip this TaskSource.
+    return RunIntent();
+  }
+  const size_t worker_count_before_add =
+      worker_count_.fetch_add(1, std::memory_order_relaxed);
+  // WillRunTask() has external synchronization to prevent concurrent calls and
+  // it is the only place where |worker_count_| is incremented. Therefore, the
+  // second reading of |worker_count_| from WillRunTask() cannot be greater than
+  // the first reading. However, since DidProcessTask() can decrement
+  // |worker_count_| concurrently with WillRunTask(), the second reading can be
+  // lower than the first reading.
+  DCHECK_LE(worker_count_before_add, worker_count_initial);
+  DCHECK_LT(worker_count_before_add, max_concurrency);
+  return MakeRunIntent(max_concurrency == worker_count_before_add + 1
+                           ? Saturated::kYes
+                           : Saturated::kNo);
+}
+
+size_t JobTaskSource::GetRemainingConcurrency() const {
+  return GetMaxConcurrency() - worker_count_.load(std::memory_order_relaxed);
+}
+
+Optional<Task> JobTaskSource::TakeTask() {
+  DCHECK_GT(worker_count_.load(std::memory_order_relaxed), 0U);
+  DCHECK(worker_task_);
+  return base::make_optional<Task>(from_here_, worker_task_, TimeDelta());
+}
+
+bool JobTaskSource::DidProcessTask(RunResult run_result) {
+  size_t worker_count_before_sub =
+      worker_count_.fetch_sub(1, std::memory_order_relaxed);
+  DCHECK_GT(worker_count_before_sub, 0U);
+  // Re-enqueue the TaskSource if the task ran and the worker count is below the
+  // max concurrency.
+  const bool must_be_queued =
+      run_result == RunResult::kSkippedAtShutdown
+          ? false
+          : worker_count_before_sub <= GetMaxConcurrency();
+  return must_be_queued;
+}
+
+SequenceSortKey JobTaskSource::GetSortKey() const {
+  return SequenceSortKey(traits_.priority(), queue_time_);
+}
+
+void JobTaskSource::Clear() {
+  worker_task_.Reset();
+}
+
+}  // namespace internal
+}  // namespace base

base/task/thread_pool/job_task_source.h

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef BASE_TASK_THREAD_POOL_JOB_TASK_SOURCE_H_
+#define BASE_TASK_THREAD_POOL_JOB_TASK_SOURCE_H_
+
+#include <stddef.h>
+
+#include <atomic>
+
+#include "base/base_export.h"
+#include "base/macros.h"
+#include "base/optional.h"
+#include "base/task/task_traits.h"
+#include "base/task/thread_pool/sequence_sort_key.h"
+#include "base/task/thread_pool/task.h"
+#include "base/task/thread_pool/task_source.h"
+
+namespace base {
+namespace internal {
+
+// A JobTaskSource generates many Tasks from a single RepeatingClosure.
+//
+// Derived classes control the intended concurrency with GetMaxConcurrency().
+// Increase in concurrency is not supported and should never happen.
+// TODO(etiennep): Support concurrency increase.
+class BASE_EXPORT JobTaskSource : public TaskSource {
+ public:
+  JobTaskSource(const Location& from_here,
+                base::RepeatingClosure task,
+                const TaskTraits& traits);
+
+  // TaskSource:
+  RunIntent WillRunTask() override;
+  ExecutionEnvironment GetExecutionEnvironment() override;
+  size_t GetRemainingConcurrency() const override;
+
+ protected:
+  ~JobTaskSource() override;
+
+  // Returns the maximum number of tasks from this TaskSource that can run
+  // concurrently. The implementation can only return values lower than or equal
+  // to previously returned values.
+  virtual size_t GetMaxConcurrency() const = 0;
+
+ private:
+  // TaskSource:
+  Optional<Task> TakeTask() override;
+  bool DidProcessTask(RunResult run_result) override;
+  SequenceSortKey GetSortKey() const override;
+  void Clear() override;
+
+  // The current number of workers concurrently running tasks from this
+  // TaskSource. "memory_order_relaxed" is sufficient to access this variable as
+  // no other state is synchronized with it.
+  std::atomic_size_t worker_count_{0U};
+
+  const Location from_here_;
+  base::RepeatingClosure worker_task_;
+  const TimeTicks queue_time_;
+
+  DISALLOW_COPY_AND_ASSIGN(JobTaskSource);
+};
+
+}  // namespace internal
+}  // namespace base
+
+#endif  // BASE_TASK_THREAD_POOL_JOB_TASK_SOURCE_H_

base/task/thread_pool/job_task_source_unittest.cc

+// Copyright 2019 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "base/task/thread_pool/job_task_source.h"
+
+#include <utility>
+
+#include "base/bind_helpers.h"
+#include "base/memory/ptr_util.h"
+#include "base/task/thread_pool/test_utils.h"
+#include "base/test/gtest_util.h"
+#include "testing/gtest/include/gtest/gtest.h"
+
+namespace base {
+namespace internal {
+
+// Verifies the normal flow of running 2 tasks in series.
+TEST(ThreadPoolJobTaskSourceTest, RunTasks) {
+  scoped_refptr<test::MockJobTaskSource> task_source =
+      MakeRefCounted<test::MockJobTaskSource>(
+          FROM_HERE, DoNothing(), TaskTraits(TaskPriority::BEST_EFFORT),
+          /* num_tasks_to_run */ 2, /* max_concurrency */ 1);
+
+  TaskSource::Transaction task_source_transaction(
+      task_source->BeginTransaction());
+
+  {
+    auto run_intent = task_source->WillRunTask();
+    EXPECT_TRUE(run_intent);
+    EXPECT_TRUE(run_intent.IsSaturated());
+
+    // An attempt to run an additional task is not allowed until this task
+    // is processed.
+    EXPECT_FALSE(task_source->WillRunTask());
+
+    auto task = task_source_transaction.TakeTask(&run_intent);
+
+    EXPECT_FALSE(task_source->WillRunTask());
+
+    std::move(task->task).Run();
+    EXPECT_TRUE(task_source_transaction.DidProcessTask(std::move(run_intent)));
+  }
+  {
+    auto run_intent = task_source->WillRunTask();
+    EXPECT_TRUE(run_intent);
+    EXPECT_TRUE(run_intent.IsSaturated());
+    auto task = task_source_transaction.TakeTask(&run_intent);
+    std::move(task->task).Run();
+    EXPECT_FALSE(task_source_transaction.DidProcessTask(std::move(run_intent)));
+  }
+}
+
+// Verifies that a job task source doesn't get reenqueued when a task is not
+// run.
+TEST(ThreadPoolJobTaskSourceTest, SkipTask) {
+  scoped_refptr<test::MockJobTaskSource> task_source =
+      MakeRefCounted<test::MockJobTaskSource>(
+          FROM_HERE, DoNothing(), TaskTraits(TaskPriority::BEST_EFFORT),
+          /* num_tasks_to_run */ 2, /* max_concurrency */ 1);
+
+  TaskSource::Transaction task_source_transaction(
+      task_source->BeginTransaction());
+
+  auto run_intent = task_source->WillRunTask();
+  EXPECT_TRUE(run_intent);
+  EXPECT_TRUE(run_intent.IsSaturated());
+  auto task = task_source_transaction.TakeTask(&run_intent);
+  EXPECT_FALSE(task_source_transaction.DidProcessTask(
+      std::move(run_intent), TaskSource::RunResult::kSkippedAtShutdown));
+}
+
+// Verifies that multiple tasks can run in parallel up to |max_concurrency|.
+TEST(ThreadPoolJobTaskSourceTest, RunTasksInParallel) {
+  scoped_refptr<test::MockJobTaskSource> task_source =
+      MakeRefCounted<test::MockJobTaskSource>(
+          FROM_HERE, DoNothing(), TaskTraits(TaskPriority::BEST_EFFORT),
+          /* num_tasks_to_run */ 3, /* max_concurrency */ 2);
+
+  TaskSource::Transaction task_source_transaction(
+      task_source->BeginTransaction());
+
+  auto run_intent_a = task_source->WillRunTask();
+  EXPECT_TRUE(run_intent_a);
+  EXPECT_FALSE(run_intent_a.IsSaturated());
+  auto task_a = task_source_transaction.TakeTask(&run_intent_a);
+
+  auto run_intent_b = task_source->WillRunTask();
+  EXPECT_TRUE(run_intent_b);
+  EXPECT_TRUE(run_intent_b.IsSaturated());
+  auto task_b = task_source_transaction.TakeTask(&run_intent_b);
+
+  // WillRunTask() should return a null RunIntent once the max concurrency is
+  // reached.
+  EXPECT_FALSE(task_source->WillRunTask());
+
+  std::move(task_a->task).Run();
+  EXPECT_TRUE(task_source_transaction.DidProcessTask(std::move(run_intent_a)));
+
+  std::move(task_b->task).Run();
+  EXPECT_TRUE(task_source_transaction.DidProcessTask(std::move(run_intent_b)));
+
+  auto run_intent_c = task_source->WillRunTask();
+  EXPECT_TRUE(run_intent_c);
+  EXPECT_TRUE(run_intent_c.IsSaturated());
+  auto task_c = task_source_transaction.TakeTask(&run_intent_c);
+
+  std::move(task_c->task).Run();
+  EXPECT_FALSE(task_source_transaction.DidProcessTask(std::move(run_intent_c)));
+}
+
+TEST(ThreadPoolJobTaskSourceTest, InvalidTakeTask) {
+  scoped_refptr<test::MockJobTaskSource> task_source =
+      MakeRefCounted<test::MockJobTaskSource>(
+          FROM_HERE, DoNothing(), TaskTraits(TaskPriority::BEST_EFFORT),
+          /* num_tasks_to_run */ 1, /* max_concurrency */ 1);
+  TaskSource::Transaction task_source_transaction(
+      task_source->BeginTransaction());
+
+  auto run_intent_a = task_source->WillRunTask();
+  auto run_intent_b = task_source->WillRunTask();
+  EXPECT_FALSE(run_intent_b);
+  // Can not be called with an invalid RunIntent.
+  EXPECT_DCHECK_DEATH(
+      { auto task = task_source_transaction.TakeTask(&run_intent_b); });
+  run_intent_a.ReleaseForTesting();
+}
+
+TEST(ThreadPoolJobTaskSourceTest, InvalidDidProcessTask) {
+  scoped_refptr<test::MockJobTaskSource> task_source =
+      MakeRefCounted<test::MockJobTaskSource>(
+          FROM_HERE, DoNothing(), TaskTraits(TaskPriority::BEST_EFFORT),
+          /* num_tasks_to_run */ 1, /* max_concurrency */ 1);
+  TaskSource::Transaction task_source_transaction(
+      task_source->BeginTransaction());
+
+  auto run_intent = task_source->WillRunTask();
+  EXPECT_TRUE(run_intent);
+  // Can not be called before TakeTask().
+  EXPECT_DCHECK_DEATH(
+      task_source_transaction.DidProcessTask(std::move(run_intent)));
+  run_intent.ReleaseForTesting();
+}
+
+}  // namespace internal
+}  // namespace base

base/task/thread_pool/sequence.cc

@@ -66,10 +66,10 @@ TaskSource::RunIntent Sequence::WillRunTask() {
   // TakeTask() and DidProcessTask() and only called if |!queue_.empty()|, which
   // means it won't race with WillPushTask()/PushTask().
   has_worker_ = true;
-  return MakeRunIntent(ConcurrencyStatus::kSaturated);
+  return MakeRunIntent(Saturated::kYes);
 }
 
-size_t Sequence::GetMaxConcurrency() const {
+size_t Sequence::GetRemainingConcurrency() const {
   return 1;
 }
 
@@ -83,7 +83,7 @@ Optional<Task> Sequence::TakeTask() {
   return std::move(next_task);
 }
 
-bool Sequence::DidProcessTask(bool /* was_run */) {
+bool Sequence::DidProcessTask(RunResult run_result) {
   // There should never be a call to DidProcessTask without an associated
   // WillRunTask().
   DCHECK(has_worker_);
@@ -92,6 +92,9 @@ bool Sequence::DidProcessTask(bool /* was_run */) {
     ReleaseTaskRunner();
     return false;
   }
+  // Let the caller re-enqueue this non-empty Sequence regardless of
+  // |run_result| so it can continue churning through this Sequence's tasks and
+  // skip/delete them in the proper scope.
   return true;
 }
 

base/task/thread_pool/sequence.h

@@ -86,7 +86,7 @@ class BASE_EXPORT Sequence : public TaskSource {
   // TaskSource:
   ExecutionEnvironment GetExecutionEnvironment() override;
   RunIntent WillRunTask() override;
-  size_t GetMaxConcurrency() const override;
+  size_t GetRemainingConcurrency() const override;
 
   // Returns a token that uniquely identifies this Sequence.
   const SequenceToken& token() const { return token_; }
@@ -100,7 +100,7 @@ class BASE_EXPORT Sequence : public TaskSource {
 
   // TaskSource:
   Optional<Task> TakeTask() override WARN_UNUSED_RESULT;
-  bool DidProcessTask(bool can_keep_running) override;
+  bool DidProcessTask(RunResult run_result) override;
   SequenceSortKey GetSortKey() const override;
   void Clear() override;
 

base/task/thread_pool/sequence_unittest.cc

@@ -64,45 +64,52 @@ TEST(ThreadPoolSequenceTest, PushTakeRemove) {
   sequence_transaction.PushTask(CreateTask(&mock_task_d));
 
   // Take the task in front of the sequence. It should be task A.
-  Optional<Task> task = sequence_transaction.TakeTask(sequence->WillRunTask());
+  auto run_intent = sequence->WillRunTask();
+  Optional<Task> task = sequence_transaction.TakeTask(&run_intent);
   ExpectMockTask(&mock_task_a, &task.value());
   EXPECT_FALSE(task->queue_time.is_null());
 
   // Remove the empty slot. Task B should now be in front.
-  EXPECT_TRUE(sequence_transaction.DidProcessTask(/* was_run */ true));
+  EXPECT_TRUE(sequence_transaction.DidProcessTask(std::move(run_intent)));
+
   EXPECT_FALSE(sequence_transaction.WillPushTask());
-  task = sequence_transaction.TakeTask(sequence->WillRunTask());
+  run_intent = sequence->WillRunTask();
+  task = sequence_transaction.TakeTask(&run_intent);
   ExpectMockTask(&mock_task_b, &task.value());
   EXPECT_FALSE(task->queue_time.is_null());
 
   // Remove the empty slot. Task C should now be in front.
-  EXPECT_TRUE(sequence_transaction.DidProcessTask(/* was_run */ true));
+  EXPECT_TRUE(sequence_transaction.DidProcessTask(std::move(run_intent)));
+
   EXPECT_FALSE(sequence_transaction.WillPushTask());
-  task = sequence_transaction.TakeTask(sequence->WillRunTask());
+  run_intent = sequence->WillRunTask();
+  task = sequence_transaction.TakeTask(&run_intent);
   ExpectMockTask(&mock_task_c, &task.value());
   EXPECT_FALSE(task->queue_time.is_null());
 
   // Remove the empty slot.
-  EXPECT_TRUE(sequence_transaction.DidProcessTask(/* was_run */ true));
+  EXPECT_TRUE(sequence_transaction.DidProcessTask(std::move(run_intent)));
 
   // Push task E in the sequence.
   EXPECT_FALSE(sequence_transaction.WillPushTask());
   sequence_transaction.PushTask(CreateTask(&mock_task_e));
 
   // Task D should be in front.
-  task = sequence_transaction.TakeTask(sequence->WillRunTask());
+  run_intent = sequence->WillRunTask();
+  task = sequence_transaction.TakeTask(&run_intent);
   ExpectMockTask(&mock_task_d, &task.value());
   EXPECT_FALSE(task->queue_time.is_null());
 
   // Remove the empty slot. Task E should now be in front.
-  EXPECT_TRUE(sequence_transaction.DidProcessTask(/* was_run */ true));
+  EXPECT_TRUE(sequence_transaction.DidProcessTask(std::move(run_intent)));
   EXPECT_FALSE(sequence_transaction.WillPushTask());
-  task = sequence_transaction.TakeTask(sequence->WillRunTask());
+  run_intent = sequence->WillRunTask();
+  task = sequence_transaction.TakeTask(&run_intent);
   ExpectMockTask(&mock_task_e, &task.value());
   EXPECT_FALSE(task->queue_time.is_null());
 
   // Remove the empty slot. The sequence should now be empty.
-  EXPECT_FALSE(sequence_transaction.DidProcessTask(/* was_run */ true));
+  EXPECT_FALSE(sequence_transaction.DidProcessTask(std::move(run_intent)));
   EXPECT_TRUE(sequence_transaction.WillPushTask());
 }
 
@@ -123,8 +130,9 @@ TEST(ThreadPoolSequenceTest, GetSortKeyBestEffort) {
 
   // Take the task from the sequence, so that its sequenced time is available
   // for the check below.
-  auto take_best_effort_task = best_effort_sequence_transaction.TakeTask(
-      best_effort_sequence->WillRunTask());
+  auto run_intent = best_effort_sequence->WillRunTask();
+  auto take_best_effort_task =
+      best_effort_sequence_transaction.TakeTask(&run_intent);
 
   // Verify the sort key.
   EXPECT_EQ(TaskPriority::BEST_EFFORT, best_effort_sort_key.priority());
@@ -132,7 +140,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyBestEffort) {
             best_effort_sort_key.next_task_sequenced_time());
 
   // DidProcessTask for correctness.
-  best_effort_sequence_transaction.DidProcessTask(/* was_run */ true);
+  best_effort_sequence_transaction.DidProcessTask(std::move(run_intent));
 }
 
 // Same as ThreadPoolSequenceTest.GetSortKeyBestEffort, but with a
@@ -153,8 +161,9 @@ TEST(ThreadPoolSequenceTest, GetSortKeyForeground) {
 
   // Take the task from the sequence, so that its sequenced time is available
   // for the check below.
-  auto take_foreground_task = foreground_sequence_transaction.TakeTask(
-      foreground_sequence->WillRunTask());
+  auto run_intent = foreground_sequence->WillRunTask();
+  auto take_foreground_task =
+      foreground_sequence_transaction.TakeTask(&run_intent);
 
   // Verify the sort key.
   EXPECT_EQ(TaskPriority::USER_VISIBLE, foreground_sort_key.priority());
@@ -162,7 +171,7 @@ TEST(ThreadPoolSequenceTest, GetSortKeyForeground) {
             foreground_sort_key.next_task_sequenced_time());
 
   // DidProcessTask for correctness.
-  foreground_sequence_transaction.DidProcessTask(/* was_run */ true);
+  foreground_sequence_transaction.DidProcessTask(std::move(run_intent));
 }
 
 // Verify that a DCHECK fires if DidProcessTask() is called on a sequence which
@@ -173,8 +182,10 @@ TEST(ThreadPoolSequenceTest, DidProcessTaskWithoutTakeTask) {
   Sequence::Transaction sequence_transaction(sequence->BeginTransaction());
   sequence_transaction.PushTask(Task(FROM_HERE, DoNothing(), TimeDelta()));
 
-  EXPECT_DCHECK_DEATH(
-      { sequence_transaction.DidProcessTask(/* was_run */ true); });
+  EXPECT_DCHECK_DEATH({
+    auto run_intent = sequence->WillRunTask();
+    sequence_transaction.DidProcessTask(std::move(run_intent));
+  });
 }
 
 // Verify that a DCHECK fires if TakeTask() is called on a sequence whose front
@@ -185,9 +196,15 @@ TEST(ThreadPoolSequenceTest, TakeEmptyFrontSlot) {
   Sequence::Transaction sequence_transaction(sequence->BeginTransaction());
   sequence_transaction.PushTask(Task(FROM_HERE, DoNothing(), TimeDelta()));
 
-  EXPECT_TRUE(sequence_transaction.TakeTask(sequence->WillRunTask()));
-  EXPECT_DCHECK_DEATH(
-      { sequence_transaction.TakeTask(sequence->WillRunTask()); });
+  {
+    auto run_intent = sequence->WillRunTask();
+    EXPECT_TRUE(sequence_transaction.TakeTask(&run_intent));
+    run_intent.ReleaseForTesting();
+  }
+  EXPECT_DCHECK_DEATH({
+    auto run_intent = sequence->WillRunTask();
+    auto task = sequence_transaction.TakeTask(&run_intent);
+  });
 }
 
 // Verify that a DCHECK fires if TakeTask() is called on an empty sequence.
@@ -195,8 +212,10 @@ TEST(ThreadPoolSequenceTest, TakeEmptySequence) {
   scoped_refptr<Sequence> sequence = MakeRefCounted<Sequence>(
       TaskTraits(ThreadPool()), nullptr, TaskSourceExecutionMode::kParallel);
   Sequence::Transaction sequence_transaction(sequence->BeginTransaction());
+  auto run_intent = sequence->WillRunTask();
   EXPECT_DCHECK_DEATH(
-      { sequence_transaction.TakeTask(sequence->WillRunTask()); });
+      { auto task = sequence_transaction.TakeTask(&run_intent); });
+  run_intent.ReleaseForTesting();
 }
 
 }  // namespace internal

base/task/thread_pool/task_source.cc

@@ -17,7 +17,8 @@ namespace internal {
 
 TaskSource::RunIntent::RunIntent(RunIntent&& other) noexcept
     : task_source_(other.task_source_),
-      concurrency_status_(other.concurrency_status_) {
+      run_step_(other.run_step_),
+      is_saturated_(other.is_saturated_) {
   other.task_source_ = nullptr;
 }
 
@@ -29,13 +30,14 @@ TaskSource::RunIntent& TaskSource::RunIntent::operator=(RunIntent&& other) {
   DCHECK_EQ(task_source_, nullptr);
   task_source_ = other.task_source_;
   other.task_source_ = nullptr;
-  concurrency_status_ = other.concurrency_status_;
+  run_step_ = other.run_step_;
+  is_saturated_ = other.is_saturated_;
   return *this;
 }
 
 TaskSource::RunIntent::RunIntent(const TaskSource* task_source,
-                                 ConcurrencyStatus concurrency_status)
-    : task_source_(task_source), concurrency_status_(concurrency_status) {}
+                                 Saturated is_saturated)
+    : task_source_(task_source), is_saturated_(is_saturated) {}
 
 TaskSource::Transaction::Transaction(TaskSource* task_source)
     : task_source_(task_source) {
@@ -54,14 +56,20 @@ TaskSource::Transaction::~Transaction() {
   }
 }
 
-Optional<Task> TaskSource::Transaction::TakeTask(RunIntent intent) {
-  DCHECK_EQ(intent.task_source_, task_source());
-  intent.Release();
+Optional<Task> TaskSource::Transaction::TakeTask(RunIntent* intent) {
+  DCHECK_EQ(intent->task_source_, task_source());
+  DCHECK_EQ(intent->run_step_, RunIntent::State::kInitial);
+  intent->run_step_ = RunIntent::State::kTaskAcquired;
   return task_source_->TakeTask();
 }
 
-bool TaskSource::Transaction::DidProcessTask(bool was_run) {
-  return task_source_->DidProcessTask(was_run);
+bool TaskSource::Transaction::DidProcessTask(RunIntent intent,
+                                             RunResult run_result) {
+  DCHECK_EQ(intent.task_source_, task_source());
+  DCHECK_EQ(intent.run_step_, RunIntent::State::kTaskAcquired);
+  intent.run_step_ = RunIntent::State::kCompleted;
+  intent.Release();
+  return task_source_->DidProcessTask(run_result);
 }
 
 SequenceSortKey TaskSource::Transaction::GetSortKey() const {
@@ -78,9 +86,8 @@ void TaskSource::Transaction::UpdatePriority(TaskPriority priority) {
   task_source_->traits_.UpdatePriority(priority);
 }
 
-TaskSource::RunIntent TaskSource::MakeRunIntent(
-    ConcurrencyStatus concurrency_status) const {
-  return RunIntent(this, concurrency_status);
+TaskSource::RunIntent TaskSource::MakeRunIntent(Saturated is_saturated) const {
+  return RunIntent(this, is_saturated);
 }
 
 void TaskSource::SetHeapHandle(const HeapHandle& handle) {
@@ -97,7 +104,9 @@ TaskSource::TaskSource(const TaskTraits& traits,
     : traits_(traits),
       task_runner_(task_runner),
       execution_mode_(execution_mode) {
-  DCHECK(task_runner_ || execution_mode_ == TaskSourceExecutionMode::kParallel);
+  DCHECK(task_runner_ ||
+         execution_mode_ == TaskSourceExecutionMode::kParallel ||
+         execution_mode_ == TaskSourceExecutionMode::kJob);
 }
 
 TaskSource::~TaskSource() = default;

base/task/thread_pool/task_source.h

@@ -29,7 +29,8 @@ enum class TaskSourceExecutionMode {
   kParallel,
   kSequenced,
   kSingleThread,
-  kMax = kSingleThread,
+  kJob,
+  kMax = kJob,
 };
 
 struct BASE_EXPORT ExecutionEnvironment {
@@ -41,32 +42,39 @@ struct BASE_EXPORT ExecutionEnvironment {
 // executed.
 //
 // In order to execute a task from this TaskSource, a worker should first make
-// sure that a task can run with WillRunTask(). TakeTask() can then be called to
-// access the next Task, and DidProcessTask() must be called after the task
-// executed. Many overlapping chains of WillRunTask(), TakeTask(), run and
-// DidProcessTask() can run concurrently, as permitted by WillRunTask(). This
-// ensure that the number of workers concurrently running tasks never go over
-// the intended concurrency.
+// sure that a task can run with WillRunTask() which returns a RunIntent.
+// TakeTask() can then be called to access the next Task, and DidProcessTask()
+// must be called after the task was processed. Many overlapping chains of
+// WillRunTask(), TakeTask(), run and DidProcessTask() can run concurrently, as
+// permitted by WillRunTask(). This ensure that the number of workers
+// concurrently running tasks never go over the intended concurrency.
 //
 // In comments below, an "empty TaskSource" is a TaskSource with no Task.
 //
 // Note: there is a known refcounted-ownership cycle in the Scheduler
-// architecture: TaskSource -> TaskRunner -> TaskSource -> ... This is
-// okay so long as the other owners of TaskSource (PriorityQueue and
-// WorkerThread in alternation and
-// ThreadGroupImpl::WorkerThreadDelegateImpl::GetWork() temporarily) keep
-// running it (and taking Tasks from it as a result). A dangling reference cycle
-// would only occur should they release their reference to it while it's not
-// empty. In other words, it is only correct for them to release it when
-// DidProcessTask() returns false.
+// architecture: TaskSource -> TaskRunner -> TaskSource -> ... This is okay so
+// long as the other owners of TaskSource (PriorityQueue and WorkerThread in
+// alternation and ThreadGroupImpl::WorkerThreadDelegateImpl::GetWork()
+// temporarily) keep running it (and taking Tasks from it as a result). A
+// dangling reference cycle would only occur should they release their reference
+// to it while it's not empty. In other words, it is only correct for them to
+// release it when DidProcessTask() returns false.
 //
 // This class is thread-safe.
 class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
+ protected:
+  // Indicates whether a TaskSource has reached its maximum intended concurrency
+  // and may not run any additional tasks.
+  enum class Saturated {
+    kYes,
+    kNo,
+  };
+
  public:
-  // Indicates whether a TaskSource may run any additional tasks.
-  enum class ConcurrencyStatus {
-    kSaturated,  // The maximum intended concurrency was reached.
-    kPartial,    // Additional tasks may concurrently run.
+  // Indicates if a task was run or skipped as a result of shutdown.
+  enum class RunResult {
+    kDidRun,
+    kSkippedAtShutdown,
   };
 
   // Result of WillRunTask(). A single task associated with a RunIntent may be
@@ -85,25 +93,36 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
     // may not run any additional tasks beyond this RunIntent as it has reached
     // its maximum concurrency. This indicates that the TaskSource no longer
     // needs to be queued.
-    bool IsSaturated() const {
-      return concurrency_status_ == ConcurrencyStatus::kSaturated;
-    }
+    bool IsSaturated() const { return is_saturated_ == Saturated::kYes; }
 
     const TaskSource* task_source() const { return task_source_; }
 
+    void ReleaseForTesting() {
+      DCHECK(task_source_);
+      task_source_ = nullptr;
+    }
+
    private:
     friend class TaskSource;
 
-    RunIntent(const TaskSource* task_source,
-              ConcurrencyStatus concurrency_status);
+    // Indicates the step of a run intent chain.
+    enum class State {
+      kInitial,       // After WillRunTask().
+      kTaskAcquired,  // After TakeTask().
+      kCompleted,     // After DidProcessTask().
+    };
+
+    RunIntent(const TaskSource* task_source, Saturated is_saturated);
 
     void Release() {
+      DCHECK_EQ(run_step_, State::kCompleted);
       DCHECK(task_source_);
       task_source_ = nullptr;
     }
 
     const TaskSource* task_source_ = nullptr;
-    ConcurrencyStatus concurrency_status_ = ConcurrencyStatus::kSaturated;
+    State run_step_ = State::kInitial;
+    Saturated is_saturated_ = Saturated::kYes;
   };
 
   // A Transaction can perform multiple operations atomically on a
@@ -118,19 +137,19 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
     operator bool() const { return !!task_source_; }
 
     // Returns the next task to run from this TaskSource. This should be called
-    // only if NeedsWorker returns true. Cannot be called on an empty
-    // TaskSource.
+    // only with a valid |intent|. Cannot be called on an empty TaskSource.
     //
     // Because this method cannot be called on an empty TaskSource, the returned
     // Optional<Task> is never nullptr. An Optional is used in preparation for
     // the merge between ThreadPool and TaskQueueManager (in Blink).
     // https://crbug.com/783309
-    Optional<Task> TakeTask(RunIntent intent);
+    Optional<Task> TakeTask(RunIntent* intent) WARN_UNUSED_RESULT;
 
-    // Must be called once the task was run or skipped. |was_run| should be true
+    // Must be called once the task was run or skipped. |run_result| indicates
     // if the task executed. Cannot be called on an empty TaskSource. Returns
     // true if the TaskSource should be queued after this operation.
-    bool DidProcessTask(bool was_run);
+    bool DidProcessTask(RunIntent intent,
+                        RunResult run_result = RunResult::kDidRun);
 
     // Returns a SequenceSortKey representing the priority of the TaskSource.
     // Cannot be called on an empty TaskSource.
@@ -180,11 +199,10 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
   // holding the TaskSource).
   virtual RunIntent WillRunTask() = 0;
 
-  // Returns the maximum number of tasks from this TaskSource that can run
-  // concurrently. The concurrency is generally controlled through
-  // WillRunTask(), but calling this directly is useful to determine the right
-  // number of workers beforehand.
-  virtual size_t GetMaxConcurrency() const = 0;
+  // Thread-safe but the returned value may immediately be obsolete. As such
+  // this should only be used as a best-effort guess of how many more workers
+  // are needed.
+  virtual size_t GetRemainingConcurrency() const = 0;
 
   // Support for IntrusiveHeap.
   void SetHeapHandle(const HeapHandle& handle);
@@ -211,9 +229,9 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
   virtual Optional<Task> TakeTask() = 0;
 
   // Informs this TaskSource that a task was processed. |was_run| indicates
-  // whether the task was allowed to run or not. Returns true if the TaskSource
+  // whether the task executed or not. Returns true if the TaskSource
   // should be queued after this operation.
-  virtual bool DidProcessTask(bool was_run) = 0;
+  virtual bool DidProcessTask(RunResult run_result) = 0;
 
   virtual SequenceSortKey GetSortKey() const = 0;
 
@@ -224,7 +242,7 @@ class BASE_EXPORT TaskSource : public RefCountedThreadSafe<TaskSource> {
 
   // Constructs and returns a RunIntent, where |is_saturated| indicates that the
   // TaskSource has reached its maximum concurrency.
-  RunIntent MakeRunIntent(ConcurrencyStatus concurrency_status) const;
+  RunIntent MakeRunIntent(Saturated is_saturated) const;
 
   // The TaskTraits of all Tasks in the TaskSource.
   TaskTraits traits_;

base/task/thread_pool/task_tracker.cc

@@ -35,7 +35,7 @@ namespace internal {
 namespace {
 
 constexpr const char* kExecutionModeString[] = {"parallel", "sequenced",
-                                                "single thread"};
+                                                "single thread", "job"};
 static_assert(
     size(kExecutionModeString) ==
         static_cast<size_t>(TaskSourceExecutionMode::kMax) + 1,
@@ -465,8 +465,7 @@ RegisteredTaskSource TaskTracker::RunAndPopNextTask(
   {
     TaskSource::Transaction task_source_transaction(
         task_source->BeginTransaction());
-    task = task_source_transaction.TakeTask(
-        std::move(run_intent_with_task_source));
+    task = task_source_transaction.TakeTask(&run_intent_with_task_source);
     traits = task_source_transaction.traits();
   }
 
@@ -481,7 +480,10 @@ RegisteredTaskSource TaskTracker::RunAndPopNextTask(
                   can_run_task);
 
     const bool task_source_must_be_queued =
-        task_source->BeginTransaction().DidProcessTask(can_run_task);
+        task_source->BeginTransaction().DidProcessTask(
+            std::move(run_intent_with_task_source),
+            can_run_task ? TaskSource::RunResult::kDidRun
+                         : TaskSource::RunResult::kSkippedAtShutdown);
 
     if (can_run_task) {
       IncrementNumTasksRun();
@@ -588,6 +590,7 @@ void TaskTracker::RunOrSkipTask(Task task,
     Optional<SequencedTaskRunnerHandle> sequenced_task_runner_handle;
     Optional<ThreadTaskRunnerHandle> single_thread_task_runner_handle;
     switch (task_source->execution_mode()) {
+      case TaskSourceExecutionMode::kJob:
       case TaskSourceExecutionMode::kParallel:
         break;
       case TaskSourceExecutionMode::kSequenced:

base/task/thread_pool/test_task_factory.cc

@@ -57,6 +57,7 @@ void TestTaskFactory::RunTaskCallback(size_t task_index,
 
   // Verify TaskRunnerHandles are set as expected in the task's scope.
   switch (execution_mode_) {
+    case TaskSourceExecutionMode::kJob:
     case TaskSourceExecutionMode::kParallel:
       EXPECT_FALSE(ThreadTaskRunnerHandle::IsSet());
       EXPECT_FALSE(SequencedTaskRunnerHandle::IsSet());

base/task/thread_pool/test_utils.cc

@@ -10,6 +10,7 @@
 #include "base/synchronization/condition_variable.h"
 #include "base/task/thread_pool/pooled_parallel_task_runner.h"
 #include "base/task/thread_pool/pooled_sequenced_task_runner.h"
+#include "base/test/bind_test_util.h"
 #include "base/threading/scoped_blocking_call.h"
 #include "base/threading/thread_restrictions.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -178,6 +179,45 @@ void MockPooledTaskRunnerDelegate::SetThreadGroup(ThreadGroup* thread_group) {
   thread_group_ = thread_group;
 }
 
+MockJobTaskSource::~MockJobTaskSource() = default;
+
+MockJobTaskSource::MockJobTaskSource(const Location& from_here,
+                                     base::RepeatingClosure worker_task,
+                                     const TaskTraits& traits,
+                                     size_t num_tasks_to_run,
+                                     size_t max_concurrency)
+    : JobTaskSource(FROM_HERE,
+                    BindLambdaForTesting([this, worker_task]() {
+                      worker_task.Run();
+                      size_t before = remaining_num_tasks_to_run_.fetch_sub(1);
+                      DCHECK_GT(before, 0U);
+                    }),
+                    traits),
+      remaining_num_tasks_to_run_(num_tasks_to_run),
+      max_concurrency_(max_concurrency) {}
+
+MockJobTaskSource::MockJobTaskSource(const Location& from_here,
+                                     base::OnceClosure worker_task,
+                                     const TaskTraits& traits)
+    : JobTaskSource(FROM_HERE,
+                    base::BindRepeating(
+                        [](MockJobTaskSource* self,
+                           base::OnceClosure&& worker_task) mutable {
+                          std::move(worker_task).Run();
+                          size_t before =
+                              self->remaining_num_tasks_to_run_.fetch_sub(1);
+                          DCHECK_EQ(before, 1U);
+                        },
+                        Unretained(this),
+                        base::Passed(std::move(worker_task))),
+                    traits),
+      remaining_num_tasks_to_run_(1),
+      max_concurrency_(1) {}
+
+size_t MockJobTaskSource::GetMaxConcurrency() const {
+  return std::min(remaining_num_tasks_to_run_.load(), max_concurrency_);
+}
+
 RegisteredTaskSource QueueAndRunTaskSource(
     TaskTracker* task_tracker,
     scoped_refptr<TaskSource> task_source) {

base/task/thread_pool/test_utils.h

@@ -5,10 +5,14 @@
 #ifndef BASE_TASK_THREAD_POOL_TEST_UTILS_H_
 #define BASE_TASK_THREAD_POOL_TEST_UTILS_H_
 
+#include <atomic>
+
+#include "base/callback.h"
 #include "base/task/common/checked_lock.h"
 #include "base/task/task_features.h"
 #include "base/task/task_traits.h"
 #include "base/task/thread_pool/delayed_task_manager.h"
+#include "base/task/thread_pool/job_task_source.h"
 #include "base/task/thread_pool/pooled_task_runner_delegate.h"
 #include "base/task/thread_pool/sequence.h"
 #include "base/task/thread_pool/task_tracker.h"
@@ -71,6 +75,33 @@ class MockPooledTaskRunnerDelegate : public PooledTaskRunnerDelegate {
   ThreadGroup* thread_group_ = nullptr;
 };
 
+// A simple JobTaskSource that will give |worker_task| a fixed number of times,
+// possibly in parallel.
+class MockJobTaskSource : public JobTaskSource {
+ public:
+  // Gives |worker_task| to requesting workers |num_tasks_to_run| times.
+  // Allowing at most |max_concurrency| workers to be running |worker_task| in
+  // parallel.
+  MockJobTaskSource(const Location& from_here,
+                    base::RepeatingClosure worker_task,
+                    const TaskTraits& traits,
+                    size_t num_tasks_to_run,
+                    size_t max_concurrency);
+
+  // Gives |worker_task| to a single requesting worker.
+  MockJobTaskSource(const Location& from_here,
+                    base::OnceClosure worker_task,
+                    const TaskTraits& traits);
+
+  size_t GetMaxConcurrency() const override;
+
+ private:
+  ~MockJobTaskSource() override;
+
+  std::atomic_size_t remaining_num_tasks_to_run_;
+  const size_t max_concurrency_;
+};
+
 // An enumeration of possible thread pool types. Used to parametrize relevant
 // thread_pool tests.
 enum class PoolType {

base/task/thread_pool/thread_pool_impl_unittest.cc

@@ -176,6 +176,8 @@ scoped_refptr<TaskRunner> CreateTaskRunnerAndExecutionMode(
       return thread_pool->CreateSingleThreadTaskRunner(
           traits, default_single_thread_task_runner_mode);
     }
+    case TaskSourceExecutionMode::kJob:
+      break;
   }
   ADD_FAILURE() << "Unknown ExecutionMode";
   return nullptr;

base/task/thread_pool/worker_thread_unittest.cc

@@ -217,10 +217,10 @@ class ThreadPoolWorkerTest : public testing::TestWithParam<int> {
         // Verify the number of Tasks in |registered_task_source|.
         auto transaction(registered_task_source->BeginTransaction());
         for (int i = 0; i < outer_->TasksPerSequence() - 1; ++i) {
-          EXPECT_TRUE(
-              transaction.TakeTask(registered_task_source->WillRunTask()));
+          auto run_intent = registered_task_source->WillRunTask();
+          EXPECT_TRUE(transaction.TakeTask(&run_intent));
           EXPECT_EQ(i == outer_->TasksPerSequence() - 2,
-                    !transaction.DidProcessTask(true));
+                    !transaction.DidProcessTask(std::move(run_intent)));
         }
 
         scoped_refptr<TaskSource> task_source =