feat(Worklets): runAsyncGuardedOnUI

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

@@ -64,9 +64,10 @@ inline jsi::Value createWorkletRuntime(
     const std::shared_ptr<MessageQueueThread> &jsQueue,
     std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy,
     const std::string &name,
-    std::shared_ptr<SerializableWorklet> &initializer) {
+    std::shared_ptr<SerializableWorklet> &initializer,
+    const std::shared_ptr<AsyncQueue> &queue) {
   const auto workletRuntime = runtimeManager->createWorkletRuntime(
-      jsiWorkletsModuleProxy, name, initializer);
+      jsiWorkletsModuleProxy, name, initializer, queue);
   return jsi::Object::createFromHostObject(originRuntime, workletRuntime);
 }
 
@@ -104,6 +105,24 @@ inline jsi::Value reportFatalErrorOnJS(
   return jsi::Value::undefined();
 }
 
+inline std::shared_ptr<AsyncQueue> extractAsyncQueue(
+    jsi::Runtime &rt,
+    const jsi::Value &value) {
+  if (!value.isObject()) {
+    return nullptr;
+  }
+  const auto object = value.asObject(rt);
+
+  if (!object.hasNativeState(rt)) {
+    return nullptr;
+  }
+
+  const auto &nativeState = object.getNativeState(rt);
+  auto asyncQueue = std::dynamic_pointer_cast<AsyncQueue>(nativeState);
+
+  return asyncQueue;
+}
+
 JSIWorkletsModuleProxy::JSIWorkletsModuleProxy(
     const bool isDevBundle,
     const std::shared_ptr<const BigStringBuffer> &script,
@@ -459,24 +478,33 @@ jsi::Value JSIWorkletsModuleProxy::get(
     return jsi::Function::createFromHostFunction(
         rt,
         propName,
-        2,
+        4,
         [clone = std::make_shared<JSIWorkletsModuleProxy>(*this)](
             jsi::Runtime &rt,
             const jsi::Value &thisValue,
             const jsi::Value *args,
             size_t count) {
-          auto name = args[0].asString(rt).utf8(rt);
+          const auto name = args[0].asString(rt).utf8(rt);
           auto serializableInitializer =
               extractSerializableOrThrow<SerializableWorklet>(
                   rt, args[1], "[Worklets] Initializer must be a worklet.");
+          const auto useDefaultQueue = args[2].asBool();
+
+          std::shared_ptr<AsyncQueue> asyncQueue;
+          if (useDefaultQueue) {
+            asyncQueue = std::make_shared<AsyncQueueImpl>(name + "_queue");
+          } else {
+            asyncQueue = extractAsyncQueue(rt, args[3]);
+          }
 
           return createWorkletRuntime(
               rt,
               clone->getRuntimeManager(),
               clone->getJSQueue(),
               clone,
               name,
-              serializableInitializer);
+              serializableInitializer,
+              asyncQueue);
         });
   }
 

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

@@ -3,6 +3,7 @@
 
 #include <worklets/NativeModules/WorkletsModuleProxy.h>
 #include <worklets/SharedItems/Shareables.h>
+#include <worklets/Tools/AsyncQueueImpl.h>
 #include <worklets/Tools/Defs.h>
 #include <worklets/WorkletRuntime/UIRuntimeDecorator.h>
 
@@ -40,8 +41,9 @@ WorkletsModuleProxy::WorkletsModuleProxy(
       script_(script),
       sourceUrl_(sourceUrl),
       runtimeManager_(std::make_shared<RuntimeManager>()),
-      uiWorkletRuntime_(
-          runtimeManager_->createUninitializedUIRuntime(jsQueue_)) {
+      uiWorkletRuntime_(runtimeManager_->createUninitializedUIRuntime(
+          jsQueue_,
+          std::make_shared<AsyncQueueUI>(uiScheduler_))) {
   /**
    * We call additional `init` method here because
    * JSIWorkletsModuleProxy needs a weak_ptr to the UI Runtime.

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

@@ -50,4 +50,11 @@ void AsyncQueueImpl::push(std::function<void()> &&job) {
   state_->cv.notify_one();
 }
 
+AsyncQueueUI::AsyncQueueUI(std::shared_ptr<UIScheduler> uiScheduler)
+    : uiScheduler_(std::move(uiScheduler)) {}
+
+void AsyncQueueUI::push(std::function<void()> &&job) {
+  uiScheduler_->scheduleOnUI(std::move(job));
+}
+
 } // namespace worklets

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

@@ -2,6 +2,7 @@
 
 #include <jsi/jsi.h>
 #include <worklets/Public/AsyncQueue.h>
+#include <worklets/Tools/UIScheduler.h>
 
 #include <atomic>
 #include <condition_variable>
@@ -30,4 +31,16 @@ class AsyncQueueImpl : public AsyncQueue {
   const std::shared_ptr<AsyncQueueState> state_;
 };
 
+class AsyncQueueUI : public AsyncQueue {
+ public:
+  explicit AsyncQueueUI(std::shared_ptr<UIScheduler> uiScheduler);
+
+  ~AsyncQueueUI() override = default;
+
+  void push(std::function<void()> &&job) override;
+
+ private:
+  std::shared_ptr<UIScheduler> uiScheduler_;
+};
+
 } // namespace worklets

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

@@ -44,12 +44,13 @@ std::shared_ptr<WorkletRuntime> RuntimeManager::getUIRuntime() {
 std::shared_ptr<WorkletRuntime> RuntimeManager::createWorkletRuntime(
     std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy,
     const std::string &name,
-    std::shared_ptr<SerializableWorklet> initializer) {
+    std::shared_ptr<SerializableWorklet> initializer,
+    const std::shared_ptr<AsyncQueue> &queue) {
   const auto runtimeId = getNextRuntimeId();
   const auto jsQueue = jsiWorkletsModuleProxy->getJSQueue();
 
   auto workletRuntime =
-      std::make_shared<WorkletRuntime>(runtimeId, jsQueue, name);
+      std::make_shared<WorkletRuntime>(runtimeId, jsQueue, name, queue);
 
   workletRuntime->init(std::move(jsiWorkletsModuleProxy));
 
@@ -65,9 +66,10 @@ std::shared_ptr<WorkletRuntime> RuntimeManager::createWorkletRuntime(
 }
 
 std::shared_ptr<WorkletRuntime> RuntimeManager::createUninitializedUIRuntime(
-    const std::shared_ptr<MessageQueueThread> &jsQueue) {
-  const auto uiRuntime =
-      std::make_shared<WorkletRuntime>(uiRuntimeId, jsQueue, uiRuntimeName);
+    const std::shared_ptr<MessageQueueThread> &jsQueue,
+    std::shared_ptr<AsyncQueue> uiAsyncQueue) {
+  const auto uiRuntime = std::make_shared<WorkletRuntime>(
+      uiRuntimeId, jsQueue, uiRuntimeName, uiAsyncQueue);
   std::unique_lock lock(weakRuntimesMutex_);
   weakRuntimes_[uiRuntimeId] = uiRuntime;
   return uiRuntime;

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

@@ -31,10 +31,12 @@ class RuntimeManager {
   std::shared_ptr<WorkletRuntime> createWorkletRuntime(
       std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy,
       const std::string &name,
-      std::shared_ptr<SerializableWorklet> initializer = nullptr);
+      std::shared_ptr<SerializableWorklet> initializer = nullptr,
+      const std::shared_ptr<AsyncQueue> &queue = nullptr);
 
   std::shared_ptr<WorkletRuntime> createUninitializedUIRuntime(
-      const std::shared_ptr<MessageQueueThread> &jsQueue);
+      const std::shared_ptr<MessageQueueThread> &jsQueue,
+      std::shared_ptr<AsyncQueue> uiAsyncQueue);
 
  private:
   uint64_t getNextRuntimeId();

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

@@ -71,11 +71,13 @@ static std::shared_ptr<jsi::Runtime> makeRuntime(
 WorkletRuntime::WorkletRuntime(
     uint64_t runtimeId,
     const std::shared_ptr<MessageQueueThread> &jsQueue,
-    const std::string &name)
+    const std::string &name,
+    const std::shared_ptr<AsyncQueue> &queue)
     : runtimeId_(runtimeId),
       runtimeMutex_(std::make_shared<std::recursive_mutex>()),
       runtime_(makeRuntime(jsQueue, name, runtimeMutex_)),
-      name_(name) {
+      name_(name),
+      queue_(queue) {
   jsi::Runtime &rt = *runtime_;
   WorkletRuntimeCollector::install(rt);
 }
@@ -130,6 +132,22 @@ void WorkletRuntime::init(
 #endif // WORKLETS_BUNDLE_MODE
 }
 
+void WorkletRuntime::runAsyncGuarded(
+    const std::shared_ptr<SerializableWorklet> &serializableWorklet) {
+  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.");
+
+  queue_->push([=, weakThis = weak_from_this()] {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    strongThis->runGuarded(serializableWorklet);
+  });
+}
+
 jsi::Value WorkletRuntime::executeSync(
     jsi::Runtime &rt,
     const jsi::Value &worklet) const {

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

@@ -3,6 +3,7 @@
 #include <cxxreact/MessageQueueThread.h>
 #include <jsi/jsi.h>
 #include <jsireact/JSIExecutor.h>
+#include <react/debug/react_native_assert.h>
 
 #include <worklets/Public/AsyncQueue.h>
 #include <worklets/SharedItems/Shareables.h>
@@ -31,7 +32,8 @@ class WorkletRuntime : public jsi::HostObject,
   explicit WorkletRuntime(
       uint64_t runtimeId,
       const std::shared_ptr<MessageQueueThread> &jsQueue,
-      const std::string &name);
+      const std::string &name,
+      const std::shared_ptr<AsyncQueue> &queue = nullptr);
 
   void init(std::shared_ptr<JSIWorkletsModuleProxy> jsiWorkletsModuleProxy);
 
@@ -49,19 +51,7 @@ class WorkletRuntime : public jsi::HostObject,
   }
 
   void runAsyncGuarded(
-      const std::shared_ptr<SerializableWorklet> &serializableWorklet) {
-    if (queue_ == nullptr) {
-      queue_ = std::make_shared<AsyncQueueImpl>(name_);
-    }
-    queue_->push([=, weakThis = weak_from_this()] {
-      auto strongThis = weakThis.lock();
-      if (!strongThis) {
-        return;
-      }
-
-      strongThis->runGuarded(serializableWorklet);
-    });
-  }
+      const std::shared_ptr<SerializableWorklet> &serializableWorklet);
 
   jsi::Value executeSync(jsi::Runtime &rt, const jsi::Value &worklet) const;
 

packages/react-native-worklets/__typetests__/createWorkletRuntime.tsx

@@ -5,36 +5,56 @@ export function createWorkletRuntimeTypeTests() {
     'worklet';
   };
 
-  // Correct usage - config object
+  // Correct usage - no parameters at all.
+  createWorkletRuntime();
+
+  // Correct usage - empty config object.
+  createWorkletRuntime({});
+
+  // Correct usage - config object with name.
   createWorkletRuntime({
     name: 'test',
+  });
+
+  // Correct usage - config object with initializer.
+  createWorkletRuntime({
     initializer,
   });
 
-  // Correct usage - deprecated positional parameters
-  createWorkletRuntime('test', () => {
-    'worklet';
-    console.log('test');
+  // Correct usage - config object with useDefaultQueue = true.
+  createWorkletRuntime({
+    useDefaultQueue: false,
   });
 
-  // @ts-expect-error - Missing name in config object
+  // Correct usage - config object with useDefaultQueue = false.
   createWorkletRuntime({
-    initializer,
+    useDefaultQueue: false,
   });
 
+  // Correct usage - config object with useDefaultQueue = false and customQueue.
+  createWorkletRuntime({
+    useDefaultQueue: false,
+    customQueue: {},
+  });
+
+  // Correct usage - deprecated positional parameters
+  createWorkletRuntime('test', initializer);
+
   // @ts-expect-error - Wrong name type in config object
   createWorkletRuntime({
     name: 123,
-    initializer,
   });
 
-  // @ts-expect-error - No parameters at all
-  createWorkletRuntime();
-
   // @ts-expect-error - Not existing parameter
   createWorkletRuntime({
     name: 'test',
     test: 'test',
     initializer,
   });
+
+  // @ts-expect-error - Using both useDefaultQueue and customQueue
+  createWorkletRuntime({
+    useDefaultQueue: true,
+    customQueue: {},
+  });
 }

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

@@ -3,7 +3,7 @@
 import { IS_JEST } from '../PlatformChecker';
 import { mockedRequestAnimationFrame } from '../runLoop/mockedRequestAnimationFrame';
 import { WorkletsError } from '../WorkletsError';
-import type { ShareableRef, WorkletRuntime } from '../workletTypes';
+import type { ShareableRef } from '../workletTypes';
 import type { IWorkletsModule } from './workletsModuleProxy';
 
 export function createJSWorkletsModule(): IWorkletsModule {
@@ -137,10 +137,7 @@ class JSWorklets implements IWorkletsModule {
     );
   }
 
-  createWorkletRuntime(
-    _name: string,
-    _initializer: ShareableRef<() => void>
-  ): WorkletRuntime {
+  createWorkletRuntime(): never {
     throw new WorkletsError(
       'createWorkletRuntime is not available in JSWorklets.'
     );

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

@@ -152,8 +152,18 @@ See https://docs.swmansion.com/react-native-worklets/docs/guides/troubleshooting
     return this.#workletsModuleProxy.executeOnUIRuntimeSync(shareable);
   }
 
-  createWorkletRuntime(name: string, initializer: ShareableRef<() => void>) {
-    return this.#workletsModuleProxy.createWorkletRuntime(name, initializer);
+  createWorkletRuntime(
+    name: string,
+    initializer: ShareableRef<() => void>,
+    useDefaultQueue: boolean,
+    customQueue: object | undefined
+  ) {
+    return this.#workletsModuleProxy.createWorkletRuntime(
+      name,
+      initializer,
+      useDefaultQueue,
+      customQueue
+    );
   }
 
   scheduleOnRuntime<T>(

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

@@ -71,7 +71,9 @@ export interface WorkletsModuleProxy {
 
   createWorkletRuntime(
     name: string,
-    initializer: ShareableRef<() => void>
+    initializer: ShareableRef<() => void>,
+    useDefaultQueue: boolean,
+    customQueue: object | undefined
   ): WorkletRuntime;
 
   scheduleOnRuntime<TValue>(