[scheduler] Registering the BrowserTaskExecutor in TLS.

This is a reland of a689aa50

Bug: 872236
Change-Id: I6ab54fe95459c294358913bf6a78a6a73e8cb449
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/1787606
Commit-Queue: Karolina Soltys <ksolt@chromium.org>
Reviewed-by: Ted Choc <tedchoc@chromium.org>
Reviewed-by: Alex Clarke <alexclarke@chromium.org>
Reviewed-by: Gabriel Charette <gab@chromium.org>
Auto-Submit: Karolina Soltys <ksolt@chromium.org>
Cr-Commit-Position: refs/heads/master@{#699818}

base/task/task_executor.cc

@@ -6,8 +6,10 @@
 
 #include <type_traits>
 
+#include "base/no_destructor.h"
 #include "base/task/task_traits.h"
 #include "base/task/task_traits_extension.h"
+#include "base/threading/thread_local.h"
 
 namespace base {
 
@@ -30,6 +32,21 @@ static_assert(
 
 }  // namespace
 
+ThreadLocalPointer<TaskExecutor>* GetTLSForCurrentTaskExecutor() {
+  static NoDestructor<ThreadLocalPointer<TaskExecutor>> instance;
+  return instance.get();
+}
+
+void SetTaskExecutorForCurrentThread(TaskExecutor* task_executor) {
+  DCHECK(!task_executor || !GetTLSForCurrentTaskExecutor()->Get() ||
+         GetTLSForCurrentTaskExecutor()->Get() == task_executor);
+  GetTLSForCurrentTaskExecutor()->Set(task_executor);
+}
+
+TaskExecutor* GetTaskExecutorForCurrentThread() {
+  return GetTLSForCurrentTaskExecutor()->Get();
+}
+
 void RegisterTaskExecutor(uint8_t extension_id, TaskExecutor* task_executor) {
   DCHECK_NE(extension_id, TaskTraitsExtensionStorage::kInvalidExtensionId);
   DCHECK_LE(extension_id, TaskTraitsExtensionStorage::kMaxExtensionId);

base/task/task_executor.h

@@ -74,6 +74,13 @@ void BASE_EXPORT RegisterTaskExecutor(uint8_t extension_id,
                                       TaskExecutor* task_executor);
 void BASE_EXPORT UnregisterTaskExecutorForTesting(uint8_t extension_id);
 
+// Stores the provided TaskExecutor in TLS for the current thread, to be used by
+// tasks with the CurrentThread() trait.
+void BASE_EXPORT SetTaskExecutorForCurrentThread(TaskExecutor* task_executor);
+
+// Returns the task executor registered for the current thread.
+BASE_EXPORT TaskExecutor* GetTaskExecutorForCurrentThread();
+
 // Determines whether a registered TaskExecutor will handle tasks with the given
 // |traits| and, if so, returns a pointer to it. Otherwise, returns |nullptr|.
 TaskExecutor* GetRegisteredTaskExecutorForTraits(const TaskTraits& traits);

content/browser/browser_main_loop_unittest.cc

@@ -75,6 +75,7 @@ TEST_F(BrowserMainLoopTest, CreateThreadsInSingleProcess) {
                     {base::ThreadPool(), base::TaskPriority::USER_VISIBLE}),
             base::SysInfo::NumberOfProcessors() - 1);
   browser_main_loop.ShutdownThreadsAndCleanUp();
+  BrowserTaskExecutor::ResetForTesting();
 }
 
 TEST_F(BrowserMainLoopTest,
@@ -104,6 +105,7 @@ TEST_F(BrowserMainLoopTest,
   content::RunAllPendingInMessageLoop(BrowserThread::IO);
 
   browser_main_loop.ShutdownThreadsAndCleanUp();
+  BrowserTaskExecutor::ResetForTesting();
 }
 
 }  // namespace content

content/browser/browser_thread_unittest.cc

@@ -57,7 +57,9 @@ class SequenceManagerThreadDelegate : public base::Thread::Delegate {
 
     BrowserTaskExecutor::CreateForTesting(
         std::move(browser_ui_thread_scheduler),
-        std::make_unique<BrowserIOThreadDelegate>());
+        std::make_unique<BrowserIOThreadDelegate>(
+            BrowserIOThreadDelegate::BrowserTaskExecutorPresent::
+                kNoForTesting));
     BrowserTaskExecutor::EnableAllQueues();
   }
 
@@ -116,6 +118,7 @@ class BrowserThreadTest : public testing::Test {
 
     BrowserThreadImpl::ResetGlobalsForTesting(BrowserThread::UI);
     BrowserThreadImpl::ResetGlobalsForTesting(BrowserThread::IO);
+    BrowserTaskExecutor::ResetForTesting();
   }
 
   // Prepares this BrowserThreadTest for Release() to be invoked. |on_release|
@@ -284,7 +287,7 @@ TEST_F(BrowserThreadTest, PostTaskAndReply) {
   run_loop.Run();
 }
 
-TEST_F(BrowserThreadTest, RunsTasksInCurrentSequencedDuringShutdown) {
+TEST_F(BrowserThreadTest, RunsTasksInCurrentSequenceDuringShutdown) {
   bool did_shutdown = false;
   base::RunLoop loop;
   UIThreadDestructionObserver observer(&did_shutdown, loop.QuitClosure());
@@ -311,7 +314,9 @@ class BrowserThreadWithCustomSchedulerTest : public testing::Test {
               QueueType::kDefault));
       BrowserTaskExecutor::CreateForTesting(
           std::move(browser_ui_thread_scheduler),
-          std::make_unique<BrowserIOThreadDelegate>());
+          std::make_unique<BrowserIOThreadDelegate>(
+              BrowserIOThreadDelegate::BrowserTaskExecutorPresent::
+                  kNoForTesting));
 
       ui_thread_ = BrowserTaskExecutor::CreateIOThread();
       BrowserTaskExecutor::InitializeIOThread();

content/browser/scheduler/browser_io_thread_delegate.cc

@@ -8,6 +8,8 @@
 #include "base/message_loop/message_pump_type.h"
 #include "base/task/sequence_manager/sequence_manager.h"
 #include "base/task/sequence_manager/task_queue.h"
+#include "base/task/task_executor.h"
+#include "content/browser/scheduler/browser_task_executor.h"
 #include "content/public/browser/browser_thread.h"
 
 namespace content {
@@ -16,17 +18,20 @@ using ::base::sequence_manager::CreateUnboundSequenceManager;
 using ::base::sequence_manager::SequenceManager;
 using ::base::sequence_manager::TaskQueue;
 
-BrowserIOThreadDelegate::BrowserIOThreadDelegate()
+BrowserIOThreadDelegate::BrowserIOThreadDelegate(
+    BrowserTaskExecutorPresent browser_task_executor_present)
     : sequence_manager_(CreateUnboundSequenceManager(
           SequenceManager::Settings::Builder()
               .SetMessagePumpType(base::MessagePumpType::IO)
-              .Build())) {
+              .Build())),
+      browser_task_executor_present_(browser_task_executor_present) {
   Init(sequence_manager_.get());
 }
 
 BrowserIOThreadDelegate::BrowserIOThreadDelegate(
     SequenceManager* sequence_manager)
-    : sequence_manager_(nullptr) {
+    : sequence_manager_(nullptr),
+      browser_task_executor_present_(BrowserTaskExecutorPresent::kYes) {
   Init(sequence_manager);
 }
 
@@ -43,7 +48,11 @@ BrowserIOThreadDelegate::GetDefaultTaskRunner() {
   return default_task_runner_;
 }
 
-BrowserIOThreadDelegate::~BrowserIOThreadDelegate() = default;
+BrowserIOThreadDelegate::~BrowserIOThreadDelegate() {
+  if (browser_task_executor_present_ == BrowserTaskExecutorPresent::kYes) {
+    base::SetTaskExecutorForCurrentThread(nullptr);
+  }
+}
 
 void BrowserIOThreadDelegate::BindToCurrentThread(
     base::TimerSlack timer_slack) {
@@ -52,6 +61,10 @@ void BrowserIOThreadDelegate::BindToCurrentThread(
       base::MessagePump::Create(base::MessagePumpType::IO));
   sequence_manager_->SetTimerSlack(timer_slack);
   sequence_manager_->SetDefaultTaskRunner(GetDefaultTaskRunner());
+
+  if (browser_task_executor_present_ == BrowserTaskExecutorPresent::kYes) {
+    base::SetTaskExecutorForCurrentThread(BrowserTaskExecutor::Get());
+  }
 }
 
 }  // namespace content

content/browser/scheduler/browser_io_thread_delegate.h

@@ -28,12 +28,20 @@ class CONTENT_EXPORT BrowserIOThreadDelegate : public base::Thread::Delegate {
  public:
   using Handle = BrowserTaskQueues::Handle;
 
+  // Normally, creating a BrowserIOThreadDelegate relies on a
+  // BrowserTaskExecutor already existing to register it as the executor for the
+  // current (IO) thread. However, some tests create it in isolation, so we need
+  // to disable registering the executor to pass checks.
+  enum class BrowserTaskExecutorPresent { kYes, kNoForTesting };
+
   static std::unique_ptr<BrowserIOThreadDelegate> CreateForTesting(
       base::sequence_manager::SequenceManager* sequence_manager) {
     return base::WrapUnique(new BrowserIOThreadDelegate(sequence_manager));
   }
 
-  BrowserIOThreadDelegate();
+  explicit BrowserIOThreadDelegate(
+      BrowserTaskExecutorPresent browser_task_executor_present =
+          BrowserTaskExecutorPresent::kYes);
   ~BrowserIOThreadDelegate() override;
 
   scoped_refptr<base::SingleThreadTaskRunner> GetDefaultTaskRunner() override;
@@ -47,6 +55,10 @@ class CONTENT_EXPORT BrowserIOThreadDelegate : public base::Thread::Delegate {
   // tests.
   void SetAllowBlockingForTesting() { allow_blocking_for_testing_ = true; }
 
+  bool browser_task_executor_present() const {
+    return browser_task_executor_present_ == BrowserTaskExecutorPresent::kYes;
+  }
+
   scoped_refptr<Handle> CreateHandle() { return task_queues_->GetHandle(); }
 
  private:
@@ -64,6 +76,8 @@ class CONTENT_EXPORT BrowserIOThreadDelegate : public base::Thread::Delegate {
 
   std::unique_ptr<BrowserTaskQueues> task_queues_;
   scoped_refptr<base::SingleThreadTaskRunner> default_task_runner_;
+
+  const BrowserTaskExecutorPresent browser_task_executor_present_;
 };
 
 }  // namespace content

content/browser/scheduler/browser_io_thread_delegate_unittest.cc

@@ -9,6 +9,7 @@
 #include "base/bind.h"
 #include "base/synchronization/waitable_event.h"
 #include "base/threading/thread.h"
+#include "content/browser/scheduler/browser_task_executor.h"
 #include "content/browser/scheduler/browser_task_queues.h"
 #include "testing/gmock/include/gmock/gmock.h"
 #include "testing/gtest/include/gtest/gtest.h"
@@ -19,7 +20,8 @@ namespace {
 TEST(BrowserIOThreadDelegateTest, CanPostTasksToThread) {
   base::Thread thread("my_thread");
 
-  auto delegate = std::make_unique<BrowserIOThreadDelegate>();
+  auto delegate = std::make_unique<BrowserIOThreadDelegate>(
+      BrowserIOThreadDelegate::BrowserTaskExecutorPresent::kNoForTesting);
   auto handle = delegate->CreateHandle();
   handle->EnableAllQueues();
 
@@ -36,10 +38,11 @@ TEST(BrowserIOThreadDelegateTest, CanPostTasksToThread) {
   event.Wait();
 }
 
-TEST(BrowserIOThreadDelegateTest, DefaultTaskRunnerIsAllwaysActive) {
+TEST(BrowserIOThreadDelegateTest, DefaultTaskRunnerIsAlwaysActive) {
   base::Thread thread("my_thread");
 
-  auto delegate = std::make_unique<BrowserIOThreadDelegate>();
+  auto delegate = std::make_unique<BrowserIOThreadDelegate>(
+      BrowserIOThreadDelegate::BrowserTaskExecutorPresent::kNoForTesting);
   auto task_runner = delegate->GetDefaultTaskRunner();
 
   base::Thread::Options options;

content/browser/scheduler/browser_task_executor.cc

@@ -74,7 +74,9 @@ BrowserTaskExecutor::BrowserTaskExecutor(
       browser_io_thread_delegate_(std::move(browser_io_thread_delegate)),
       browser_io_thread_handle_(browser_io_thread_delegate_->CreateHandle()) {}
 
-BrowserTaskExecutor::~BrowserTaskExecutor() = default;
+BrowserTaskExecutor::~BrowserTaskExecutor() {
+  base::SetTaskExecutorForCurrentThread(nullptr);
+}
 
 // static
 void BrowserTaskExecutor::Create() {
@@ -104,23 +106,35 @@ void BrowserTaskExecutor::CreateInternal(
   g_browser_task_executor->browser_ui_thread_handle_
       ->EnableAllExceptBestEffortQueues();
 
+  // Here we register the BrowserTaskExecutor for the UI thread; registration
+  // for the IO thread happens in BrowserIOThreadDelegate::BindToCurrentThread.
+  base::SetTaskExecutorForCurrentThread(g_browser_task_executor);
+
 #if defined(OS_ANDROID)
   base::PostTaskAndroid::SignalNativeSchedulerReady();
 #endif
 }
 
+// static
+BrowserTaskExecutor* BrowserTaskExecutor::Get() {
+  DCHECK(g_browser_task_executor);
+  return g_browser_task_executor;
+}
+
 // static
 void BrowserTaskExecutor::ResetForTesting() {
 #if defined(OS_ANDROID)
   base::PostTaskAndroid::SignalNativeSchedulerShutdown();
 #endif
-
   if (g_browser_task_executor) {
+    RunAllPendingTasksOnThreadForTesting(BrowserThread::UI);
+    RunAllPendingTasksOnThreadForTesting(BrowserThread::IO);
     base::UnregisterTaskExecutorForTesting(
         BrowserTaskTraitsExtension::kExtensionId);
     delete g_browser_task_executor;
     g_browser_task_executor = nullptr;
   }
+  base::SetTaskExecutorForCurrentThread(nullptr);
 }
 
 // static
@@ -164,6 +178,12 @@ void BrowserTaskExecutor::RunAllPendingTasksOnThreadForTesting(
           ->ScheduleRunAllPendingTasksForTesting(run_loop.QuitClosure());
       break;
     case BrowserThread::IO: {
+      // In tests there may not be a functional IO thread.
+      if (!g_browser_task_executor->browser_io_thread_delegate_ ||
+          !g_browser_task_executor->browser_io_thread_delegate_
+               ->browser_task_executor_present()) {
+        return;
+      }
       g_browser_task_executor->browser_io_thread_handle_
           ->ScheduleRunAllPendingTasksForTesting(run_loop.QuitClosure());
       break;

content/browser/scheduler/browser_task_executor.h

@@ -148,6 +148,7 @@ class CONTENT_EXPORT BrowserTaskExecutor : public base::TaskExecutor {
 #endif  // DCHECK_IS_ON()
 
  private:
+  friend class BrowserIOThreadDelegate;
   friend class BrowserTaskExecutorTest;
 
   static void CreateInternal(
@@ -162,6 +163,10 @@ class CONTENT_EXPORT BrowserTaskExecutor : public base::TaskExecutor {
   FRIEND_TEST_ALL_PREFIXES(BrowserTaskExecutorTest,
                            BestEffortTasksRunAfterStartup);
 
+  // For Get();
+  FRIEND_TEST_ALL_PREFIXES(BrowserTaskExecutorTest,
+                           RegisterExecutorForBothThreads);
+
   explicit BrowserTaskExecutor(
       std::unique_ptr<BrowserUIThreadScheduler> browser_ui_thread_scheduler,
       std::unique_ptr<BrowserIOThreadDelegate> browser_io_thread_delegate);
@@ -170,6 +175,8 @@ class CONTENT_EXPORT BrowserTaskExecutor : public base::TaskExecutor {
   scoped_refptr<base::SingleThreadTaskRunner> GetTaskRunner(
       const base::TaskTraits& traits) const;
 
+  static BrowserTaskExecutor* Get();
+
   std::unique_ptr<BrowserUIThreadScheduler> browser_ui_thread_scheduler_;
   scoped_refptr<BrowserUIThreadScheduler::Handle> browser_ui_thread_handle_;
 

content/browser/scheduler/browser_task_executor_unittest.cc

@@ -47,6 +47,21 @@ class BrowserTaskExecutorTest : public testing::Test {
 using StrictMockTask =
     testing::StrictMock<base::MockCallback<base::RepeatingCallback<void()>>>;
 
+TEST_F(BrowserTaskExecutorTest, RegisterExecutorForBothThreads) {
+  base::PostTask(FROM_HERE, {BrowserThread::UI}, base::BindOnce([]() {
+                   EXPECT_EQ(BrowserTaskExecutor::Get(),
+                             base::GetTaskExecutorForCurrentThread());
+                 }));
+
+  base::PostTask(FROM_HERE, {BrowserThread::IO}, base::BindOnce([]() {
+                   EXPECT_EQ(BrowserTaskExecutor::Get(),
+                             base::GetTaskExecutorForCurrentThread());
+                 }));
+
+  BrowserTaskExecutor::RunAllPendingTasksOnThreadForTesting(BrowserThread::UI);
+  BrowserTaskExecutor::RunAllPendingTasksOnThreadForTesting(BrowserThread::IO);
+}
+
 TEST_F(BrowserTaskExecutorTest, RunAllPendingTasksForTestingOnUI) {
   StrictMockTask task_1;
   StrictMockTask task_2;
@@ -179,7 +194,9 @@ class BrowserTaskExecutorWithCustomSchedulerTest : public testing::Test {
               QueueType::kDefault));
       BrowserTaskExecutor::CreateForTesting(
           std::move(browser_ui_thread_scheduler),
-          std::make_unique<BrowserIOThreadDelegate>());
+          std::make_unique<BrowserIOThreadDelegate>(
+              BrowserIOThreadDelegate::BrowserTaskExecutorPresent::
+                  kNoForTesting));
     }
   };