feat(Worklets): Worklet Runtime C++ scheduling API

packages/react-native-reanimated/Common/cpp/reanimated/AnimatedSensor/AnimatedSensorModule.cpp

@@ -57,7 +57,7 @@ jsi::Value AnimatedSensorModule::registerSensor(
         }
         value.setProperty(uiRuntime, "interfaceOrientation", orientationDegrees);
 
-        uiWorkletRuntime->runGuarded(serializableHandler, value);
+        uiWorkletRuntime->runSync(serializableHandler, value);
       });
   if (sensorId != -1) {
     sensorsIds_.insert(sensorId);

packages/react-native-reanimated/Common/cpp/reanimated/NativeModules/ReanimatedModuleProxy.cpp

@@ -1273,7 +1273,7 @@ jsi::Value ReanimatedModuleProxy::subscribeForKeyboardEvents(
         if (!strongThis) {
           return;
         }
-        strongThis->workletsModuleProxy_->getUIWorkletRuntime()->runGuarded(
+        strongThis->workletsModuleProxy_->getUIWorkletRuntime()->runSync(
             serializableHandler, jsi::Value(keyboardState), jsi::Value(height));
       },
       isStatusBarTranslucent.getBool(),

packages/react-native-worklets/Common/cpp/worklets/NativeModules/JSIWorkletsModuleProxy.cpp

@@ -50,7 +50,7 @@ inline void scheduleOnUI(
     const auto scope = jsi::Scope(uiWorkletRuntime->getJSIRuntime());
 #endif // JS_RUNTIME_HERMES
 
-    uiWorkletRuntime->runGuarded(serializableWorklet);
+    uiWorkletRuntime->runSync(serializableWorklet);
   });
 }
 
@@ -59,7 +59,10 @@ inline jsi::Value executeOnUIRuntimeSync(
     jsi::Runtime &rt,
     const jsi::Value &worklet) {
   if (auto uiWorkletRuntime = weakUIWorkletRuntime.lock()) {
-    return uiWorkletRuntime->executeSync(rt, worklet);
+    auto serializableWorklet = extractSerializableOrThrow<SerializableWorklet>(
+        rt, worklet, "[Worklets] Only worklets can be executed on UI runtime.");
+    auto serializedResult = uiWorkletRuntime->runSyncSerialized(serializableWorklet);
+    return serializedResult->toJSValue(rt);
   }
   return jsi::Value::undefined();
 }
@@ -85,10 +88,9 @@ inline jsi::Value propagateModuleUpdate(
   const auto runtimes = runtimeManager->getAllRuntimes();
 
   for (auto runtime : runtimes) {
-    runtime->executeSync([code, sourceUrl](jsi::Runtime &rt) {
+    runtime->runSync([code, sourceUrl](jsi::Runtime &rt) -> void {
       const auto buffer = std::make_shared<jsi::StringBuffer>(code);
       rt.evaluateJavaScript(buffer, sourceUrl);
-      return jsi::Value::undefined();
     });
   }
   return jsi::Value::undefined();

packages/react-native-worklets/Common/cpp/worklets/SharedItems/Serializable.h

@@ -19,20 +19,26 @@ jsi::Function getValueUnpacker(jsi::Runtime &rt);
 jsi::Function getCallGuard(jsi::Runtime &rt);
 #endif // NDEBUG
 
-// If possible, please use `WorkletRuntime::runGuarded` instead.
+/** If possible, please use `WorkletRuntime::runSync` instead. */
 template <typename... Args>
-inline jsi::Value runOnRuntimeGuarded(jsi::Runtime &rt, const jsi::Value &function, Args &&...args) {
+inline jsi::Value runOnRuntimeGuarded(jsi::Runtime &rt, const jsi::Function &function, Args &&...args) {
   // We only use callGuard in debug mode, otherwise we call the provided
   // function directly. CallGuard provides a way of capturing exceptions in
   // JavaScript and propagating them to the main React Native thread such that
   // they can be presented using RN's LogBox.
 #ifndef NDEBUG
   return getCallGuard(rt).call(rt, function, args...);
 #else
-  return function.asObject(rt).asFunction(rt).call(rt, args...);
+  return function..call(rt, args...);
 #endif // NDEBUG
 }
 
+/** If possible, please use `WorkletRuntime::runSync` instead. */
+template <typename... Args>
+inline jsi::Value runOnRuntimeGuarded(jsi::Runtime &rt, const jsi::Value &function, Args &&...args) {
+  return runOnRuntimeGuarded(rt, function.asObject(rt).asFunction(rt), std::forward<Args>(args)...);
+}
+
 inline void cleanupIfRuntimeExists(jsi::Runtime *rt, std::unique_ptr<jsi::Value> &value) {
   if (rt != nullptr && !WorkletRuntimeRegistry::isRuntimeAlive(rt)) {
     // The below use of unique_ptr.release prevents the smart pointer from

packages/react-native-worklets/Common/cpp/worklets/Tools/WorkletEventHandler.cpp

@@ -9,7 +9,7 @@ void WorkletEventHandler::process(
     const std::shared_ptr<WorkletRuntime> &workletRuntime,
     const double eventTimestamp,
     const jsi::Value &eventValue) const {
-  workletRuntime->runGuarded(handlerFunction_, jsi::Value(eventTimestamp), eventValue);
+  workletRuntime->runSync(handlerFunction_, jsi::Value(eventTimestamp), eventValue);
 }
 
 uint64_t WorkletEventHandler::getHandlerId() const {

packages/react-native-worklets/Common/cpp/worklets/WorkletRuntime/RuntimeManager.cpp

@@ -49,7 +49,7 @@ std::shared_ptr<WorkletRuntime> RuntimeManager::createWorkletRuntime(
   workletRuntime->init(std::move(jsiWorkletsModuleProxy));
 
   if (initializer) {
-    workletRuntime->runGuarded(initializer);
+    workletRuntime->runSync(initializer);
   }
 
   registerRuntime(runtimeId, workletRuntime);

packages/react-native-worklets/Common/cpp/worklets/WorkletRuntime/WorkletRuntime.cpp

@@ -135,44 +135,68 @@ void WorkletRuntime::init(std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsMod
 #endif // WORKLETS_BUNDLE_MODE
 }
 
-void WorkletRuntime::runAsyncGuarded(const std::shared_ptr<SerializableWorklet> &worklet) {
-  react_native_assert(
-      "[Worklets] Tried to invoke `runAsyncGuarded` on a Worklet Runtime but "
-      "the async queue is not set. Recreate the runtime with a valid async queue.");
+/* #region schedule */
 
-  queue_->push([worklet, weakThis = weak_from_this()] {
+void WorkletRuntime::schedule(jsi::Function &&function) const {
+  react_native_assert(
+      queue_ &&
+      "[Worklets] Tried to invoke `schedule` on a Worklet Runtime but the "
+      "async queue is not set. Recreate the runtime with a valid async queue.");
+  queue_->push([function = std::make_shared<jsi::Function>(std::move(function)), weakThis = weak_from_this()]() {
     auto strongThis = weakThis.lock();
     if (!strongThis) {
       return;
     }
 
-    strongThis->runGuarded(worklet);
+    strongThis->runSync(*function);
   });
 }
 
-jsi::Value WorkletRuntime::executeSync(jsi::Runtime &rt, const jsi::Value &worklet) const {
-  auto serializableWorklet = extractSerializableOrThrow<SerializableWorklet>(
-      rt, worklet, "[Worklets] Only worklets can be executed synchronously on UI runtime.");
-  auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
-  jsi::Runtime &uiRuntime = getJSIRuntime();
-  auto result = runGuarded(serializableWorklet);
-  auto serializableResult = extractSerializableOrThrow(uiRuntime, result);
-  lock.unlock();
-  return serializableResult->toJSValue(rt);
+void WorkletRuntime::schedule(std::shared_ptr<SerializableWorklet> worklet) const {
+  react_native_assert(
+      queue_ &&
+      "[Worklets] Tried to invoke `schedule` on a Worklet Runtime but the "
+      "async queue is not set. Recreate the runtime with a valid async queue.");
+
+  queue_->push([worklet = std::move(worklet), weakThis = weak_from_this()] {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    strongThis->runSync(worklet);
+  });
 }
 
-jsi::Value WorkletRuntime::executeSync(std::function<jsi::Value(jsi::Runtime &)> &&job) const {
-  auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
-  jsi::Runtime &uiRuntime = getJSIRuntime();
-  return job(uiRuntime);
+void WorkletRuntime::schedule(std::function<void()> job) const {
+  react_native_assert(
+      queue_ &&
+      "[Worklets] Tried to invoke `schedule` on a Worklet Runtime but the "
+      "async queue is not set. Recreate the runtime with a valid async queue.");
+
+  queue_->push(std::move(job));
 }
 
-jsi::Value WorkletRuntime::executeSync(const std::function<jsi::Value(jsi::Runtime &)> &job) const {
-  auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
-  jsi::Runtime &uiRuntime = getJSIRuntime();
-  return job(uiRuntime);
+void WorkletRuntime::schedule(std::function<void(jsi::Runtime &)> job) const {
+  react_native_assert(
+      queue_ &&
+      "[Worklets] Tried to invoke `schedule` on a Worklet Runtime but the "
+      "async queue is not set. Recreate the runtime with a valid async queue.");
+
+  queue_->push([job = std::move(job), weakThis = weak_from_this()]() {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    auto lock = std::unique_lock<std::recursive_mutex>(*strongThis->runtimeMutex_);
+    jsi::Runtime &runtime = strongThis->getJSIRuntime();
+    job(runtime);
+  });
 }
 
+/* #endregion */
+
 jsi::Value WorkletRuntime::get(jsi::Runtime &rt, const jsi::PropNameID &propName) {
   auto name = propName.utf8(rt);
   if (name == "toString") {
@@ -215,7 +239,7 @@ void scheduleOnRuntime(
       rt,
       serializableWorkletValue,
       "[Worklets] Function passed to `_scheduleOnRuntime` is not a serializable worklet.");
-  workletRuntime->runAsyncGuarded(serializableWorklet);
+  workletRuntime->schedule(serializableWorklet);
 }
 
 #if REACT_NATIVE_MINOR_VERSION >= 81
@@ -231,4 +255,27 @@ std::weak_ptr<WorkletRuntime> WorkletRuntime::getWeakRuntimeFromJSIRuntime(jsi::
 }
 #endif // REACT_NATIVE_MINOR_VERSION >= 81
 
+/* #region deprecated */
+
+void WorkletRuntime::runAsyncGuarded(const std::shared_ptr<SerializableWorklet> &worklet) {
+  schedule(worklet);
+}
+
+jsi::Value WorkletRuntime::executeSync(jsi::Runtime &caller, const jsi::Value &worklet) const {
+  auto serializableWorklet = extractSerializableOrThrow<SerializableWorklet>(
+      caller, worklet, "[Worklets] Only worklets can be executed synchronously on UI runtime.");
+  auto result = runSyncSerialized(serializableWorklet);
+  return result->toJSValue(caller);
+}
+
+jsi::Value WorkletRuntime::executeSync(std::function<jsi::Value(jsi::Runtime &)> &&job) const {
+  return runSync(job);
+}
+
+jsi::Value WorkletRuntime::executeSync(const std::function<jsi::Value(jsi::Runtime &)> &job) const {
+  return runSync(job);
+}
+
+/* #endregion */
+
 } // namespace worklets

packages/react-native-worklets/Common/cpp/worklets/WorkletRuntime/WorkletRuntime.h

@@ -14,6 +14,7 @@
 
 #include <memory>
 #include <string>
+#include <type_traits>
 #include <utility>
 #include <vector>
 
@@ -22,56 +23,134 @@ using namespace react;
 
 namespace worklets {
 
+template <typename TCallable>
+concept ImplicitlySerializableCallable = std::is_assignable_v<const jsi::Function &, TCallable> ||
+    std::is_assignable_v<const std::shared_ptr<SerializableWorklet> &, TCallable>;
+
+template <typename TCallable>
+concept RuntimeCallable = requires(TCallable &&callable, jsi::Runtime &rt) {
+  // NOLINTNEXTLINE(readability/braces) cpplint doesn't understand concepts
+  { callable(rt) };
+} || ImplicitlySerializableCallable<TCallable>;
+
 /**
  * Forward declaration to avoid circular dependencies.
  */
 class JSIWorkletsModuleProxy;
 
 class WorkletRuntime : public jsi::HostObject, public std::enable_shared_from_this<WorkletRuntime> {
  public:
-  explicit WorkletRuntime(
-      uint64_t runtimeId,
-      const std::shared_ptr<MessageQueueThread> &jsQueue,
-      const std::string &name,
-      const std::shared_ptr<AsyncQueue> &queue = nullptr,
-      bool enableEventLoop = true);
+  void schedule(jsi::Function &&function) const;
+  void schedule(std::shared_ptr<SerializableWorklet> worklet) const;
+  void schedule(std::function<void()> job) const;
+  void schedule(std::function<void(jsi::Runtime &)> job) const;
 
-  void init(std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy);
+  /* #region runSync */
 
-  jsi::Runtime &getJSIRuntime() const {
-    return *runtime_;
+  template <RuntimeCallable TCallable, typename... Args>
+  std::invoke_result_t<TCallable, Args...> inline runSync(TCallable &&callable, Args &&...args) const;
+  template <typename... Args>
+  jsi::Value inline runSync(const jsi::Function &function, Args &&...args) const {
+    auto &rt = *runtime_;
+    return runOnRuntimeGuarded(rt, function, std::forward<Args>(args)...);
   }
-
   template <typename... Args>
-  inline jsi::Value runGuarded(const std::shared_ptr<SerializableWorklet> &serializableWorklet, Args &&...args) const {
+  jsi::Value inline runSync(const std::shared_ptr<SerializableWorklet> &worklet, Args &&...args) const {
     jsi::Runtime &rt = *runtime_;
-    return runOnRuntimeGuarded(rt, serializableWorklet->toJSValue(rt), std::forward<Args>(args)...);
+    return runSync(worklet->toJSValue(rt).asObject(rt).asFunction(rt), std::forward<Args>(args)...);
+  }
+  template <RuntimeCallable TCallable>
+  std::invoke_result_t<TCallable, jsi::Runtime &> inline runSync(TCallable &&job) const {
+    jsi::Runtime &rt = getJSIRuntime();
+    auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
+    return job(rt);
   }
 
-  void runAsyncGuarded(const std::shared_ptr<SerializableWorklet> &worklet);
-
-  jsi::Value executeSync(jsi::Runtime &rt, const jsi::Value &worklet) const;
-
-  jsi::Value executeSync(std::function<jsi::Value(jsi::Runtime &)> &&job) const;
+  /* #endregion */
 
-  jsi::Value executeSync(const std::function<jsi::Value(jsi::Runtime &)> &job) const;
+  /* #region runSyncSerialized */
 
-  std::string toString() const {
-    return "[WorkletRuntime \"" + name_ + "\"]";
+  template <ImplicitlySerializableCallable TCallable, typename... Args>
+  std::shared_ptr<Serializable> inline runSyncSerialized(TCallable &&callable, Args &&...args) const;
+  template <typename... Args>
+  std::shared_ptr<Serializable> inline runSyncSerialized(const jsi::Function &function, Args &&...args) const {
+    jsi::Runtime &rt = getJSIRuntime();
+    auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
+    auto result = runSync(function, std::forward<Args>(args)...);
+    auto serializableResult = extractSerializableOrThrow(
+        rt,
+        result,
+        "[Worklets] Function passed to `runSyncSerialized`"
+        "must return a value serialized with `createSerializable`.");
+    return serializableResult;
+  }
+  template <typename... Args>
+  std::shared_ptr<Serializable> inline runSyncSerialized(
+      const std::shared_ptr<SerializableWorklet> &worklet,
+      Args &&...args) const {
+    jsi::Runtime &rt = getJSIRuntime();
+    auto lock = std::unique_lock<std::recursive_mutex>(*runtimeMutex_);
+    auto result = runSync(worklet, std::forward<Args>(args)...);
+    auto serializableResult = extractSerializableOrThrow(
+        rt,
+        result,
+        "[Worklets] Worklet passed to `runSyncSerialized`"
+        "must return a value serialized with `createSerializable`.");
+    return serializableResult;
   }
 
+  /* #endregion */
+
   jsi::Value get(jsi::Runtime &rt, const jsi::PropNameID &propName) override;
 
   std::vector<jsi::PropNameID> getPropertyNames(jsi::Runtime &rt) override;
 
-  [[nodiscard]] auto getRuntimeId() const -> uint64_t {
+  [[nodiscard]] inline std::string toString() const noexcept {
+    return "[WorkletRuntime \"" + name_ + "\"]";
+  }
+
+  [[nodiscard]] inline jsi::Runtime &getJSIRuntime() const noexcept {
+    return *runtime_;
+  }
+
+  [[nodiscard]] inline uint64_t getRuntimeId() const noexcept {
     return runtimeId_;
   }
 
-  [[nodiscard]] auto getRuntimeName() const -> std::string {
+  [[nodiscard]] inline std::string getRuntimeName() const noexcept {
     return name_;
   }
 
+  explicit WorkletRuntime(
+      uint64_t runtimeId,
+      const std::shared_ptr<MessageQueueThread> &jsQueue,
+      const std::string &name,
+      const std::shared_ptr<AsyncQueue> &queue = nullptr,
+      bool enableEventLoop = true);
+
+  void init(std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy);
+
+  /* #region deprecated */
+
+  /** @deprecated Use `runSync` instead. */
+  template <RuntimeCallable TCallable, typename... Args>
+  inline jsi::Value runGuarded(TCallable &&callable, Args &&...args) const {
+    return runSync(std::forward<TCallable>(callable), std::forward<Args>(args)...);
+  }
+
+  /** @deprecated Use `schedule` instead. */
+  void runAsyncGuarded(const std::shared_ptr<SerializableWorklet> &worklet);
+
+  /** @deprecated Use `runSyncSerialized` and extract to `jsi::Value` with
+   * `extractSerializableOrThrow` instead. */
+  jsi::Value executeSync(jsi::Runtime &rt, const jsi::Value &worklet) const;
+  /** @deprecated Use `runSync` instead. */
+  jsi::Value executeSync(std::function<jsi::Value(jsi::Runtime &)> &&job) const;
+  /** @deprecated Use `runSync` instead. */
+  jsi::Value executeSync(const std::function<jsi::Value(jsi::Runtime &)> &job) const;
+
+  /* #endregion */
+
 #if REACT_NATIVE_MINOR_VERSION >= 81
   /**
    * Retrieves a weak reference to the WorkletRuntime associated with the
@@ -99,5 +178,4 @@ void scheduleOnRuntime(
     jsi::Runtime &rt,
     const jsi::Value &workletRuntimeValue,
     const jsi::Value &serializableWorkletValue);
-
 } // namespace worklets