[MBI] Make ThreadTaskRunnerHandle::Get() return currently active...

[MBI] Make ThreadTaskRunnerHandle::Get() return currently active AgentGroupSchedulerImpl's default task runer.

[related design doc]
https://docs.google.com/document/d/1y-vHkrD1z2RtyWYwT6rJkSLHClYNjDDpUbTtBU7l95A

Bug: 1105403
Change-Id: I4a215c8cd93f10e894071af6836027289d8225a4
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2413728Reviewed-by: Dominic Farolino <dom@chromium.org>
Reviewed-by: Tal Pressman <talp@chromium.org>
Reviewed-by: Kouhei Ueno <kouhei@chromium.org>
Reviewed-by: Kentaro Hara <haraken@chromium.org>
Commit-Queue: Minoru Chikamune <chikamune@chromium.org>
Cr-Commit-Position: refs/heads/master@{#810043}

chrome/renderer/autofill/password_autofill_agent_browsertest.cc

@@ -1264,6 +1264,7 @@ TEST_F(PasswordAutofillAgentTest, SendPasswordFormsTest_ReloadTab) {
   std::string url_string = "data:text/html;charset=utf-8,";
   url_string.append(kNonVisibleFormHTML);
   Reload(GURL(url_string));
+  base::RunLoop().RunUntilIdle();
   EXPECT_TRUE(fake_driver_.called_password_forms_parsed());
   ASSERT_TRUE(fake_driver_.form_data_parsed());
   EXPECT_FALSE(fake_driver_.form_data_parsed()->empty());

third_party/blink/renderer/platform/scheduler/main_thread/agent_group_scheduler_impl.cc

@@ -10,10 +10,12 @@
 namespace blink {
 namespace scheduler {
 
-MainThreadTaskQueue::QueueCreationParams DefaultTaskQueueCreationParams() {
+MainThreadTaskQueue::QueueCreationParams DefaultTaskQueueCreationParams(
+    AgentGroupSchedulerImpl* agent_group_scheduler_impl) {
   return MainThreadTaskQueue::QueueCreationParams(
              MainThreadTaskQueue::QueueType::kDefault)
-      .SetShouldMonitorQuiescence(true);
+      .SetShouldMonitorQuiescence(true)
+      .SetAgentGroupScheduler(agent_group_scheduler_impl);
 }
 
 static AgentGroupSchedulerImpl* g_current_agent_group_scheduler_impl;
@@ -34,16 +36,14 @@ void AgentGroupSchedulerImpl::SetCurrent(
 AgentGroupSchedulerImpl::AgentGroupSchedulerImpl(
     MainThreadSchedulerImpl* main_thread_scheduler)
     : default_task_queue_(main_thread_scheduler->NewTaskQueue(
-          DefaultTaskQueueCreationParams())),
+          DefaultTaskQueueCreationParams(this))),
       default_task_runner_(default_task_queue_->CreateTaskRunner(
           TaskType::kMainThreadTaskQueueDefault)),
       main_thread_scheduler_(main_thread_scheduler) {}
 
 AgentGroupSchedulerImpl::~AgentGroupSchedulerImpl() {
   default_task_queue_->ShutdownTaskQueue();
-  if (main_thread_scheduler_) {
   main_thread_scheduler_->RemoveAgentGroupScheduler(this);
-  }
 }
 
 }  // namespace scheduler

third_party/blink/renderer/platform/scheduler/main_thread/main_thread_scheduler_impl.cc

@@ -2447,11 +2447,19 @@ void MainThreadSchedulerImpl::OnTaskStarted(
     MainThreadTaskQueue* queue,
     const base::sequence_manager::Task& task,
     const TaskQueue::TaskTiming& task_timing) {
-  if (queue && queue->GetFrameScheduler()) {
-    AgentGroupSchedulerImpl::SetCurrent(
-        queue->GetFrameScheduler()->GetAgentGroupScheduler());
+  // Switch current active scheduler
+  {
+    AgentGroupSchedulerImpl* agent_group_scheduler =
+        queue ? queue->GetAgentGroupScheduler() : nullptr;
+    AgentGroupSchedulerImpl::SetCurrent(agent_group_scheduler);
+    if (agent_group_scheduler) {
+      sequence_manager_->SetDefaultTaskRunner(
+          agent_group_scheduler->DefaultTaskRunner());
     } else {
-    AgentGroupSchedulerImpl::SetCurrent(nullptr);
+      // If there is no proper AgentGroupScheduler, it means that the
+      // current scheduler is MainThreadScheduler.
+      sequence_manager_->SetDefaultTaskRunner(helper_.DefaultTaskRunner());
+    }
   }
 
   main_thread_only().running_queues.push(queue);
@@ -2517,6 +2525,7 @@ void MainThreadSchedulerImpl::OnTaskCompleted(
   // AgentGroupSchedulerImpl::GetCurrent() should return nullptr when
   // it's running thread global task runners.
   AgentGroupSchedulerImpl::SetCurrent(nullptr);
+  sequence_manager_->SetDefaultTaskRunner(helper_.DefaultTaskRunner());
 }
 
 void MainThreadSchedulerImpl::RecordTaskUkm(

third_party/blink/renderer/platform/scheduler/main_thread/main_thread_task_queue.cc

@@ -187,10 +187,25 @@ void MainThreadTaskQueue::ShutdownTaskQueue() {
   TaskQueue::ShutdownTaskQueue();
 }
 
+AgentGroupSchedulerImpl* MainThreadTaskQueue::GetAgentGroupScheduler() {
+  DCHECK(task_runner()->BelongsToCurrentThread());
+  if (agent_group_scheduler_) {
+    DCHECK(!frame_scheduler_);
+    return agent_group_scheduler_;
+  }
+  if (frame_scheduler_) {
+    return frame_scheduler_->GetAgentGroupScheduler();
+  }
+  // If this MainThreadTaskQueue was created for MainThreadSchedulerImpl, this
+  // queue will not be associated with AgentGroupScheduler or FrameScheduler.
+  return nullptr;
+}
+
 void MainThreadTaskQueue::ClearReferencesToSchedulers() {
   if (main_thread_scheduler_)
     main_thread_scheduler_->OnShutdownTaskQueue(this);
   main_thread_scheduler_ = nullptr;
+  agent_group_scheduler_ = nullptr;
   frame_scheduler_ = nullptr;
 }
 

third_party/blink/renderer/platform/scheduler/main_thread/main_thread_task_queue.h

@@ -11,6 +11,7 @@
 #include "base/task/sequence_manager/task_queue.h"
 #include "base/task/sequence_manager/task_queue_impl.h"
 #include "net/base/request_priority.h"
+#include "third_party/blink/renderer/platform/scheduler/main_thread/agent_group_scheduler_impl.h"
 #include "third_party/blink/renderer/platform/scheduler/public/frame_scheduler.h"
 #include "third_party/blink/renderer/platform/scheduler/public/web_scheduling_priority.h"
 
@@ -291,10 +292,17 @@ class PLATFORM_EXPORT MainThreadTaskQueue
 
     // Forwarded calls to |spec|.
 
+    QueueCreationParams SetAgentGroupScheduler(
+        AgentGroupSchedulerImpl* scheduler) {
+      agent_group_scheduler = scheduler;
+      return *this;
+    }
+
     QueueCreationParams SetFrameScheduler(FrameSchedulerImpl* scheduler) {
       frame_scheduler = scheduler;
       return *this;
     }
+
     QueueCreationParams SetShouldMonitorQuiescence(bool should_monitor) {
       spec = spec.SetShouldMonitorQuiescence(should_monitor);
       return *this;
@@ -313,6 +321,7 @@ class PLATFORM_EXPORT MainThreadTaskQueue
 
     QueueType queue_type;
     base::sequence_manager::TaskQueue::Spec spec;
+    AgentGroupSchedulerImpl* agent_group_scheduler;
     FrameSchedulerImpl* frame_scheduler;
     QueueTraits queue_traits;
     bool freeze_when_keep_active;
@@ -385,6 +394,8 @@ class PLATFORM_EXPORT MainThreadTaskQueue
   // Override base method to notify MainThreadScheduler about shutdown queue.
   void ShutdownTaskQueue() override;
 
+  AgentGroupSchedulerImpl* GetAgentGroupScheduler();
+
   FrameSchedulerImpl* GetFrameScheduler() const;
 
   scoped_refptr<base::SingleThreadTaskRunner> CreateTaskRunner(
@@ -442,6 +453,8 @@ class PLATFORM_EXPORT MainThreadTaskQueue
   // Needed to notify renderer scheduler about completed tasks.
   MainThreadSchedulerImpl* main_thread_scheduler_;  // NOT OWNED
 
+  AgentGroupSchedulerImpl* agent_group_scheduler_{nullptr};  // NOT OWNED
+
   // Set in the constructor. Cleared in ClearReferencesToSchedulers(). Can never
   // be set to a different value afterwards (except in tests).
   FrameSchedulerImpl* frame_scheduler_;  // NOT OWNED