handle re-entry in RuntimeScheduler::executeNowOnTheSameThread (#44606)

Summary:
Pull Request resolved: https://github.com/facebook/react-native/pull/44606

changelog: [internal]

Protect RuntimeScheduler::executeNowOnTheSameThread against re-entry to prevent deadlocks.

Reviewed By: NickGerleman

Differential Revision: D57514317

fbshipit-source-id: 88309aa27fedefcce7bf8ce5be9a382fdf72cfae

packages/react-native/ReactCommon/react/renderer/runtimescheduler/RuntimeScheduler_Legacy.cpp

@@ -101,20 +101,30 @@ void RuntimeScheduler_Legacy::executeNowOnTheSameThread(
     RawCallback&& callback) {
   SystraceSection s("RuntimeScheduler::executeNowOnTheSameThread");
 
-  runtimeAccessRequests_ += 1;
-  executeSynchronouslyOnSameThread_CAN_DEADLOCK(
-      runtimeExecutor_,
-      [this, callback = std::move(callback)](jsi::Runtime& runtime) {
-        SystraceSection s2(
-            "RuntimeScheduler::executeNowOnTheSameThread callback");
+  static thread_local jsi::Runtime* runtimePtr = nullptr;
 
-        runtimeAccessRequests_ -= 1;
-        {
-          ScopedShadowTreeRevisionLock revisionLock(
-              shadowTreeRevisionConsistencyManager_);
-          callback(runtime);
-        }
-      });
+  if (runtimePtr == nullptr) {
+    runtimeAccessRequests_ += 1;
+    executeSynchronouslyOnSameThread_CAN_DEADLOCK(
+        runtimeExecutor_, [this, &callback](jsi::Runtime& runtime) {
+          SystraceSection s2(
+              "RuntimeScheduler::executeNowOnTheSameThread callback");
+
+          runtimeAccessRequests_ -= 1;
+          {
+            ScopedShadowTreeRevisionLock revisionLock(
+                shadowTreeRevisionConsistencyManager_);
+            runtimePtr = &runtime;
+            callback(runtime);
+            runtimePtr = nullptr;
+          }
+        });
+  } else {
+    // Protecting against re-entry into `executeNowOnTheSameThread` from within
+    // `executeNowOnTheSameThread`. Without accounting for re-rentry, a deadlock
+    // will occur when trying to gain access to the runtime.
+    return callback(*runtimePtr);
+  }
 
   // Resume work loop if needed. In synchronous mode
   // only expired tasks are executed. Tasks with lower priority

packages/react-native/ReactCommon/react/renderer/runtimescheduler/RuntimeScheduler_Modern.cpp

@@ -91,24 +91,33 @@ void RuntimeScheduler_Modern::executeNowOnTheSameThread(
     RawCallback&& callback) {
   SystraceSection s("RuntimeScheduler::executeNowOnTheSameThread");
 
-  syncTaskRequests_++;
+  static thread_local jsi::Runtime* runtimePtr = nullptr;
 
-  executeSynchronouslyOnSameThread_CAN_DEADLOCK(
-      runtimeExecutor_,
-      [this, callback = std::move(callback)](jsi::Runtime& runtime) mutable {
-        SystraceSection s2(
-            "RuntimeScheduler::executeNowOnTheSameThread callback");
+  auto currentTime = now_();
+  auto priority = SchedulerPriority::ImmediatePriority;
+  auto expirationTime = currentTime + timeoutForSchedulerPriority(priority);
+  Task task{priority, std::move(callback), expirationTime};
 
-        syncTaskRequests_--;
+  if (runtimePtr == nullptr) {
+    syncTaskRequests_++;
+    executeSynchronouslyOnSameThread_CAN_DEADLOCK(
+        runtimeExecutor_,
+        [this, currentTime, &task](jsi::Runtime& runtime) mutable {
+          SystraceSection s2(
+              "RuntimeScheduler::executeNowOnTheSameThread callback");
 
-        auto currentTime = now_();
-        auto priority = SchedulerPriority::ImmediatePriority;
-        auto expirationTime =
-            currentTime + timeoutForSchedulerPriority(priority);
+          syncTaskRequests_--;
+          runtimePtr = &runtime;
+          executeTask(runtime, task, currentTime);
+          runtimePtr = nullptr;
+        });
 
-        auto task = Task{priority, std::move(callback), expirationTime};
-        executeTask(runtime, task, currentTime);
-      });
+  } else {
+    // Protecting against re-entry into `executeNowOnTheSameThread` from within
+    // `executeNowOnTheSameThread`. Without accounting for re-rentry, a deadlock
+    // will occur when trying to gain access to the runtime.
+    return executeTask(*runtimePtr, task, currentTime);
+  }
 
   bool shouldScheduleWorkLoop = false;
 

packages/react-native/ReactCommon/react/renderer/runtimescheduler/tests/RuntimeSchedulerTest.cpp

@@ -841,6 +841,48 @@ TEST_P(RuntimeSchedulerTest, sameThreadTaskCreatesImmediatePriorityTask) {
   EXPECT_EQ(stubQueue_->size(), 0);
 }
 
+TEST_P(RuntimeSchedulerTest, syncAccessReentryProtection) {
+  bool didRunFirstSyncTask = false;
+  bool didRunSecondSyncTask = false;
+  std::thread t1([this, &didRunFirstSyncTask, &didRunSecondSyncTask]() {
+    runtimeScheduler_->executeNowOnTheSameThread(
+        [this, &didRunFirstSyncTask, &didRunSecondSyncTask](
+            jsi::Runtime& /*runtime*/) {
+          didRunFirstSyncTask = true;
+          runtimeScheduler_->executeNowOnTheSameThread(
+              [&didRunSecondSyncTask](jsi::Runtime& /*runtime*/) {
+                EXPECT_FALSE(didRunSecondSyncTask);
+                didRunSecondSyncTask = true;
+              });
+
+          EXPECT_TRUE(didRunSecondSyncTask);
+        });
+  });
+
+  auto hasTask = stubQueue_->waitForTask();
+
+  EXPECT_TRUE(hasTask);
+  EXPECT_FALSE(didRunFirstSyncTask);
+  EXPECT_FALSE(didRunSecondSyncTask);
+  EXPECT_EQ(stubQueue_->size(), 1);
+
+  stubQueue_->tick();
+  t1.join();
+
+  if (GetParam()) {
+    EXPECT_EQ(stubQueue_->size(), 0);
+  } else {
+    // The legacy RuntimeScheduler always schedules a task and within the task
+    // it checks if the task queue is empty. Unlike the modern RuntimeScheduler,
+    // which bails out before scheduling a task.
+    EXPECT_EQ(stubQueue_->size(), 1);
+  }
+
+  EXPECT_TRUE(didRunFirstSyncTask);
+  EXPECT_TRUE(didRunSecondSyncTask);
+  EXPECT_FALSE(runtimeScheduler_->getShouldYield());
+}
+
 TEST_P(RuntimeSchedulerTest, sameThreadTaskCreatesLowPriorityTask) {
   bool didRunSynchronousTask = false;
   bool didRunSubsequentTask = false;