Implement task continuation

Summary:
Changelog: [internal]

Scheduler's callback have option to add more work inside callback. This work stays on top of the priority queue and gives React ability to flush all work synchronously if need.

This diff adds use of `shouldYield_` to the workLoop. For now, it always evaluates to false. In the future when we allow access to the scheduler to native, it will allow yielding.

Relevant code in JavaScript:
https://github.com/facebook/react/blob/master/packages/scheduler/src/forks/SchedulerNoDOM.js#L190

Reviewed By: mdvacca

Differential Revision: D27823528

fbshipit-source-id: 016101e41eb7c41c2ac5abb55f803814867b8517

Author: sammy-SC

ReactCommon/react/renderer/runtimescheduler/RuntimeScheduler.cpp

@@ -32,8 +32,23 @@ std::shared_ptr<Task> RuntimeScheduler::scheduleTask(
 
       while (!taskQueue_.empty()) {
         auto topPriority = taskQueue_.top();
-        taskQueue_.pop();
-        (*topPriority)(runtime);
+        auto now = now_();
+        auto didUserCallbackTimeout = topPriority->expirationTime <= now;
+
+        if (!didUserCallbackTimeout && shouldYield_) {
+          // This task hasn't expired and we need to yield.
+          break;
+        }
+
+        auto result = topPriority->execute(runtime);
+
+        if (result.isObject() &&
+            result.getObject(runtime).isFunction(runtime)) {
+          topPriority->callback =
+              result.getObject(runtime).getFunction(runtime);
+        } else {
+          taskQueue_.pop();
+        }
       }
     });
   }
@@ -42,7 +57,7 @@ std::shared_ptr<Task> RuntimeScheduler::scheduleTask(
 }
 
 void RuntimeScheduler::cancelTask(const std::shared_ptr<Task> &task) {
-  task->cancel();
+  task->callback.reset();
 }
 
 bool RuntimeScheduler::getShouldYield() const {

ReactCommon/react/renderer/runtimescheduler/Task.cpp

@@ -13,33 +13,19 @@ Task::Task(
     SchedulerPriority priority,
     jsi::Function callback,
     std::chrono::steady_clock::time_point expirationTime)
-    : priority_(priority),
-      callback_(std::move(callback)),
-      expirationTime_(expirationTime) {}
+    : priority(priority),
+      callback(std::move(callback)),
+      expirationTime(expirationTime) {}
 
-SchedulerPriority Task::getPriority() const {
-  return priority_;
-}
-
-RuntimeSchedulerClock::time_point Task::getExpirationTime() const {
-  return expirationTime_;
-}
-
-void Task::cancel() {
-  // Null out the callback to indicate the task has been canceled. (Can't
-  // remove from the priority_queue because you can't remove arbitrary nodes
-  // from an array based heap, only the first one.)
-  callback_.reset();
-}
-
-void Task::operator()(jsi::Runtime &runtime) const {
+jsi::Value Task::execute(jsi::Runtime &runtime) const {
   // Cancelled task doesn't have a callback.
-  if (callback_) {
+  if (callback) {
     // Callback in JavaScript is expecting a single bool parameter.
     // React team plans to remove it and it is safe to pass in
     // hardcoded false value.
-    callback_.value().call(runtime, {false});
+    return callback.value().call(runtime, {false});
   }
+  return jsi::Value::undefined();
 }
 
 } // namespace facebook::react

ReactCommon/react/renderer/runtimescheduler/Task.h

@@ -14,33 +14,32 @@
 
 namespace facebook::react {
 
-class Task final {
- public:
+class RuntimeScheduler;
+class TaskPriorityComparer;
+
+struct Task final {
   Task(
       SchedulerPriority priority,
       jsi::Function callback,
       std::chrono::steady_clock::time_point expirationTime);
 
-  SchedulerPriority getPriority() const;
-
-  RuntimeSchedulerClock::time_point getExpirationTime() const;
-
-  void cancel();
+ private:
+  friend RuntimeScheduler;
+  friend TaskPriorityComparer;
 
-  void operator()(jsi::Runtime &runtime) const;
+  SchedulerPriority priority;
+  better::optional<jsi::Function> callback;
+  RuntimeSchedulerClock::time_point expirationTime;
 
- private:
-  SchedulerPriority priority_;
-  better::optional<jsi::Function> callback_;
-  RuntimeSchedulerClock::time_point expirationTime_;
+  jsi::Value execute(jsi::Runtime &runtime) const;
 };
 
 class TaskPriorityComparer {
  public:
   inline bool operator()(
       std::shared_ptr<Task> const &lhs,
       std::shared_ptr<Task> const &rhs) {
-    return lhs->getExpirationTime() > rhs->getExpirationTime();
+    return lhs->expirationTime > rhs->expirationTime;
   }
 };
 

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

@@ -43,7 +43,7 @@ class RuntimeSchedulerTest : public testing::Test {
   }
 
   jsi::Function createHostFunctionFromLambda(
-      std::function<void(bool)> callback) {
+      std::function<jsi::Value(bool)> callback) {
     return jsi::Function::createFromHostFunction(
         *runtime_,
         jsi::PropNameID::forUtf8(*runtime_, ""),
@@ -55,8 +55,7 @@ class RuntimeSchedulerTest : public testing::Test {
             size_t) noexcept -> jsi::Value {
           ++hostFunctionCallCount_;
           auto didUserCallbackTimeout = arguments[0].getBool();
-          callback(didUserCallbackTimeout);
-          return jsi::Value::undefined();
+          return callback(didUserCallbackTimeout);
         });
   }
 
@@ -93,6 +92,7 @@ TEST_F(RuntimeSchedulerTest, scheduleSingleTask) {
       createHostFunctionFromLambda([&didRunTask](bool didUserCallbackTimeout) {
         didRunTask = true;
         EXPECT_FALSE(didUserCallbackTimeout);
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -113,6 +113,7 @@ TEST_F(RuntimeSchedulerTest, scheduleImmediatePriorityTask) {
       createHostFunctionFromLambda([&didRunTask](bool didUserCallbackTimeout) {
         didRunTask = true;
         EXPECT_FALSE(didUserCallbackTimeout);
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -135,6 +136,7 @@ TEST_F(RuntimeSchedulerTest, taskExpiration) {
         // Task has timed out but the parameter is deprecated and `false` is
         // hardcoded.
         EXPECT_FALSE(didUserCallbackTimeout);
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -157,6 +159,7 @@ TEST_F(RuntimeSchedulerTest, scheduleTwoTasksWithSamePriority) {
   auto callbackOne =
       createHostFunctionFromLambda([this, &firstTaskCallOrder](bool) {
         firstTaskCallOrder = hostFunctionCallCount_;
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -166,6 +169,7 @@ TEST_F(RuntimeSchedulerTest, scheduleTwoTasksWithSamePriority) {
   auto callbackTwo =
       createHostFunctionFromLambda([this, &secondTaskCallOrder](bool) {
         secondTaskCallOrder = hostFunctionCallCount_;
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -188,6 +192,7 @@ TEST_F(RuntimeSchedulerTest, scheduleTwoTasksWithDifferentPriorities) {
   auto callbackOne =
       createHostFunctionFromLambda([this, &lowPriorityTaskCallOrder](bool) {
         lowPriorityTaskCallOrder = hostFunctionCallCount_;
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -197,6 +202,7 @@ TEST_F(RuntimeSchedulerTest, scheduleTwoTasksWithDifferentPriorities) {
   auto callbackTwo = createHostFunctionFromLambda(
       [this, &userBlockingPriorityTaskCallOrder](bool) {
         userBlockingPriorityTaskCallOrder = hostFunctionCallCount_;
+        return jsi::Value::undefined();
       });
 
   runtimeScheduler_->scheduleTask(
@@ -216,8 +222,10 @@ TEST_F(RuntimeSchedulerTest, scheduleTwoTasksWithDifferentPriorities) {
 
 TEST_F(RuntimeSchedulerTest, cancelTask) {
   bool didRunTask = false;
-  auto callback =
-      createHostFunctionFromLambda([&didRunTask](bool) { didRunTask = true; });
+  auto callback = createHostFunctionFromLambda([&didRunTask](bool) {
+    didRunTask = true;
+    return jsi::Value::undefined();
+  });
 
   auto task = runtimeScheduler_->scheduleTask(
       SchedulerPriority::NormalPriority, std::move(callback));
@@ -233,4 +241,36 @@ TEST_F(RuntimeSchedulerTest, cancelTask) {
   EXPECT_EQ(stubQueue_->size(), 0);
 }
 
+TEST_F(RuntimeSchedulerTest, continuationTask) {
+  bool didRunTask = false;
+  bool didContinuationTask = false;
+
+  auto callback = createHostFunctionFromLambda([&](bool) {
+    didRunTask = true;
+    return jsi::Function::createFromHostFunction(
+        *runtime_,
+        jsi::PropNameID::forUtf8(*runtime_, ""),
+        1,
+        [&](jsi::Runtime &runtime,
+            jsi::Value const &,
+            jsi::Value const *arguments,
+            size_t) noexcept -> jsi::Value {
+          didContinuationTask = true;
+          return jsi::Value::undefined();
+        });
+  });
+
+  auto task = runtimeScheduler_->scheduleTask(
+      SchedulerPriority::NormalPriority, std::move(callback));
+
+  EXPECT_FALSE(didRunTask);
+  EXPECT_EQ(stubQueue_->size(), 1);
+
+  stubQueue_->tick();
+
+  EXPECT_TRUE(didRunTask);
+  EXPECT_TRUE(didContinuationTask);
+  EXPECT_EQ(stubQueue_->size(), 0);
+}
+
 } // namespace facebook::react