feat(Worklets): runAsyncGuardedOnUI

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(const std::shared_ptr<UIScheduler> &uiScheduler)
+    : uiScheduler_(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(const 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

@@ -66,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,
+    const 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

@@ -35,7 +35,8 @@ class RuntimeManager {
       const std::shared_ptr<AsyncQueue> &queue = nullptr);
 
   std::shared_ptr<WorkletRuntime> createUninitializedUIRuntime(
-      const std::shared_ptr<MessageQueueThread> &jsQueue);
+      const std::shared_ptr<MessageQueueThread> &jsQueue,
+      const std::shared_ptr<AsyncQueue> &uiAsyncQueue);
 
  private:
   uint64_t getNextRuntimeId();

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

@@ -132,6 +132,22 @@ void WorkletRuntime::init(
 #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.");
+
+  queue_->push([worklet, weakThis = weak_from_this()] {
+    auto strongThis = weakThis.lock();
+    if (!strongThis) {
+      return;
+    }
+
+    strongThis->runGuarded(worklet);
+  });
+}
+
 jsi::Value WorkletRuntime::executeSync(
     jsi::Runtime &rt,
     const jsi::Value &worklet) const {

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

@@ -50,20 +50,7 @@ class WorkletRuntime : public jsi::HostObject,
         rt, serializableWorklet->toJSValue(rt), std::forward<Args>(args)...);
   }
 
-  void 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);
-    });
-  }
+  void runAsyncGuarded(const std::shared_ptr<SerializableWorklet> &worklet);
 
   jsi::Value executeSync(jsi::Runtime &rt, const jsi::Value &worklet) const;