feat(Worklets): runOnRuntimeAsync

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

@@ -226,6 +226,8 @@ std::vector<jsi::PropNameID> JSIWorkletsModuleProxy::getPropertyNames(
       jsi::PropNameID::forAscii(rt, "synchronizableLock"));
   propertyNames.emplace_back(
       jsi::PropNameID::forAscii(rt, "synchronizableUnlock"));
+  propertyNames.emplace_back(
+      jsi::PropNameID::forAscii(rt, "runOnRuntimeAsync"));
 
 #ifdef WORKLETS_BUNDLE_MODE
   propertyNames.emplace_back(
@@ -467,6 +469,22 @@ jsi::Value JSIWorkletsModuleProxy::get(
         });
   }
 
+  if (name == "runOnRuntimeAsync") {
+    return jsi::Function::createFromHostFunction(
+        rt,
+        propName,
+        4,
+        [jsScheduler = jsScheduler_](
+            jsi::Runtime &rt,
+            const jsi::Value &thisValue,
+            const jsi::Value *args,
+            size_t count) {
+          runOnRuntimeAsync(
+              rt, jsScheduler, args[0], args[1], args[2], args[3]);
+          return jsi::Value::undefined();
+        });
+  }
+
   if (name == "scheduleOnUI") {
     return jsi::Function::createFromHostFunction(
         rt,

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

@@ -0,0 +1,36 @@
+#include <worklets/Tools/Promise.h>
+
+namespace worklets {
+
+void Promise::resolve(const jsi::Value &result) {
+  auto resolve = resolve_->toJSValue(rt_).asObject(rt_).asFunction(rt_);
+  resolve.call(rt_, result);
+}
+
+void Promise::reject(
+    const std::string &message,
+    const std::string &stack,
+    const std::string &name,
+    const std::string &jsEngine) {
+  const auto reject = reject_->toJSValue(rt_).asObject(rt_).asFunction(rt_);
+  const auto errorInstance = rt_.global()
+                                 .getPropertyAsFunction(rt_, "Error")
+                                 .callAsConstructor(rt_)
+                                 .asObject(rt_);
+
+  errorInstance.setProperty(
+      rt_, "message", jsi::String::createFromUtf8(rt_, message));
+
+  errorInstance.setProperty(
+      rt_, "stack", jsi::String::createFromUtf8(rt_, stack));
+
+  errorInstance.setProperty(
+      rt_, "name", jsi::String::createFromUtf8(rt_, name));
+
+  errorInstance.setProperty(
+      rt_, "jsEngine", jsi::String::createFromUtf8(rt_, jsEngine));
+
+  reject.call(rt_, errorInstance);
+}
+
+} // namespace worklets

packages/react-native-worklets/Common/cpp/worklets/Tools/Promise.h

@@ -0,0 +1,27 @@
+#include <jsi/jsi.h>
+#include <worklets/SharedItems/Serializable.h>
+
+using namespace facebook;
+
+namespace worklets {
+
+struct Promise {
+  Promise(
+      jsi::Runtime &rt,
+      std::shared_ptr<Serializable> resolve,
+      std::shared_ptr<Serializable> reject)
+      : resolve_(resolve), reject_(reject), rt_(rt) {}
+
+  void resolve(const jsi::Value &result);
+  void reject(
+      const std::string &message,
+      const std::string &stack,
+      const std::string &name,
+      const std::string &jsEngine);
+
+  std::shared_ptr<Serializable> resolve_;
+  std::shared_ptr<Serializable> reject_;
+  jsi::Runtime &rt_;
+};
+
+} // namespace worklets

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

@@ -22,7 +22,6 @@ std::shared_ptr<WorkletRuntime> RuntimeManager::getRuntime(
   return nullptr;
 }
 
-#ifdef WORKLETS_BUNDLE_MODE
 std::shared_ptr<WorkletRuntime> RuntimeManager::getRuntime(
     jsi::Runtime *runtime) {
   std::shared_lock lock(weakRuntimesMutex_);
@@ -31,7 +30,6 @@ std::shared_ptr<WorkletRuntime> RuntimeManager::getRuntime(
   }
   return nullptr;
 }
-#endif // WORKLETS_BUNDLE_MODE
 
 std::vector<std::shared_ptr<WorkletRuntime>> RuntimeManager::getAllRuntimes() {
   std::shared_lock lock(weakRuntimesMutex_);
@@ -97,9 +95,7 @@ void RuntimeManager::registerRuntime(
   std::unique_lock lock(weakRuntimesMutex_);
   weakRuntimes_[runtimeId] = workletRuntime;
   nameToRuntimeId_[name] = runtimeId;
-#ifdef WORKLETS_BUNDLE_MODE
   runtimeAddressToRuntimeId_[&workletRuntime->getJSIRuntime()] = runtimeId;
-#endif // WORKLETS_BUNDLE_MODE
 }
 
 } // namespace worklets

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

@@ -24,9 +24,7 @@ class RuntimeManager {
  public:
   std::shared_ptr<WorkletRuntime> getRuntime(uint64_t runtimeId);
   std::shared_ptr<WorkletRuntime> getRuntime(const std::string &name);
-#ifdef WORKLETS_BUNDLE_MODE
   std::shared_ptr<WorkletRuntime> getRuntime(jsi::Runtime *runtime);
-#endif // WORKLETS_BUNDLE_MODE
 
   std::vector<std::shared_ptr<WorkletRuntime>> getAllRuntimes();
 
@@ -55,9 +53,7 @@ class RuntimeManager {
   std::map<uint64_t, std::weak_ptr<WorkletRuntime>> weakRuntimes_;
   std::shared_mutex weakRuntimesMutex_;
   std::map<std::string, uint64_t> nameToRuntimeId_;
-#ifdef WORKLETS_BUNDLE_MODE
   std::map<jsi::Runtime *, uint64_t> runtimeAddressToRuntimeId_;
-#endif // WORKLETS_BUNDLE_MODE
 };
 
 } // namespace worklets

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

@@ -1,9 +1,11 @@
 #include <worklets/NativeModules/JSIWorkletsModuleProxy.h>
 #include <worklets/Resources/Unpackers.h>
+#include <worklets/SharedItems/Serializable.h>
 #include <worklets/Tools/Defs.h>
 #include <worklets/Tools/JSISerializer.h>
 #include <worklets/Tools/JSLogger.h>
 #include <worklets/Tools/WorkletsJSIUtils.h>
+#include <worklets/WorkletRuntime/RuntimeManager.h>
 #include <worklets/WorkletRuntime/WorkletRuntime.h>
 #include <worklets/WorkletRuntime/WorkletRuntimeCollector.h>
 #include <worklets/WorkletRuntime/WorkletRuntimeDecorator.h>
@@ -239,4 +241,62 @@ void scheduleOnRuntime(
   workletRuntime->runAsyncGuarded(serializableWorklet);
 }
 
+void runOnRuntimeAsync(
+    jsi::Runtime &rt,
+    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const jsi::Value &workletRuntimeValue,
+    const jsi::Value &serializableWorkletValue,
+    const jsi::Value &serializableResolveValue,
+    const jsi::Value &serializableRejectValue) {
+  auto workletRuntime = extractWorkletRuntime(rt, workletRuntimeValue);
+  auto serializableWorklet = extractSerializableOrThrow<SerializableWorklet>(
+      rt,
+      serializableWorkletValue,
+      "[Worklets] Function `worklet` passed to `runOnRuntimeAsync` is not a serializable worklet.");
+  auto serializableResolve = extractSerializableOrThrow<
+      SerializableRemoteFunction>(
+      rt,
+      serializableResolveValue,
+      "[Worklets] Function `resolve` passed to `runOnRuntimeAsync` is not a serializable function.");
+  auto serializableReject = extractSerializableOrThrow<
+      SerializableRemoteFunction>(
+      rt,
+      serializableRejectValue,
+      "[Worklets] Function `reject` passed to `runOnRuntimeAsync` is not a serializable function.");
+
+  auto promise = std::make_shared<worklets::Promise>(
+      rt, serializableResolve, serializableReject);
+
+  auto eventLoop = workletRuntime->getEventLoop();
+
+  if (!eventLoop) {
+    throw std::runtime_error(
+        "[Worklets] Event Loop is not enabled for the (" +
+        workletRuntime->getRuntimeName() +
+        ") Worklet Runtime. Please enable it to use `runOnRuntimeAsync`.");
+  }
+
+  eventLoop->pushTask([serializableWorklet, &jsScheduler, promise](
+                          jsi::Runtime &workletRt) {
+    try {
+      auto result = serializableWorklet->toJSValue(workletRt)
+                        .asObject(workletRt)
+                        .asFunction(workletRt)
+                        .call(workletRt);
+      auto serializableResult = extractSerializableOrThrow(workletRt, result);
+      jsScheduler->scheduleOnJS(
+          [serializableResult, promise](jsi::Runtime &rnRt) {
+            auto resolveResult = serializableResult->toJSValue(rnRt);
+            promise->resolve(resolveResult);
+          });
+    } catch (const jsi::JSError &error) {
+      const auto &message = error.getMessage();
+      const auto &stack = error.getStack();
+      jsScheduler->scheduleOnJS([promise, message, stack](jsi::Runtime &rnRt) {
+        promise->reject(message, stack, "WorkletsError", "Worklets");
+      });
+    }
+  });
+}
+
 } // namespace worklets

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

@@ -10,6 +10,7 @@
 #include <worklets/RunLoop/EventLoop.h>
 #include <worklets/SharedItems/Serializable.h>
 #include <worklets/Tools/JSScheduler.h>
+#include <worklets/Tools/Promise.h>
 #include <worklets/WorkletRuntime/RuntimeData.h>
 
 #include <memory>
@@ -77,6 +78,10 @@ class WorkletRuntime : public jsi::HostObject,
     return name_;
   }
 
+  [[nodiscard]] auto getEventLoop() const -> std::shared_ptr<EventLoop> {
+    return eventLoop_;
+  }
+
  private:
   const uint64_t runtimeId_;
   const std::shared_ptr<std::recursive_mutex> runtimeMutex_;
@@ -97,4 +102,11 @@ void scheduleOnRuntime(
     const jsi::Value &workletRuntimeValue,
     const jsi::Value &serializableWorkletValue);
 
+void runOnRuntimeAsync(
+    jsi::Runtime &rt,
+    const std::shared_ptr<JSScheduler> &jsScheduler,
+    const jsi::Value &workletRuntimeValue,
+    const jsi::Value &serializableWorkletValue,
+    const jsi::Value &serializableResolveValue,
+    const jsi::Value &serializableRejectValue);
 } // namespace worklets

packages/react-native-worklets/src/WorkletsModule/JSWorklets.ts

@@ -79,6 +79,12 @@ class JSWorklets implements IWorkletsModule {
     );
   }
 
+  runOnRuntimeAsync(): never {
+    throw new WorkletsError(
+      'runOnRuntimeAsync should never be called in JSWorklets.'
+    );
+  }
+
   createSerializableSet(): never {
     throw new WorkletsError(
       'createSerializableSet should never be called in JSWorklets.'

packages/react-native-worklets/src/WorkletsModule/NativeWorklets.ts

@@ -160,6 +160,20 @@ See https://docs.swmansion.com/react-native-worklets/docs/guides/troubleshooting
     return this.#workletsModuleProxy.scheduleOnUI(serializable);
   }
 
+  runOnRuntimeAsync<TValue, TReturn>(
+    workletRuntime: WorkletRuntime,
+    worklet: SerializableRef<TValue>,
+    resolve: SerializableRef<(value: TReturn) => void>,
+    reject: SerializableRef<(error: Error) => void>
+  ): void {
+    return this.#workletsModuleProxy.runOnRuntimeAsync(
+      workletRuntime,
+      worklet,
+      resolve,
+      reject
+    );
+  }
+
   executeOnUIRuntimeSync<TValue, TReturn>(
     serializable: SerializableRef<TValue>
   ): TReturn {

packages/react-native-worklets/src/WorkletsModule/workletsModuleProxy.ts

@@ -71,6 +71,13 @@ export interface WorkletsModuleProxy {
 
   scheduleOnUI<TValue>(serializable: SerializableRef<TValue>): void;
 
+  runOnRuntimeAsync<TValue, TReturn>(
+    workletRuntime: WorkletRuntime,
+    worklet: SerializableRef<TValue>,
+    resolve: SerializableRef<(value: TReturn) => void>,
+    reject: SerializableRef<(error: Error) => void>
+  ): void;
+
   executeOnUIRuntimeSync<TValue, TReturn>(
     serializable: SerializableRef<TValue>
   ): TReturn;

packages/react-native-worklets/src/index.ts

@@ -21,6 +21,7 @@ export { getRuntimeKind, RuntimeKind } from './runtimeKind';
 export {
   createWorkletRuntime,
   runOnRuntime,
+  runOnRuntimeAsync,
   scheduleOnRuntime,
 } from './runtimes';
 export { createSerializable, isSerializableRef } from './serializable';

packages/react-native-worklets/src/runtimes.ts

@@ -170,6 +170,32 @@ export function scheduleOnRuntime<Args extends unknown[], ReturnValue>(
   runOnRuntime(workletRuntime, worklet)(...args);
 }
 
+export function runOnRuntimeAsync<Args extends unknown[], ReturnValue>(
+  workletRuntime: WorkletRuntime,
+  worklet: (...args: Args) => ReturnValue,
+  ...args: Args
+): Promise<ReturnValue> {
+  'worklet';
+  if (globalThis.__RUNTIME_KIND !== RuntimeKind.ReactNative) {
+    throw new WorkletsError(
+      '`runOnRuntimeAsync` cannot be called on the Worker Runtime.'
+    );
+  }
+
+  return new Promise<ReturnValue>((resolve, reject) => {
+    WorkletsModule.runOnRuntimeAsync(
+      workletRuntime,
+      createSerializable(() => {
+        'worklet';
+        const result = worklet(...args);
+        return makeShareableCloneOnUIRecursive(result);
+      }),
+      createSerializable(resolve),
+      createSerializable(reject)
+    );
+  });
+}
+
 /** Configuration object for creating a worklet runtime. */
 export type WorkletRuntimeConfig = {
   /** The name of the worklet runtime. */

packages/react-native-worklets/src/threads.ts

@@ -427,7 +427,7 @@ function flushUIQueue(): void {
         queue.forEach(([workletFunction, workletArgs, jobResolve]) => {
           const result = workletFunction(...workletArgs);
           if (jobResolve) {
-            runOnJS(jobResolve)(result);
+            scheduleOnRN(jobResolve, result);
           }
         });
         callMicrotasks();