Yield for each access to the runtime

Summary:
changelog: [internal]

With multiple requests to the runtime, we need to make sure they are all granted before React continues with rendering. A boolean is not enough to track this. The first lambda that has VM will set it to false and subsequent requests will have to wait for React to finish rendering.

To prevent this, we can count how many lambdas are pending access to the runtime.

Reviewed By: ShikaSD

Differential Revision: D33792734

fbshipit-source-id: f785fae3575470179dd69acc6a466211b79b633b

Author: sammy-SC

ReactCommon/react/renderer/runtimescheduler/RuntimeScheduler.cpp

@@ -23,10 +23,10 @@ RuntimeScheduler::RuntimeScheduler(
 void RuntimeScheduler::scheduleWork(
     std::function<void(jsi::Runtime &)> callback) const {
   if (enableYielding_) {
-    shouldYield_ = true;
+    runtimeAccessRequests_ += 1;
     runtimeExecutor_(
         [this, callback = std::move(callback)](jsi::Runtime &runtime) {
-          shouldYield_ = false;
+          runtimeAccessRequests_ -= 1;
           callback(runtime);
           startWorkLoop(runtime);
         });
@@ -51,7 +51,7 @@ std::shared_ptr<Task> RuntimeScheduler::scheduleTask(
 }
 
 bool RuntimeScheduler::getShouldYield() const noexcept {
-  return shouldYield_;
+  return runtimeAccessRequests_ > 0;
 }
 
 bool RuntimeScheduler::getIsSynchronous() const noexcept {
@@ -76,11 +76,11 @@ void RuntimeScheduler::setEnableYielding(bool enableYielding) {
 
 void RuntimeScheduler::executeNowOnTheSameThread(
     std::function<void(jsi::Runtime &runtime)> callback) {
-  shouldYield_ = true;
+  runtimeAccessRequests_ += 1;
   executeSynchronouslyOnSameThread_CAN_DEADLOCK(
       runtimeExecutor_,
       [this, callback = std::move(callback)](jsi::Runtime &runtime) {
-        shouldYield_ = false;
+        runtimeAccessRequests_ -= 1;
         isSynchronous_ = true;
         callback(runtime);
         isSynchronous_ = false;

ReactCommon/react/renderer/runtimescheduler/RuntimeScheduler.h

@@ -117,7 +117,12 @@ class RuntimeScheduler final {
 
   RuntimeExecutor const runtimeExecutor_;
   mutable SchedulerPriority currentPriority_{SchedulerPriority::NormalPriority};
-  mutable std::atomic_bool shouldYield_{false};
+
+  /*
+   * Counter indicating how many access to the runtime have been requested.
+   */
+  mutable std::atomic<uint_fast8_t> runtimeAccessRequests_{0};
+
   mutable std::atomic_bool isSynchronous_{false};
 
   void startWorkLoop(jsi::Runtime &runtime) const;

ReactCommon/react/renderer/runtimescheduler/tests/RuntimeSchedulerTest.cpp

@@ -43,6 +43,7 @@ class RuntimeSchedulerTest : public testing::Test {
 
     runtimeScheduler_ =
         std::make_unique<RuntimeScheduler>(runtimeExecutor, stubNow);
+    runtimeScheduler_->setEnableYielding(true);
   }
 
   jsi::Function createHostFunctionFromLambda(
@@ -317,6 +318,7 @@ TEST_F(RuntimeSchedulerTest, getCurrentPriorityLevel) {
 }
 
 TEST_F(RuntimeSchedulerTest, scheduleWork) {
+  runtimeScheduler_->setEnableYielding(false);
   bool wasCalled = false;
   runtimeScheduler_->scheduleWork(
       [&](jsi::Runtime const &) { wasCalled = true; });
@@ -334,7 +336,6 @@ TEST_F(RuntimeSchedulerTest, scheduleWork) {
 }
 
 TEST_F(RuntimeSchedulerTest, scheduleWorkWithYielding) {
-  runtimeScheduler_->setEnableYielding(true);
   bool wasCalled = false;
   runtimeScheduler_->scheduleWork(
       [&](jsi::Runtime const &) { wasCalled = true; });
@@ -353,8 +354,6 @@ TEST_F(RuntimeSchedulerTest, scheduleWorkWithYielding) {
 }
 
 TEST_F(RuntimeSchedulerTest, normalTaskYieldsToPlatformEvent) {
-  runtimeScheduler_->setEnableYielding(true);
-
   bool didRunJavaScriptTask = false;
   bool didRunPlatformWork = false;
 
@@ -382,8 +381,6 @@ TEST_F(RuntimeSchedulerTest, normalTaskYieldsToPlatformEvent) {
 }
 
 TEST_F(RuntimeSchedulerTest, expiredTaskDoesntYieldToPlatformEvent) {
-  runtimeScheduler_->setEnableYielding(true);
-
   bool didRunJavaScriptTask = false;
   bool didRunPlatformWork = false;
 
@@ -412,8 +409,6 @@ TEST_F(RuntimeSchedulerTest, expiredTaskDoesntYieldToPlatformEvent) {
 }
 
 TEST_F(RuntimeSchedulerTest, immediateTaskDoesntYieldToPlatformEvent) {
-  runtimeScheduler_->setEnableYielding(true);
-
   bool didRunJavaScriptTask = false;
   bool didRunPlatformWork = false;
 
@@ -603,4 +598,42 @@ TEST_F(RuntimeSchedulerTest, sameThreadTaskCreatesLowPriorityTask) {
   EXPECT_EQ(stubQueue_->size(), 0);
 }
 
+TEST_F(RuntimeSchedulerTest, twoThreadsRequestAccessToTheRuntime) {
+  bool didRunSynchronousTask = false;
+  bool didRunWork = false;
+
+  runtimeScheduler_->scheduleWork(
+      [&didRunWork](jsi::Runtime &) { didRunWork = true; });
+
+  std::thread t1([this, &didRunSynchronousTask]() {
+    runtimeScheduler_->executeNowOnTheSameThread(
+        [&didRunSynchronousTask](jsi::Runtime &runtime) {
+          didRunSynchronousTask = true;
+        });
+  });
+
+  auto hasTask = stubQueue_->waitForTasks(2, 1ms);
+
+  EXPECT_TRUE(hasTask);
+  EXPECT_FALSE(didRunWork);
+  EXPECT_FALSE(didRunSynchronousTask);
+  EXPECT_TRUE(runtimeScheduler_->getShouldYield());
+  EXPECT_EQ(stubQueue_->size(), 2);
+
+  stubQueue_->tick();
+
+  EXPECT_TRUE(didRunWork);
+  EXPECT_FALSE(didRunSynchronousTask);
+  EXPECT_TRUE(runtimeScheduler_->getShouldYield());
+  EXPECT_EQ(stubQueue_->size(), 1);
+
+  stubQueue_->tick();
+
+  t1.join();
+
+  EXPECT_TRUE(didRunWork);
+  EXPECT_TRUE(didRunSynchronousTask);
+  EXPECT_FALSE(runtimeScheduler_->getShouldYield());
+}
+
 } // namespace facebook::react

ReactCommon/react/renderer/runtimescheduler/tests/StubQueue.h

@@ -53,6 +53,15 @@ class StubQueue {
         lock, timeout, [this]() { return !callbackQueue_.empty(); });
   }
 
+  bool waitForTasks(
+      std::size_t numberOfTasks,
+      std::chrono::duration<double> timeout) const {
+    std::unique_lock<std::mutex> lock(mutex_);
+    return signal_.wait_for(lock, timeout, [this, numberOfTasks]() {
+      return numberOfTasks == callbackQueue_.size();
+    });
+  }
+
  private:
   mutable std::condition_variable signal_;
   mutable std::mutex mutex_;