refactor the task_runner and task_runner_checker

fml/memory/task_runner_checker.cc

@@ -13,18 +13,21 @@ TaskRunnerChecker::~TaskRunnerChecker() = default;
 
 bool TaskRunnerChecker::RunsOnCreationTaskRunner() const {
   FML_CHECK(fml::MessageLoop::IsInitializedForCurrentThread());
-
   const auto current_queue_id = MessageLoop::GetCurrentTaskQueueId();
+  return RunsOnTheSameThread(current_queue_id, initialized_queue_id_);
+};
 
-  if (current_queue_id == initialized_queue_id_) {
+bool TaskRunnerChecker::RunsOnTheSameThread(TaskQueueId queue_a,
+                                            TaskQueueId queue_b) {
+  if (queue_a == queue_b) {
     return true;
   }
 
   auto queues = MessageLoopTaskQueues::GetInstance();
-  if (queues->Owns(current_queue_id, initialized_queue_id_)) {
+  if (queues->Owns(queue_a, queue_b)) {
     return true;
   }
-  if (queues->Owns(initialized_queue_id_, current_queue_id)) {
+  if (queues->Owns(queue_b, queue_a)) {
     return true;
   }
   return false;

fml/memory/task_runner_checker.h

@@ -18,6 +18,8 @@ class TaskRunnerChecker final {
 
   bool RunsOnCreationTaskRunner() const;
 
+  static bool RunsOnTheSameThread(TaskQueueId queue_a, TaskQueueId queue_b);
+
  private:
   TaskQueueId initialized_queue_id_;
 

fml/memory/task_runner_checker_unittest.cc

@@ -4,15 +4,16 @@
 
 #define FML_USED_ON_EMBEDDER
 
-#include <thread>
+#include <gtest/gtest.h>
 
+#include <thread>
 #include "flutter/fml/memory/task_runner_checker.h"
+
 #include "flutter/fml/message_loop.h"
+#include "flutter/fml/raster_thread_merger.h"
 #include "flutter/fml/synchronization/count_down_latch.h"
 #include "flutter/fml/synchronization/waitable_event.h"
 
-#include <gtest/gtest.h>
-
 namespace fml {
 namespace testing {
 
@@ -41,5 +42,78 @@ TEST(TaskRunnerCheckerTests, FailsTheCheckIfOnDifferentTaskRunner) {
   EXPECT_EQ(checker.RunsOnCreationTaskRunner(), true);
 }
 
+TEST(TaskRunnerCheckerTests, SameTaskRunnerRunsOnTheSameThread) {
+  fml::MessageLoop& loop1 = fml::MessageLoop::GetCurrent();
+  fml::MessageLoop& loop2 = fml::MessageLoop::GetCurrent();
+  TaskQueueId a = loop1.GetTaskRunner()->GetTaskQueueId();
+  TaskQueueId b = loop2.GetTaskRunner()->GetTaskQueueId();
+  EXPECT_EQ(TaskRunnerChecker::RunsOnTheSameThread(a, b), true);
+}
+
+TEST(TaskRunnerCheckerTests, RunsOnDifferentThreadsReturnsFalse) {
+  fml::MessageLoop& loop1 = fml::MessageLoop::GetCurrent();
+  TaskQueueId a = loop1.GetTaskRunner()->GetTaskQueueId();
+  fml::AutoResetWaitableEvent latch;
+  std::thread anotherThread([&]() {
+    fml::MessageLoop::EnsureInitializedForCurrentThread();
+    fml::MessageLoop& loop2 = fml::MessageLoop::GetCurrent();
+    TaskQueueId b = loop2.GetTaskRunner()->GetTaskQueueId();
+    EXPECT_EQ(TaskRunnerChecker::RunsOnTheSameThread(a, b), false);
+    latch.Signal();
+  });
+  latch.Wait();
+  anotherThread.join();
+}
+
+TEST(TaskRunnerCheckerTests, MergedTaskRunnersRunsOnTheSameThread) {
+  fml::MessageLoop* loop1 = nullptr;
+  fml::AutoResetWaitableEvent latch1;
+  fml::AutoResetWaitableEvent term1;
+  std::thread thread1([&loop1, &latch1, &term1]() {
+    fml::MessageLoop::EnsureInitializedForCurrentThread();
+    loop1 = &fml::MessageLoop::GetCurrent();
+    latch1.Signal();
+    term1.Wait();
+  });
+
+  fml::MessageLoop* loop2 = nullptr;
+  fml::AutoResetWaitableEvent latch2;
+  fml::AutoResetWaitableEvent term2;
+  std::thread thread2([&loop2, &latch2, &term2]() {
+    fml::MessageLoop::EnsureInitializedForCurrentThread();
+    loop2 = &fml::MessageLoop::GetCurrent();
+    latch2.Signal();
+    term2.Wait();
+  });
+
+  latch1.Wait();
+  latch2.Wait();
+  fml::TaskQueueId qid1 = loop1->GetTaskRunner()->GetTaskQueueId();
+  fml::TaskQueueId qid2 = loop2->GetTaskRunner()->GetTaskQueueId();
+  const auto raster_thread_merger_ =
+      fml::MakeRefCounted<fml::RasterThreadMerger>(qid1, qid2);
+  const int kNumFramesMerged = 5;
+
+  raster_thread_merger_->MergeWithLease(kNumFramesMerged);
+
+  // merged, running on the same thread
+  EXPECT_EQ(TaskRunnerChecker::RunsOnTheSameThread(qid1, qid2), true);
+
+  for (int i = 0; i < kNumFramesMerged; i++) {
+    ASSERT_TRUE(raster_thread_merger_->IsMerged());
+    raster_thread_merger_->DecrementLease();
+  }
+
+  ASSERT_FALSE(raster_thread_merger_->IsMerged());
+
+  // un-merged, not running on the same thread
+  EXPECT_EQ(TaskRunnerChecker::RunsOnTheSameThread(qid1, qid2), false);
+
+  term1.Signal();
+  term2.Signal();
+  thread1.join();
+  thread2.join();
+}
+
 }  // namespace testing
 }  // namespace fml

fml/task_runner.cc

@@ -5,6 +5,7 @@
 #define FML_USED_ON_EMBEDDER
 
 #include "flutter/fml/task_runner.h"
+#include "flutter/fml/memory/task_runner_checker.h"
 
 #include <utility>
 
@@ -47,19 +48,8 @@ bool TaskRunner::RunsTasksOnCurrentThread() {
   const auto current_queue_id = MessageLoop::GetCurrentTaskQueueId();
   const auto loop_queue_id = loop_->GetTaskQueueId();
 
-  if (current_queue_id == loop_queue_id) {
-    return true;
-  }
-
-  auto queues = MessageLoopTaskQueues::GetInstance();
-  if (queues->Owns(current_queue_id, loop_queue_id)) {
-    return true;
-  }
-  if (queues->Owns(loop_queue_id, current_queue_id)) {
-    return true;
-  }
-
-  return false;
+  return TaskRunnerChecker::RunsOnTheSameThread(current_queue_id,
+                                                loop_queue_id);
 }
 
 void TaskRunner::RunNowOrPostTask(fml::RefPtr<fml::TaskRunner> runner,