Feat/worklet runner performance improvement

packages/react-native-audio-api/common/cpp/audioapi/HostObjects/BaseAudioContextHostObject.cpp

@@ -77,14 +77,15 @@ JSI_HOST_FUNCTION_IMPL(BaseAudioContextHostObject, createWorkletSourceNode) {
           runtime, args[0]);
   std::weak_ptr<worklets::WorkletRuntime> workletRuntime;
   auto shouldUseUiRuntime = args[1].getBool();
+  auto shouldLockRuntime = shouldUseUiRuntime;
   if (shouldUseUiRuntime) {
     workletRuntime = context_->runtimeRegistry_.uiRuntime;
   } else {
     workletRuntime = context_->runtimeRegistry_.audioRuntime;
   }
 
-  auto workletSourceNode =
-      context_->createWorkletSourceNode(shareableWorklet, workletRuntime);
+  auto workletSourceNode = context_->createWorkletSourceNode(
+      shareableWorklet, workletRuntime, shouldLockRuntime);
   auto workletSourceNodeHostObject =
       std::make_shared<WorkletSourceNodeHostObject>(workletSourceNode);
   return jsi::Object::createFromHostObject(
@@ -101,6 +102,7 @@ JSI_HOST_FUNCTION_IMPL(BaseAudioContextHostObject, createWorkletNode) {
 
   std::weak_ptr<worklets::WorkletRuntime> workletRuntime;
   auto shouldUseUiRuntime = args[1].getBool();
+  auto shouldLockRuntime = shouldUseUiRuntime;
   if (shouldUseUiRuntime) {
     workletRuntime = context_->runtimeRegistry_.uiRuntime;
   } else {
@@ -110,7 +112,11 @@ JSI_HOST_FUNCTION_IMPL(BaseAudioContextHostObject, createWorkletNode) {
   auto inputChannelCount = static_cast<size_t>(args[3].getNumber());
 
   auto workletNode = context_->createWorkletNode(
-      shareableWorklet, workletRuntime, bufferLength, inputChannelCount);
+      shareableWorklet,
+      workletRuntime,
+      bufferLength,
+      inputChannelCount,
+      shouldLockRuntime);
   auto workletNodeHostObject =
       std::make_shared<WorkletNodeHostObject>(workletNode);
   return jsi::Object::createFromHostObject(runtime, workletNodeHostObject);
@@ -128,14 +134,15 @@ JSI_HOST_FUNCTION_IMPL(
 
   std::weak_ptr<worklets::WorkletRuntime> workletRuntime;
   auto shouldUseUiRuntime = args[1].getBool();
+  auto shouldLockRuntime = shouldUseUiRuntime;
   if (shouldUseUiRuntime) {
     workletRuntime = context_->runtimeRegistry_.uiRuntime;
   } else {
     workletRuntime = context_->runtimeRegistry_.audioRuntime;
   }
 
-  auto workletProcessingNode =
-      context_->createWorkletProcessingNode(shareableWorklet, workletRuntime);
+  auto workletProcessingNode = context_->createWorkletProcessingNode(
+      shareableWorklet, workletRuntime, shouldLockRuntime);
   auto workletProcessingNodeHostObject =
       std::make_shared<WorkletProcessingNodeHostObject>(workletProcessingNode);
   return jsi::Object::createFromHostObject(

packages/react-native-audio-api/common/cpp/audioapi/core/BaseAudioContext.cpp

@@ -67,9 +67,11 @@ std::shared_ptr<AudioDestinationNode> BaseAudioContext::getDestination() {
 
 std::shared_ptr<WorkletSourceNode> BaseAudioContext::createWorkletSourceNode(
     std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime) {
+    std::weak_ptr<worklets::WorkletRuntime> runtime,
+    bool shouldLockRuntime) {
+  WorkletsRunner workletRunner(runtime, shareableWorklet, shouldLockRuntime);
   auto workletSourceNode =
-      std::make_shared<WorkletSourceNode>(this, shareableWorklet, runtime);
+      std::make_shared<WorkletSourceNode>(this, std::move(workletRunner));
   nodeManager_->addSourceNode(workletSourceNode);
   return workletSourceNode;
 }
@@ -78,19 +80,23 @@ std::shared_ptr<WorkletNode> BaseAudioContext::createWorkletNode(
     std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
     std::weak_ptr<worklets::WorkletRuntime> runtime,
     size_t bufferLength,
-    size_t inputChannelCount) {
+    size_t inputChannelCount,
+    bool shouldLockRuntime) {
+  WorkletsRunner workletRunner(runtime, shareableWorklet, shouldLockRuntime);
   auto workletNode = std::make_shared<WorkletNode>(
-      this, shareableWorklet, runtime, bufferLength, inputChannelCount);
+      this, bufferLength, inputChannelCount, std::move(workletRunner));
   nodeManager_->addProcessingNode(workletNode);
   return workletNode;
 }
 
 std::shared_ptr<WorkletProcessingNode>
 BaseAudioContext::createWorkletProcessingNode(
     std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime) {
+    std::weak_ptr<worklets::WorkletRuntime> runtime,
+    bool shouldLockRuntime) {
+  WorkletsRunner workletRunner(runtime, shareableWorklet, shouldLockRuntime);
   auto workletProcessingNode =
-      std::make_shared<WorkletProcessingNode>(this, shareableWorklet, runtime);
+      std::make_shared<WorkletProcessingNode>(this, std::move(workletRunner));
   nodeManager_->addProcessingNode(workletProcessingNode);
   return workletProcessingNode;
 }

packages/react-native-audio-api/common/cpp/audioapi/core/BaseAudioContext.h

@@ -47,9 +47,20 @@ class BaseAudioContext {
   std::shared_ptr<AudioDestinationNode> getDestination();
 
   std::shared_ptr<RecorderAdapterNode> createRecorderAdapter();
-  std::shared_ptr<WorkletSourceNode> createWorkletSourceNode(std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet, std::weak_ptr<worklets::WorkletRuntime> runtime);
-  std::shared_ptr<WorkletNode> createWorkletNode(std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet, std::weak_ptr<worklets::WorkletRuntime> runtime, size_t bufferLength, size_t inputChannelCount);
-  std::shared_ptr<WorkletProcessingNode> createWorkletProcessingNode(std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet, std::weak_ptr<worklets::WorkletRuntime> runtime);
+  std::shared_ptr<WorkletSourceNode> createWorkletSourceNode(
+    std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
+    std::weak_ptr<worklets::WorkletRuntime> runtime,
+    bool shouldLockRuntime = true);
+  std::shared_ptr<WorkletNode> createWorkletNode(
+    std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
+    std::weak_ptr<worklets::WorkletRuntime> runtime,
+    size_t bufferLength,
+    size_t inputChannelCount,
+    bool shouldLockRuntime = true);
+  std::shared_ptr<WorkletProcessingNode> createWorkletProcessingNode(
+    std::shared_ptr<worklets::SerializableWorklet> &shareableWorklet,
+    std::weak_ptr<worklets::WorkletRuntime> runtime,
+    bool shouldLockRuntime = true);
   std::shared_ptr<OscillatorNode> createOscillator();
   std::shared_ptr<ConstantSourceNode> createConstantSource();
   std::shared_ptr<StreamerNode> createStreamer();

packages/react-native-audio-api/common/cpp/audioapi/core/effects/WorkletNode.cpp

@@ -4,13 +4,11 @@ namespace audioapi {
 
 WorkletNode::WorkletNode(
     BaseAudioContext *context,
-    std::shared_ptr<worklets::SerializableWorklet> &worklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime,
     size_t bufferLength,
-    size_t inputChannelCount)
+    size_t inputChannelCount,
+    WorkletsRunner &&runtime)
     : AudioNode(context),
-      workletRunner_(runtime),
-      shareableWorklet_(worklet),
+      workletRunner_(std::move(runtime)),
       bufferLength_(bufferLength),
       inputChannelCount_(inputChannelCount),
       curBuffIndex_(0) {
@@ -40,35 +38,36 @@ std::shared_ptr<AudioBus> WorkletNode::processNode(
     curBuffIndex_ += shouldProcess;
 
     /// If we filled the entire buffer, we need to execute the worklet
-    if (curBuffIndex_ == bufferLength_) {
-      // Reset buffer index, channel buffers and execute worklet
-      curBuffIndex_ = 0;
-      workletRunner_.executeOnRuntimeGuardedSync(
-          [this, channelCount_](jsi::Runtime &uiRuntimeRaw) {
-            /// Arguments preparation
-            auto jsArray = jsi::Array(uiRuntimeRaw, channelCount_);
-            for (size_t ch = 0; ch < channelCount_; ch++) {
-              auto audioArray = std::make_shared<AudioArray>(bufferLength_);
-              audioArray->copy(bus_->getChannel(ch));
-              auto sharedAudioArray =
-                  std::make_shared<AudioArrayBuffer>(audioArray);
-              auto sharedAudioArraySize = sharedAudioArray->size();
-              auto arrayBuffer =
-                  jsi::ArrayBuffer(uiRuntimeRaw, std::move(sharedAudioArray));
-              arrayBuffer.setExternalMemoryPressure(
-                  uiRuntimeRaw, sharedAudioArraySize);
-              jsArray.setValueAtIndex(uiRuntimeRaw, ch, std::move(arrayBuffer));
-            }
+    if (curBuffIndex_ != bufferLength_) {
+      continue;
+    }
+    // Reset buffer index, channel buffers and execute worklet
+    curBuffIndex_ = 0;
+    workletRunner_.executeOnRuntimeSync([this, channelCount_](
+                                            jsi::Runtime &uiRuntimeRaw) {
+      /// Arguments preparation
+      auto jsArray = jsi::Array(uiRuntimeRaw, channelCount_);
+      for (size_t ch = 0; ch < channelCount_; ch++) {
+        auto audioArray = std::make_shared<AudioArray>(bufferLength_);
+        audioArray->copy(bus_->getChannel(ch));
+        auto sharedAudioArray = std::make_shared<AudioArrayBuffer>(audioArray);
+        auto sharedAudioArraySize = sharedAudioArray->size();
+        auto arrayBuffer =
+            jsi::ArrayBuffer(uiRuntimeRaw, std::move(sharedAudioArray));
+        arrayBuffer.setExternalMemoryPressure(
+            uiRuntimeRaw, sharedAudioArraySize);
+        jsArray.setValueAtIndex(uiRuntimeRaw, ch, std::move(arrayBuffer));
+      }
 
-            bus_->zero();
+      bus_->zero();
 
-            workletRunner_.executeWorklet(
-                shareableWorklet_,
-                std::move(jsArray),
-                jsi::Value(uiRuntimeRaw, static_cast<int>(channelCount_)));
-            return jsi::Value::undefined();
-          });
-    }
+      /// Call the worklet
+      workletRunner_.callUnsafe(
+          std::move(jsArray),
+          jsi::Value(uiRuntimeRaw, static_cast<int>(channelCount_)));
+
+      return jsi::Value::undefined();
+    });
   }
 
   return processingBus;

packages/react-native-audio-api/common/cpp/audioapi/core/effects/WorkletNode.h

@@ -19,10 +19,9 @@ class WorkletNode : public AudioNode {
  public:
   explicit WorkletNode(
       BaseAudioContext *context,
-      std::shared_ptr<worklets::SerializableWorklet> &worklet,
-      std::weak_ptr<worklets::WorkletRuntime> runtime,
       size_t bufferLength,
-      size_t inputChannelCount
+      size_t inputChannelCount,
+      WorkletsRunner &&workletRunner
   ) : AudioNode(context) {}
 
  protected:
@@ -36,10 +35,9 @@ class WorkletNode : public AudioNode {
  public:
   explicit WorkletNode(
       BaseAudioContext *context,
-      std::shared_ptr<worklets::SerializableWorklet> &worklet,
-      std::weak_ptr<worklets::WorkletRuntime> runtime,
       size_t bufferLength,
-      size_t inputChannelCount
+      size_t inputChannelCount,
+      WorkletsRunner &&workletRunner
   );
 
   ~WorkletNode() override = default;
@@ -50,7 +48,6 @@ class WorkletNode : public AudioNode {
 
  private:
   WorkletsRunner workletRunner_;
-  std::shared_ptr<worklets::SerializableWorklet> shareableWorklet_;
   std::shared_ptr<AudioBus> bus_;
 
   /// @brief Length of the byte buffer that will be passed to the AudioArrayBuffer

packages/react-native-audio-api/common/cpp/audioapi/core/effects/WorkletProcessingNode.cpp

@@ -5,9 +5,8 @@ namespace audioapi {
 
 WorkletProcessingNode::WorkletProcessingNode(
     BaseAudioContext *context,
-    std::shared_ptr<worklets::SerializableWorklet> &worklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime)
-    : AudioNode(context), workletRunner_(runtime), shareableWorklet_(worklet) {
+    WorkletsRunner &&workletRunner)
+    : AudioNode(context), workletRunner_(std::move(workletRunner)) {
   isInitialized_ = true;
 
   // Pre-allocate buffers for max 128 frames and 2 channels (stereo)
@@ -42,8 +41,8 @@ std::shared_ptr<AudioBus> WorkletProcessingNode::processNode(
   }
 
   // Execute the worklet
-  auto result = workletRunner_.executeOnRuntimeGuardedSync(
-      [this, channelCount, framesToProcess](jsi::Runtime &rt) {
+  auto result = workletRunner_.executeOnRuntimeSync(
+      [this, channelCount, framesToProcess](jsi::Runtime &rt) -> jsi::Value {
         auto inputJsArray = jsi::Array(rt, channelCount);
         auto outputJsArray = jsi::Array(rt, channelCount);
 
@@ -58,14 +57,14 @@ std::shared_ptr<AudioBus> WorkletProcessingNode::processNode(
           outputJsArray.setValueAtIndex(rt, ch, outputArrayBuffer);
         }
 
-        return workletRunner_
-            .executeWorklet(
-                shareableWorklet_,
-                inputJsArray,
-                outputJsArray,
-                jsi::Value(rt, static_cast<int>(framesToProcess)),
-                jsi::Value(rt, this->context_->getCurrentTime()))
-            .value_or(jsi::Value::undefined());
+        // We call unsafely here because we are already on the runtime thread
+        // and the runtime is locked by executeOnRuntimeSync (if
+        // shouldLockRuntime is true)
+        return workletRunner_.callUnsafe(
+            inputJsArray,
+            outputJsArray,
+            jsi::Value(rt, static_cast<int>(framesToProcess)),
+            jsi::Value(rt, this->context_->getCurrentTime()));
       });
 
   // Copy processed output data back to the processing bus or zero on failure

packages/react-native-audio-api/common/cpp/audioapi/core/effects/WorkletProcessingNode.h

@@ -18,8 +18,7 @@ class WorkletProcessingNode : public AudioNode {
  public:
   explicit WorkletProcessingNode(
       BaseAudioContext *context,
-      std::shared_ptr<worklets::SerializableWorklet> &worklet,
-      std::weak_ptr<worklets::WorkletRuntime> runtime
+      WorkletsRunner &&workletRunner
   ) : AudioNode(context) {}
 
  protected:
@@ -33,16 +32,14 @@ class WorkletProcessingNode : public AudioNode {
  public:
   explicit WorkletProcessingNode(
       BaseAudioContext *context,
-      std::shared_ptr<worklets::SerializableWorklet> &worklet,
-      std::weak_ptr<worklets::WorkletRuntime> runtime
+      WorkletsRunner &&workletRunner
   );
 
  protected:
   std::shared_ptr<AudioBus> processNode(const std::shared_ptr<AudioBus>& processingBus, int framesToProcess) override;
 
  private:
   WorkletsRunner workletRunner_;
-  std::shared_ptr<worklets::SerializableWorklet> shareableWorklet_;
   std::vector<std::shared_ptr<AudioArrayBuffer>> inputBuffsHandles_;
   std::vector<std::shared_ptr<AudioArrayBuffer>> outputBuffsHandles_;
 };

packages/react-native-audio-api/common/cpp/audioapi/core/sources/WorkletSourceNode.cpp

@@ -5,11 +5,9 @@ namespace audioapi {
 
 WorkletSourceNode::WorkletSourceNode(
     BaseAudioContext *context,
-    std::shared_ptr<worklets::SerializableWorklet> &worklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime)
+    WorkletsRunner &&workletRunner)
     : AudioScheduledSourceNode(context),
-      workletRunner_(runtime),
-      shareableWorklet_(worklet) {
+      workletRunner_(std::move(workletRunner)) {
   isInitialized_ = true;
 
   // Prepare buffers for audio processing
@@ -42,21 +40,22 @@ std::shared_ptr<AudioBus> WorkletSourceNode::processNode(
 
   size_t outputChannelCount = processingBus->getNumberOfChannels();
 
-  auto result = workletRunner_.executeOnRuntimeGuardedSync(
+  auto result = workletRunner_.executeOnRuntimeSync(
       [this, nonSilentFramesToProcess, startOffset](jsi::Runtime &rt) {
         auto jsiArray = jsi::Array(rt, this->outputBuffsHandles_.size());
         for (size_t i = 0; i < this->outputBuffsHandles_.size(); ++i) {
           auto arrayBuffer = jsi::ArrayBuffer(rt, this->outputBuffsHandles_[i]);
           jsiArray.setValueAtIndex(rt, i, arrayBuffer);
         }
-        return workletRunner_
-            .executeWorklet(
-                shareableWorklet_,
-                jsiArray,
-                jsi::Value(rt, static_cast<int>(nonSilentFramesToProcess)),
-                jsi::Value(rt, this->context_->getCurrentTime()),
-                jsi::Value(rt, static_cast<int>(startOffset)))
-            .value_or(jsi::Value::undefined());
+
+        // We call unsafely here because we are already on the runtime thread
+        // and the runtime is locked by executeOnRuntimeSync (if
+        // shouldLockRuntime is true)
+        return workletRunner_.callUnsafe(
+            jsiArray,
+            jsi::Value(rt, static_cast<int>(nonSilentFramesToProcess)),
+            jsi::Value(rt, this->context_->getCurrentTime()),
+            jsi::Value(rt, static_cast<int>(startOffset)));
       });
 
   // If the worklet execution failed, zero the output

packages/react-native-audio-api/common/cpp/audioapi/core/sources/WorkletSourceNode.h

@@ -18,8 +18,7 @@ class WorkletSourceNode : public AudioScheduledSourceNode {
  public:
   explicit WorkletSourceNode(
     BaseAudioContext *context,
-    std::shared_ptr<worklets::SerializableWorklet> &worklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime
+    WorkletsRunner &&workletRunner
   ) : AudioScheduledSourceNode(context) {}
 
  protected:
@@ -31,15 +30,13 @@ class WorkletSourceNode : public AudioScheduledSourceNode {
  public:
   explicit WorkletSourceNode(
     BaseAudioContext *context,
-    std::shared_ptr<worklets::SerializableWorklet> &worklet,
-    std::weak_ptr<worklets::WorkletRuntime> runtime
+    WorkletsRunner &&workletRunner
   );
 
  protected:
   std::shared_ptr<AudioBus> processNode(const std::shared_ptr<AudioBus>& processingBus, int framesToProcess) override;
  private:
   WorkletsRunner workletRunner_;
-  std::shared_ptr<worklets::SerializableWorklet> shareableWorklet_;
   std::vector<std::shared_ptr<AudioArrayBuffer>> outputBuffsHandles_;
 };
 #endif // RN_AUDIO_API_TEST