Add (Sequenced|SingleThread)TaskRunner::GetCurrentBestEffort

Adds accessors to SequencedTaskRunner and SingleThreadTaskRunner to
return a best-effort task runner for the current sequence/thread. This
can be used as a convenience from code in layers like components/ that
can't refer to explicit BrowserThreads.

Currently this only returns a best-effort task runner if called from a
task scheduled by SequenceManager, such as a BrowserThread. If not it
falls back to calling GetCurrentDefault().

Also adds HasCurrentBestEffort() methods that can be used to check
whether GetCurrentBestEffort() will return the current default.

R=gab

Bug: 441949788
Change-Id: Ie165e61fcf4eddc1ddebcb810558a165d7a28d50
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/7121638
Reviewed-by: Stefan Zager <[email protected]>
Auto-Submit: Joe Mason <[email protected]>
Commit-Queue: Joe Mason <[email protected]>
Reviewed-by: Kent Tamura <[email protected]>
Reviewed-by: Gabriel Charette <[email protected]>
Cr-Commit-Position: refs/heads/main@{#1550107}
NOKEYCHECK=True
GitOrigin-RevId: 223b0999980933f46fa5e43f16e2056b3fbce0b3

Author: JoeNotCharlesGoogle

task/sequence_manager/sequence_manager_impl.cc

@@ -1084,6 +1084,25 @@ WeakPtr<SequenceManagerImpl> SequenceManagerImpl::GetWeakPtr() {
   return weak_factory_.GetWeakPtr();
 }
 
+// static
+scoped_refptr<SingleThreadTaskRunner>
+SequenceManagerImpl::GetCurrentBestEffortTaskRunner(
+    PassKey<SingleThreadTaskRunner>) {
+  if (SequenceManagerImpl* current = GetCurrent()) {
+    if (std::optional<TaskQueue::QueuePriority> best_effort_priority =
+            current->GetBestEffortPriority()) {
+      // Return the first queue with the right priority.
+      for (internal::TaskQueueImpl* task_queue :
+           current->main_thread_only().active_queues) {
+        if (task_queue->GetQueuePriority() == *best_effort_priority) {
+          return task_queue->task_runner();
+        }
+      }
+    }
+  }
+  return nullptr;
+}
+
 void SequenceManagerImpl::SetDefaultTaskRunner(
     scoped_refptr<SingleThreadTaskRunner> task_runner) {
   controller_->SetDefaultTaskRunner(task_runner);

task/sequence_manager/sequence_manager_impl.h

@@ -202,6 +202,11 @@ class BASE_EXPORT SequenceManagerImpl
   // How frequently to perform housekeeping tasks (sweeping canceled tasks etc).
   static constexpr TimeDelta kReclaimMemoryInterval = Seconds(30);
 
+  // Allows SingleThreadTaskRunner to find the best-effort task queue for the
+  // current thread. Returns nullptr if there isn't one.
+  static scoped_refptr<SingleThreadTaskRunner> GetCurrentBestEffortTaskRunner(
+      PassKey<SingleThreadTaskRunner>);
+
  protected:
   static std::unique_ptr<ThreadControllerImpl>
   CreateThreadControllerImplForCurrentThread(const TickClock* clock);

task/sequenced_task_runner.cc

@@ -8,6 +8,7 @@
 
 #include "base/functional/bind.h"
 #include "base/task/default_delayed_task_handle_delegate.h"
+#include "base/task/single_thread_task_runner.h"
 #include "base/task/thread_pool/thread_pool_instance.h"
 #include "base/time/time.h"
 
@@ -98,6 +99,23 @@ bool SequencedTaskRunner::HasCurrentDefault() {
   return !!current_default_handle && !!current_default_handle->task_runner_;
 }
 
+// static
+scoped_refptr<SequencedTaskRunner> SequencedTaskRunner::GetCurrentBestEffort() {
+  // Currently only threads that multiplex several task queues have current
+  // best-effort task runners.
+  if (SingleThreadTaskRunner::HasCurrentBestEffort()) {
+    return SingleThreadTaskRunner::GetCurrentBestEffort();
+  }
+  return GetCurrentDefault();
+}
+
+// static
+bool SequencedTaskRunner::HasCurrentBestEffort() {
+  // Currently only threads that multiplex several task queues have current
+  // best-effort task runners.
+  return SingleThreadTaskRunner::HasCurrentBestEffort();
+}
+
 SequencedTaskRunner::CurrentDefaultHandle::CurrentDefaultHandle(
     scoped_refptr<SequencedTaskRunner> task_runner)
     : CurrentDefaultHandle(std::move(task_runner), MayAlreadyExist{}) {

task/sequenced_task_runner.h

@@ -343,6 +343,22 @@ class BASE_EXPORT SequencedTaskRunner : public TaskRunner {
   //    (which includes any thread that runs a MessagePump).
   [[nodiscard]] static bool HasCurrentDefault();
 
+  // Returns a SequencedTaskRunner for the current task. If possible, the
+  // task runner will schedule tasks with BEST_EFFORT TaskPriority. If not, it
+  // returns the same value as GetCurrentDefault(). See the comments on
+  // HasCurrentBestEffort() for more details.
+  [[nodiscard]] static scoped_refptr<SequencedTaskRunner>
+  GetCurrentBestEffort();
+
+  // Returns true if the current task is running on a sequence that multiplexes
+  // multiple task queues (eg. BrowserThread::UI). If so, GetCurrentBestEffort()
+  // will return the task runner for the lowest-priority task queue. Otherwise
+  // it will call GetCurrentDefault(). So if this and GetCurrentDefault() both
+  // return false, it's not safe to call GetCurrentBestEffort().
+  // TODO(crbug.com/441949788): It would also be possible to return true on a
+  // non-multiplexing sequence that only runs BEST_EFFORT tasks. Implement this.
+  [[nodiscard]] static bool HasCurrentBestEffort();
+
   class BASE_EXPORT CurrentDefaultHandle {
    public:
     // Sets the value returned by `SequencedTaskRunner::GetCurrentDefault()` to
@@ -366,7 +382,7 @@ class BASE_EXPORT SequencedTaskRunner : public TaskRunner {
     // SingleThreadTaskRunner::CurrentHandleOverrideForTesting in unit tests to
     // avoid the friend requirement.
     friend class SingleThreadTaskRunner;
-    FRIEND_TEST_ALL_PREFIXES(SequencedTaskRunnerCurrentDefaultHandleTest,
+    FRIEND_TEST_ALL_PREFIXES(SequencedTaskRunnerCurrentDefaultHandleDeathTest,
                              OverrideWithNull);
     FRIEND_TEST_ALL_PREFIXES(SequencedTaskRunnerCurrentDefaultHandleTest,
                              OverrideWithNonNull);

task/sequenced_task_runner_unittest.cc

@@ -4,8 +4,10 @@
 
 #include "base/task/sequenced_task_runner.h"
 
+#include <array>
 #include <utility>
 
+#include "base/barrier_closure.h"
 #include "base/functional/bind.h"
 #include "base/functional/callback.h"
 #include "base/functional/callback_helpers.h"
@@ -16,6 +18,8 @@
 #include "base/run_loop.h"
 #include "base/sequence_checker_impl.h"
 #include "base/task/sequenced_task_runner.h"
+#include "base/task/single_thread_task_runner.h"
+#include "base/task/task_traits.h"
 #include "base/task/thread_pool.h"
 #include "base/test/bind.h"
 #include "base/test/gtest_util.h"
@@ -215,6 +219,9 @@ class SequencedTaskRunnerCurrentDefaultHandleTest : public ::testing::Test {
   test::TaskEnvironment task_environment_;
 };
 
+using SequencedTaskRunnerCurrentDefaultHandleDeathTest =
+    SequencedTaskRunnerCurrentDefaultHandleTest;
+
 }  // namespace
 
 TEST_F(SequencedTaskRunnerCurrentDefaultHandleTest, FromTaskEnvironment) {
@@ -239,7 +246,7 @@ TEST_F(SequencedTaskRunnerCurrentDefaultHandleTest,
 
 // Verify that `CurrentDefaultHandle` can be used to set the current default
 // `SequencedTaskRunner` to null in a scope that already has a default.
-TEST_F(SequencedTaskRunnerCurrentDefaultHandleTest, OverrideWithNull) {
+TEST_F(SequencedTaskRunnerCurrentDefaultHandleDeathTest, OverrideWithNull) {
   EXPECT_TRUE(SequencedTaskRunner::HasCurrentDefault());
   auto tr1 = SequencedTaskRunner::GetCurrentDefault();
   EXPECT_TRUE(tr1);
@@ -292,4 +299,167 @@ TEST(SequencedTaskRunnerCurrentDefaultHandleTestWithoutTaskEnvironment,
   EXPECT_FALSE(SingleThreadTaskRunner::HasCurrentDefault());
 }
 
+TEST(SequencedTaskRunnerCurrentBestEffortTest, SequenceManagerSingleQueue) {
+  // TaskEnvironment wraps a SequenceManager with a single default task queue.
+  test::TaskEnvironment task_env;
+  EXPECT_FALSE(SequencedTaskRunner::HasCurrentBestEffort());
+  ASSERT_TRUE(SequencedTaskRunner::HasCurrentDefault());
+
+  // Should fall back to returning the default task runner when no best-effort
+  // task runner is set.
+  EXPECT_EQ(SequencedTaskRunner::GetCurrentBestEffort(),
+            SequencedTaskRunner::GetCurrentDefault());
+}
+
+TEST(SequencedTaskRunnerCurrentBestEffortTest, SequenceManagerManyQueues) {
+  // TaskEnvironmentWithMainThreadPriorities wraps a SequenceManager with
+  // several task queues.
+  test::TaskEnvironmentWithMainThreadPriorities task_env;
+  EXPECT_TRUE(SequencedTaskRunner::HasCurrentBestEffort());
+  ASSERT_TRUE(SequencedTaskRunner::HasCurrentDefault());
+  ASSERT_TRUE(SequencedTaskRunner::GetCurrentBestEffort());
+
+  // The best-effort task runner should NOT be the default.
+  EXPECT_NE(SequencedTaskRunner::GetCurrentBestEffort(),
+            SequencedTaskRunner::GetCurrentDefault());
+
+  // All should run tasks on the same sequence. They differ only in priority.
+
+  // Use SequenceCheckerImpl to make sure it's not a no-op in Release builds.
+  SequenceCheckerImpl sequence_checker;
+
+  RunLoop run_loop;
+  auto quit_closure = BarrierClosure(2, run_loop.QuitClosure());
+  SequencedTaskRunner::GetCurrentBestEffort()->PostTask(
+      FROM_HERE, BindLambdaForTesting([&] {
+                   EXPECT_TRUE(sequence_checker.CalledOnValidSequence())
+                       << "GetCurrentBestEffort";
+                 }).Then(quit_closure));
+  SequencedTaskRunner::GetCurrentDefault()->PostTask(
+      FROM_HERE, BindLambdaForTesting([&] {
+                   EXPECT_TRUE(sequence_checker.CalledOnValidSequence())
+                       << "GetCurrentDefault";
+                 }).Then(quit_closure));
+  run_loop.Run();
+}
+
+TEST(SequencedTaskRunnerCurrentBestEffortTest, ThreadPoolSingleThreadTask) {
+  test::TaskEnvironmentWithMainThreadPriorities task_env;
+
+  EXPECT_TRUE(SequencedTaskRunner::HasCurrentBestEffort());
+  auto task_env_best_effort_task_runner =
+      SequencedTaskRunner::GetCurrentBestEffort();
+
+  // Use SequenceCheckerImpl to make sure it's not a no-op in Release builds.
+  SequenceCheckerImpl sequence_checker;
+
+  // Test GetCurrentBestEffort() from a default-priority task and a BEST_EFFORT
+  // task.
+  constexpr auto kPriorities =
+      std::to_array({TaskPriority::USER_BLOCKING, TaskPriority::BEST_EFFORT});
+
+  for (auto priority : kPriorities) {
+    ThreadPool::CreateSingleThreadTaskRunner({priority})
+        ->PostTask(FROM_HERE,
+                   BindLambdaForTesting([&] {
+                     SCOPED_TRACE(priority);
+
+                     // TODO(crbug.com/441949788): Even if this is on a
+                     // BEST_EFFORT task runner, HasCurrentBestEffort() returns
+                     // false because it only supports SequenceManager task
+                     // queues.
+                     EXPECT_FALSE(SequencedTaskRunner::HasCurrentBestEffort());
+                     ASSERT_TRUE(SequencedTaskRunner::HasCurrentDefault());
+
+                     // Should fall back to returning the default task runner
+                     // when no best-effort task runner is set.
+                     EXPECT_EQ(SequencedTaskRunner::GetCurrentBestEffort(),
+                               SequencedTaskRunner::GetCurrentDefault());
+
+                     // It should NOT be the TaskEnvironment's best-effort task
+                     // runner.
+                     EXPECT_NE(SequencedTaskRunner::GetCurrentBestEffort(),
+                               task_env_best_effort_task_runner);
+                     EXPECT_FALSE(sequence_checker.CalledOnValidSequence());
+                   }).Then(task_env.QuitClosure()));
+    task_env.RunUntilQuit();
+  }
+}
+
+TEST(SequencedTaskRunnerCurrentBestEffortTest, ThreadPoolSequencedTask) {
+  test::TaskEnvironmentWithMainThreadPriorities task_env;
+
+  EXPECT_TRUE(SequencedTaskRunner::HasCurrentBestEffort());
+  auto task_env_best_effort_task_runner =
+      SequencedTaskRunner::GetCurrentBestEffort();
+
+  // Use SequenceCheckerImpl to make sure it's not a no-op in Release builds.
+  SequenceCheckerImpl sequence_checker;
+
+  // Test GetCurrentBestEffort() from a default-priority task and a BEST_EFFORT
+  // task.
+  constexpr auto kPriorities =
+      std::to_array({TaskPriority::USER_BLOCKING, TaskPriority::BEST_EFFORT});
+
+  for (auto priority : kPriorities) {
+    ThreadPool::CreateSequencedTaskRunner({priority})
+        ->PostTask(FROM_HERE,
+                   BindLambdaForTesting([&] {
+                     SCOPED_TRACE(priority);
+
+                     // TODO(crbug.com/441949788): Even if this is on a
+                     // BEST_EFFORT task runner, HasCurrentBestEffort() returns
+                     // false because it only supports SequenceManager task
+                     // queues.
+                     EXPECT_FALSE(SequencedTaskRunner::HasCurrentBestEffort());
+                     ASSERT_TRUE(SequencedTaskRunner::HasCurrentDefault());
+
+                     // Should fall back to returning the default task runner
+                     // when no best-effort task runner is set.
+                     EXPECT_EQ(SequencedTaskRunner::GetCurrentBestEffort(),
+                               SequencedTaskRunner::GetCurrentDefault());
+
+                     // It should NOT be the TaskEnvironment's best-effort task
+                     // runner.
+                     EXPECT_NE(SequencedTaskRunner::GetCurrentBestEffort(),
+                               task_env_best_effort_task_runner);
+                     EXPECT_FALSE(sequence_checker.CalledOnValidSequence());
+                   }).Then(task_env.QuitClosure()));
+    task_env.RunUntilQuit();
+  }
+}
+
+TEST(SequencedTaskRunnerCurrentBestEffortTest, ThreadPoolUnsequencedTask) {
+  test::TaskEnvironmentWithMainThreadPriorities task_env;
+
+  EXPECT_TRUE(SequencedTaskRunner::HasCurrentBestEffort());
+
+  // Test GetCurrentBestEffort() from a default-priority task and a BEST_EFFORT
+  // task.
+  constexpr auto kPriorities =
+      std::to_array({TaskPriority::USER_BLOCKING, TaskPriority::BEST_EFFORT});
+
+  for (auto priority : kPriorities) {
+    ThreadPool::PostTask(
+        FROM_HERE, {priority},
+        BindLambdaForTesting([&] {
+          SCOPED_TRACE(priority);
+
+          // The current task isn't bound to a sequence.
+          EXPECT_FALSE(SequencedTaskRunner::HasCurrentBestEffort());
+          EXPECT_FALSE(SequencedTaskRunner::HasCurrentDefault());
+        }).Then(task_env.QuitClosure()));
+    task_env.RunUntilQuit();
+  }
+}
+
+TEST(SequencedTaskRunnerCurrentBestEffortDeathTest, NoContext) {
+  EXPECT_FALSE(SequencedTaskRunner::HasCurrentBestEffort());
+  EXPECT_FALSE(SequencedTaskRunner::HasCurrentDefault());
+  // Ensure that GetCurrentBestEffort() doesn't return a value when
+  // HasCurrentDefault() is false.
+  EXPECT_CHECK_DEATH(
+      { auto task_runner = SequencedTaskRunner::GetCurrentBestEffort(); });
+}
+
 }  // namespace base

task/single_thread_task_runner.cc

@@ -14,11 +14,15 @@
 #include "base/functional/bind.h"
 #include "base/lazy_instance.h"
 #include "base/run_loop.h"
+#include "base/task/sequence_manager/sequence_manager_impl.h"
+#include "base/types/pass_key.h"
 
 namespace base {
 
 namespace {
 
+using SequenceManagerImpl = sequence_manager::internal::SequenceManagerImpl;
+
 constinit thread_local SingleThreadTaskRunner::CurrentDefaultHandle*
     current_default_handle = nullptr;
 
@@ -68,6 +72,22 @@ bool SingleThreadTaskRunner::HasCurrentDefault() {
          !!GetCurrentDefaultHandle()->task_runner_;
 }
 
+// static
+scoped_refptr<SingleThreadTaskRunner>
+SingleThreadTaskRunner::GetCurrentBestEffort() {
+  if (auto task_runner = SequenceManagerImpl::GetCurrentBestEffortTaskRunner(
+          PassKey<SingleThreadTaskRunner>())) {
+    return task_runner;
+  }
+  return GetCurrentDefault();
+}
+
+// static
+bool SingleThreadTaskRunner::HasCurrentBestEffort() {
+  return !!SequenceManagerImpl::GetCurrentBestEffortTaskRunner(
+      PassKey<SingleThreadTaskRunner>());
+}
+
 // static
 const scoped_refptr<SingleThreadTaskRunner>&
 SingleThreadTaskRunner::GetMainThreadDefault() {
@@ -86,6 +106,21 @@ bool SingleThreadTaskRunner::HasMainThreadDefault() {
          !!main_thread_default_handle->task_runner_;
 }
 
+// static
+scoped_refptr<SingleThreadTaskRunner>
+SingleThreadTaskRunner::GetMainThreadBestEffort() {
+  if (HasMainThreadBestEffort()) {
+    return main_thread_default_handle->best_effort_task_runner_;
+  }
+  return GetMainThreadDefault();
+}
+
+// static
+bool SingleThreadTaskRunner::HasMainThreadBestEffort() {
+  return !!main_thread_default_handle &&
+         !!main_thread_default_handle->best_effort_task_runner_;
+}
+
 SingleThreadTaskRunner::CurrentDefaultHandle::CurrentDefaultHandle(
     scoped_refptr<SingleThreadTaskRunner> task_runner)
     : CurrentDefaultHandle(std::move(task_runner), MayAlreadyExist{}) {
@@ -125,6 +160,11 @@ SingleThreadTaskRunner::CurrentHandleOverrideForTesting::
 SingleThreadTaskRunner::MainThreadDefaultHandle::MainThreadDefaultHandle(
     scoped_refptr<SingleThreadTaskRunner> task_runner)
     : task_runner_(std::move(task_runner)),
+      // `task_runner` belongs to this thread, so if there's a BEST_EFFORT task
+      // runner for the thread GetCurrentBestEffortTaskRunner will return it.
+      best_effort_task_runner_(
+          SequenceManagerImpl::GetCurrentBestEffortTaskRunner(
+              PassKey<SingleThreadTaskRunner>())),
       previous_handle_(main_thread_default_handle) {
   CHECK(!main_thread_default_handle || can_override);
   main_thread_default_handle = this;