Correctly join thread in TaskDispatchThread (#53943)

Summary:
Pull Request resolved: #53943

Changelog: [Internal]

Thread join should happen only during TaskDispatchThread dtor, otherwise if quit is called on the same thread (thread_), then thread_ will get detached, which can lead to use of already deleted fields (queue_, running_) of TaskDispatchThread.

But we also need to make sure that on quit() call, loop is getting stopped, otherwise in loop, already running task will use already deleted fields of captured object. If quit() is called as task on thread_, then there is no need to wait for future because loop will organically end since running_ is false.

Reviewed By: rshest

Differential Revision: D83044012

fbshipit-source-id: a2da97a89093c47c64b693e3c9e5d6f7297c038e

Author: Nurtau Toganbay

packages/react-native/ReactCxxPlatform/react/threading/TaskDispatchThread.cpp

@@ -50,6 +50,10 @@ TaskDispatchThread::TaskDispatchThread(
 
 TaskDispatchThread::~TaskDispatchThread() noexcept {
   quit();
+
+  if (thread_.joinable()) {
+    thread_.join();
+  }
 }
 
 bool TaskDispatchThread::isOnThread() noexcept {
@@ -87,17 +91,14 @@ void TaskDispatchThread::runSync(TaskFn&& task) noexcept {
 }
 
 void TaskDispatchThread::quit() noexcept {
-  if (!running_) {
+  bool expected = true;
+  if (!running_.compare_exchange_strong(expected, false)) {
     return;
   }
-  running_ = false;
+
   loopCv_.notify_one();
-  if (thread_.joinable()) {
-    if (!isOnThread()) {
-      thread_.join();
-    } else {
-      thread_.detach();
-    }
+  if (!isOnThread()) {
+    loopStoppedPromise_.get_future().wait();
   }
 }
 
@@ -108,6 +109,7 @@ void TaskDispatchThread::loop() noexcept {
   while (running_) {
     std::unique_lock<std::mutex> lock(queueLock_);
     loopCv_.wait(lock, [&]() { return !running_ || !queue_.empty(); });
+
     while (!queue_.empty()) {
       auto task = queue_.top();
       auto now = std::chrono::system_clock::now();
@@ -120,11 +122,11 @@ void TaskDispatchThread::loop() noexcept {
       } else {
         // Shutting down, skip all the remaining tasks
         queue_ = {};
-        return;
+        break;
       }
 
-      // We should check whether the task thread is still running at this point
-      // (which may not anymore be the case since the previous check)
+      // We should check whether the task thread is still running at this
+      // point (which may not anymore be the case since the previous check)
       if (running_) {
         lock.unlock();
         task.fn();
@@ -133,6 +135,8 @@ void TaskDispatchThread::loop() noexcept {
       queue_.pop();
     }
   }
+
+  loopStoppedPromise_.set_value();
 }
 
 } //  namespace facebook::react

packages/react-native/ReactCxxPlatform/react/threading/TaskDispatchThread.h

@@ -10,6 +10,7 @@
 #include <chrono>
 #include <condition_variable>
 #include <functional>
+#include <future>
 #include <mutex>
 #include <queue>
 #include <thread>
@@ -71,6 +72,8 @@ class TaskDispatchThread {
   std::atomic<bool> running_{true};
   std::string threadName_;
   std::thread thread_;
+
+  std::promise<void> loopStoppedPromise_;
 };
 
 } // namespace facebook::react

packages/react-native/ReactCxxPlatform/react/threading/tests/TaskDispatchThreadTests.cpp

@@ -15,27 +15,28 @@ namespace facebook::react {
 
 class TaskDispatchThreadTest : public ::testing::Test {
  protected:
-  TaskDispatchThread dispatcher;
+  std::unique_ptr<TaskDispatchThread> dispatcher{
+      std::make_unique<TaskDispatchThread>()};
 };
 
 // Test: isOnThread returns true inside the looper thread
 TEST_F(TaskDispatchThreadTest, IsOnThreadReturnsTrueInLooper) {
   bool result = false;
-  dispatcher.runSync([&] { result = dispatcher.isOnThread(); });
+  dispatcher->runSync([&] { result = dispatcher->isOnThread(); });
   EXPECT_TRUE(result);
 }
 
 // Test: isRunning returns true before quit, false after
 TEST_F(TaskDispatchThreadTest, IsRunningFlag) {
-  EXPECT_TRUE(dispatcher.isRunning());
-  dispatcher.quit();
-  EXPECT_FALSE(dispatcher.isRunning());
+  EXPECT_TRUE(dispatcher->isRunning());
+  dispatcher->quit();
+  EXPECT_FALSE(dispatcher->isRunning());
 }
 
 // Test: runAsync executes the task
 TEST_F(TaskDispatchThreadTest, RunAsyncExecutesTask) {
   std::atomic<int> counter{0};
-  dispatcher.runAsync([&] { counter++; });
+  dispatcher->runAsync([&] { counter++; });
   // Wait for task to complete
   std::this_thread::sleep_for(std::chrono::milliseconds(50));
   EXPECT_EQ(counter.load(), 1);
@@ -44,14 +45,14 @@ TEST_F(TaskDispatchThreadTest, RunAsyncExecutesTask) {
 // Test: runSync executes the task and blocks until done
 TEST_F(TaskDispatchThreadTest, RunSyncExecutesTask) {
   std::atomic<int> counter{0};
-  dispatcher.runSync([&] { counter++; });
+  dispatcher->runSync([&] { counter++; });
   EXPECT_EQ(counter.load(), 1);
 }
 
 // Test: runAsync with delay
 TEST_F(TaskDispatchThreadTest, RunAsyncWithDelay) {
   std::atomic<int> counter{0};
-  dispatcher.runAsync([&] { counter++; }, std::chrono::milliseconds(100));
+  dispatcher->runAsync([&] { counter++; }, std::chrono::milliseconds(100));
   std::this_thread::sleep_for(std::chrono::milliseconds(50));
   EXPECT_EQ(counter.load(), 0); // Not yet executed
   std::this_thread::sleep_for(std::chrono::milliseconds(70));
@@ -61,9 +62,9 @@ TEST_F(TaskDispatchThreadTest, RunAsyncWithDelay) {
 // Test: Multiple delayed tasks execute in order
 TEST_F(TaskDispatchThreadTest, MultipleDelayedTasksOrder) {
   std::vector<int> results;
-  dispatcher.runAsync(
+  dispatcher->runAsync(
       [&] { results.push_back(1); }, std::chrono::milliseconds(50));
-  dispatcher.runAsync(
+  dispatcher->runAsync(
       [&] { results.push_back(2); }, std::chrono::milliseconds(100));
   std::this_thread::sleep_for(std::chrono::milliseconds(120));
   ASSERT_EQ(results.size(), 2);
@@ -75,48 +76,48 @@ TEST_F(TaskDispatchThreadTest, MultipleDelayedTasksOrder) {
 TEST_F(TaskDispatchThreadTest, RunSyncBlocksUntilDone) {
   std::atomic<bool> started{false};
   std::atomic<bool> finished{false};
-  dispatcher.runAsync([&] {
+  dispatcher->runAsync([&] {
     started = true;
     std::this_thread::sleep_for(std::chrono::milliseconds(50));
     finished = true;
   });
-  dispatcher.runSync([&] {
+  dispatcher->runSync([&] {
     EXPECT_TRUE(started);
     EXPECT_TRUE(finished);
   });
 }
 
 // Test: quit prevents further tasks from running
 TEST_F(TaskDispatchThreadTest, QuitPreventsFurtherTasks) {
-  dispatcher.quit();
+  dispatcher->quit();
   std::atomic<int> counter{0};
-  dispatcher.runAsync([&] { counter++; });
+  dispatcher->runAsync([&] { counter++; });
   std::this_thread::sleep_for(std::chrono::milliseconds(50));
   EXPECT_EQ(counter.load(), 0);
 }
 
 // Test: Multiple runSync tasks execute serially
 TEST_F(TaskDispatchThreadTest, MultipleRunSyncSerialExecution) {
   std::vector<int> results;
-  dispatcher.runSync([&] { results.push_back(1); });
-  dispatcher.runSync([&] { results.push_back(2); });
+  dispatcher->runSync([&] { results.push_back(1); });
+  dispatcher->runSync([&] { results.push_back(2); });
   EXPECT_EQ(results.size(), 2);
   EXPECT_EQ(results[0], 1);
   EXPECT_EQ(results[1], 2);
 }
 
 // Test: Edge case - runSync after quit should not execute
 TEST_F(TaskDispatchThreadTest, RunSyncAfterQuitDoesNotExecute) {
-  dispatcher.quit();
+  dispatcher->quit();
   std::atomic<int> counter{0};
-  dispatcher.runSync([&] { counter++; });
+  dispatcher->runSync([&] { counter++; });
   EXPECT_EQ(counter.load(), 0);
 }
 
 // Test: Thread safety - runAsync from multiple threads
 TEST_F(TaskDispatchThreadTest, RunAsyncFromMultipleThreads) {
   std::atomic<int> counter{0};
-  auto task = [&] { dispatcher.runAsync([&] { counter++; }); };
+  auto task = [&] { dispatcher->runAsync([&] { counter++; }); };
   std::thread t1(task);
   std::thread t2(task);
   std::thread t3(task);
@@ -127,4 +128,23 @@ TEST_F(TaskDispatchThreadTest, RunAsyncFromMultipleThreads) {
   EXPECT_EQ(counter.load(), 3);
 }
 
+TEST_F(TaskDispatchThreadTest, QuitInTaskShouldntBeBlockedForever) {
+  dispatcher->runSync([&] { dispatcher->quit(); });
+}
+
+TEST_F(TaskDispatchThreadTest, QuitShouldWaitAlreadyRunningTask) {
+  {
+    std::unique_ptr<int> counter = std::make_unique<int>(0);
+    dispatcher->runAsync([&] {
+      std::this_thread::sleep_for(std::chrono::milliseconds(300));
+      *counter = 1;
+    });
+    std::this_thread::sleep_for(std::chrono::milliseconds(50));
+    // if quit doesn't wait for running task, then *counter will access already
+    // deleted object
+    dispatcher->quit();
+  }
+  // forcing dispatcher to join thread
+  dispatcher.reset();
+}
 } // namespace facebook::react