feat(Worklets): allow passing user-implemented AsyncQueue to createWorkletRuntime

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/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));
 

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

@@ -31,7 +31,8 @@ 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);

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);
 }

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);
 
@@ -50,9 +52,9 @@ class WorkletRuntime : public jsi::HostObject,
 
   void runAsyncGuarded(
       const std::shared_ptr<SerializableWorklet> &serializableWorklet) {
-    if (queue_ == nullptr) {
-      queue_ = std::make_shared<AsyncQueueImpl>(name_);
-    }
+    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) {

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>(

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

@@ -16,14 +16,14 @@ import type { WorkletFunction, WorkletRuntime } from './workletTypes';
  * Lets you create a new JS runtime which can be used to run worklets possibly
  * on different threads than JS or UI thread.
  *
- * @param config - Configuration object containing runtime name and optional
- *   initializer
+ * @param config - Runtime configuration object - {@link WorkletRuntimeConfig}.
  * @returns WorkletRuntime which is a
  *   `jsi::HostObject<worklets::WorkletRuntime>` - {@link WorkletRuntime}
  * @see https://docs.swmansion.com/react-native-reanimated/docs/threading/createWorkletRuntime
  */
+// @ts-expect-error Public API overload.
 export function createWorkletRuntime(
-  config: WorkletRuntimeConfig
+  config?: WorkletRuntimeConfig
 ): WorkletRuntime;
 
 /**
@@ -40,26 +40,36 @@ export function createWorkletRuntime(
  *   `jsi::HostObject<worklets::WorkletRuntime>` - {@link WorkletRuntime}
  * @see https://docs.swmansion.com/react-native-reanimated/docs/threading/createWorkletRuntime
  */
-// @ts-expect-error Check `runOnUI` overload.
 export function createWorkletRuntime(
-  name: string,
+  name?: string,
   initializer?: () => void
 ): WorkletRuntime;
 
 export function createWorkletRuntime(
-  nameOrConfig: string | WorkletRuntimeConfig,
+  nameOrConfig?: string | WorkletRuntimeConfigInternal,
   initializer?: WorkletFunction<[], void>
 ): WorkletRuntime {
   const runtimeBoundCapturableConsole = getMemorySafeCapturableConsole();
 
   let name: string;
   let initializerFn: (() => void) | undefined;
+  let useDefaultQueue = true;
+  let customQueue: object | undefined;
   if (typeof nameOrConfig === 'string') {
     name = nameOrConfig;
     initializerFn = initializer;
   } else {
-    name = nameOrConfig.name;
-    initializerFn = nameOrConfig.initializer;
+    // TODO: Make anonymous name globally unique.
+    name = nameOrConfig?.name ?? 'anonymous';
+    initializerFn = nameOrConfig?.initializer;
+    useDefaultQueue = nameOrConfig?.useDefaultQueue ?? true;
+    customQueue = nameOrConfig?.customQueue;
+  }
+
+  if (initializerFn && !isWorkletFunction(initializerFn)) {
+    throw new WorkletsError(
+      'The initializer passed to `createWorkletRuntime` is not a worklet.'
+    );
   }
 
   return WorkletsModule.createWorkletRuntime(
@@ -70,7 +80,9 @@ export function createWorkletRuntime(
       registerWorkletsError();
       setupConsole(runtimeBoundCapturableConsole);
       initializerFn?.();
-    })
+    }),
+    useDefaultQueue,
+    customQueue
   );
 }
 
@@ -92,7 +104,7 @@ export function runOnRuntime<Args extends unknown[], ReturnValue>(
   }
   if (globalThis._WORKLET) {
     return (...args) =>
-      global._scheduleOnRuntime(
+      globalThis._scheduleOnRuntime(
         workletRuntime,
         makeShareableCloneOnUIRecursive(() => {
           'worklet';
@@ -110,7 +122,46 @@ export function runOnRuntime<Args extends unknown[], ReturnValue>(
     );
 }
 
+/** Configuration object for creating a worklet runtime. */
 export type WorkletRuntimeConfig = {
-  name: string;
+  /** The name of the worklet runtime. */
+  name?: string;
+  /**
+   * A worklet that will be run immediately after the runtime is created and
+   * before any other worklets.
+   */
   initializer?: () => void;
+} & (
+  | {
+      /**
+       * If true, the runtime will use the default queue implementation for
+       * scheduling worklets. Defaults to true.
+       */
+      useDefaultQueue?: true;
+      /**
+       * An optional custom queue to be used for scheduling worklets.
+       *
+       * The queue has to implement the C++ `AsyncQueue` interface from
+       * `<worklets/Public/AsyncQueue.h>`.
+       */
+      customQueue?: never;
+    }
+  | {
+      /**
+       * If true, the runtime will use the default queue implementation for
+       * scheduling worklets. Defaults to true.
+       */
+      useDefaultQueue: false;
+      /**
+       * An optional custom queue to be used for scheduling worklets.
+       *
+       * The queue has to implement the C++ `AsyncQueue` interface from
+       * `<worklets/Public/AsyncQueue.h>`.
+       */
+      customQueue?: object;
+    }
+);
+
+type WorkletRuntimeConfigInternal = WorkletRuntimeConfig & {
+  initializer?: WorkletFunction<[], void>;
 };