feat: add zero-copy ArrayBuffer for Java TM

Author: paradowstack

packages/react-native/ReactAndroid/src/main/jni/react/jni/JByteBufferMutableBuffer.h

@@ -39,6 +39,11 @@ class JByteBufferMutableBuffer final : public jsi::MutableBuffer {
   JByteBufferMutableBuffer(JByteBufferMutableBuffer &&) = delete;
   JByteBufferMutableBuffer &operator=(JByteBufferMutableBuffer &&) = delete;
 
+  jni::local_ref<jni::JByteBuffer> getJavaByteBuffer() const
+  {
+    return jni::make_local(byteBuffer_);
+  }
+
   size_t size() const override
   {
     return byteBuffer_->getDirectSize();

packages/react-native/ReactCommon/react/nativemodule/core/platform/android/ReactCommon/JavaTurboModule.cpp

@@ -20,6 +20,7 @@
 #include <ReactCommon/TurboModule.h>
 #include <ReactCommon/TurboModulePerfLogger.h>
 #include <jsi/JSIDynamic.h>
+#include <react/bridging/ArrayBuffer.h>
 #include <react/bridging/Bridging.h>
 #include <react/debug/react_native_assert.h>
 #include <react/featureflags/ReactNativeFeatureFlags.h>
@@ -63,6 +64,15 @@ JavaTurboModule::~JavaTurboModule() {
 
 namespace {
 
+// Aliased ArrayBuffer arg passed without copying. nativeOwner keeps native
+// bytes alive for the invocation, or is null when JS heap bytes were borrowed 
+// for a synchronous call.
+struct AliasedArrayBuffer {
+  const uint8_t* bytes{nullptr};
+  size_t size{0};
+  std::shared_ptr<jsi::MutableBuffer> nativeOwner;
+};
+
 struct JNIArgs {
   JNIArgs(size_t count) : args(count) {}
   JNIArgs(const JNIArgs&) = delete;
@@ -73,6 +83,7 @@ struct JNIArgs {
 
   std::vector<jvalue> args;
   std::vector<jobject> globalRefs;
+  std::vector<AliasedArrayBuffer> aliasedArrayBuffers;
 
   ~JNIArgs() {
     JNIEnv* env = jni::Environment::current();
@@ -315,6 +326,8 @@ JNIArgs convertJSIArgsToJNIArgs(
   auto& jargs = jniArgs.args;
   auto& globalRefs = jniArgs.globalRefs;
 
+  auto isSyncInvocation = valueKind != VoidKind && valueKind != PromiseKind;
+
   auto makeGlobalIfNecessary = [&](jobject obj) {
     if (valueKind == VoidKind || valueKind == PromiseKind) {
       jobject globalObj = env->NewGlobalRef(obj);
@@ -457,17 +470,42 @@ JNIArgs convertJSIArgsToJNIArgs(
             rt,
             "JavaTurboModule::convertJSIArgsToJNIArgs: ArrayBuffer exceeds maximum size.");
       }
-      auto data = arrayBuffer.data(rt);
-      // ArrayBuffer arguments are always copied into a Java-owned direct
-      // ByteBuffer, so Java fully owns the bytes. Borrowing the JS bytes is
-      // never safe — even on a synchronous call the module may retain the
-      // buffer or hand it to an async method, and JS may garbage-collect the
-      // source ArrayBuffer, leaving Java with a dangling view.
-      auto buffer = jni::JByteBuffer::allocateDirect(static_cast<jint>(size));
-      if (size > 0) {
-        // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy
-        std::memcpy(buffer->getDirectBytes(), data, size);
+
+      // Empty buffers cannot use NewDirectByteBuffer; allocate an empty direct buffer.
+      if (size == 0) {
+        jarg->l = makeGlobalIfNecessary(
+            jni::JByteBuffer::allocateDirect(0).release());
+        continue;
+      }
+
+      // Native-backed: alias and retain MutableBuffer for the invocation.
+      if (auto mutableBuffer = arrayBuffer.tryGetMutableBuffer(rt)) {
+        // Round-trip a Java-origin buffer as the same ByteBuffer object.
+        auto* javaBacked =
+            dynamic_cast<JByteBufferMutableBuffer*>(mutableBuffer.get());
+        auto byteBuffer = javaBacked != nullptr
+            ? javaBacked->getJavaByteBuffer()
+            : jni::JByteBuffer::wrapBytes(mutableBuffer->data(), size);
+        jniArgs.aliasedArrayBuffers.push_back(
+            {byteBuffer->getDirectBytes(), size, std::move(mutableBuffer)});
+        jarg->l = makeGlobalIfNecessary(byteBuffer.release());
+        continue;
+      }
+
+      // JS heap bytes: borrow for sync calls only.
+      if (isSyncInvocation) {
+        auto byteBuffer =
+            jni::JByteBuffer::wrapBytes(arrayBuffer.data(rt), size);
+        jniArgs.aliasedArrayBuffers.push_back(
+            {byteBuffer->getDirectBytes(), size, nullptr});
+        jarg->l = makeGlobalIfNecessary(byteBuffer.release());
+        continue;
       }
+
+      // Async: copy so bytes outlive the JS call.
+      auto buffer = jni::JByteBuffer::allocateDirect(static_cast<jint>(size));
+      // @lint-ignore CLANGSECURITY facebook-security-vulnerable-memcpy
+      std::memcpy(buffer->getDirectBytes(), arrayBuffer.data(rt), size);
       jarg->l = makeGlobalIfNecessary(buffer.release());
     } else {
       throw JavaTurboModuleInvalidArgumentTypeException(
@@ -1016,12 +1054,30 @@ jsi::Value JavaTurboModule::invokeJavaMethod(
               runtime,
               "Only direct ByteBuffers (ByteBuffer.allocateDirect) can be returned from a TurboModule.");
         }
-        // Zero-copy: JByteBufferMutableBuffer takes a global reference that
-        // pins the ByteBuffer's memory for the lifetime of the JS ArrayBuffer,
-        // and its destructor attaches the current thread before releasing that
-        // ref, so JS GC finalization on any thread is safe.
-        auto nativeBuffer =
-            std::make_shared<JByteBufferMutableBuffer>(jByteBuffer);
+        auto* returnedBytes = jByteBuffer->getDirectBytes();
+        auto returnedSize = jByteBuffer->getDirectSize();
+        std::shared_ptr<jsi::MutableBuffer> nativeBuffer;
+
+        // Exact argument match: reuse nativeOwner or copy borrowed bytes.
+        for (const auto& aliased : jniArgs.aliasedArrayBuffers) {
+          if (aliased.bytes != returnedBytes || aliased.size != returnedSize) {
+            continue;
+          }
+          if (aliased.nativeOwner != nullptr) {
+            nativeBuffer = aliased.nativeOwner;
+          } else {
+            nativeBuffer =
+                std::make_shared<detail::OwnedBytesBuffer>(std::vector<uint8_t>(
+                    returnedBytes, returnedBytes + returnedSize));
+          }
+          break;
+        }
+
+        if (nativeBuffer == nullptr) {
+          nativeBuffer =
+              std::make_shared<JByteBufferMutableBuffer>(jByteBuffer);
+        }
+
         returnValue = {
             runtime, jsi::ArrayBuffer{runtime, std::move(nativeBuffer)}};
       }

packages/react-native/ReactCommon/react/nativemodule/samples/platform/android/SampleTurboModule.kt

@@ -155,9 +155,15 @@ public class SampleTurboModule(private val context: ReactApplicationContext) :
     return map
   }
 
+  // Mutating the argument updates the JS ArrayBuffer in place.
   @DoNotStrip
   @Suppress("unused")
   override fun getArrayBuffer(buffer: ByteBuffer?): ByteBuffer? {
+    if (buffer != null) {
+      for (i in 0 until buffer.capacity()) {
+        buffer.put(i, (buffer.get(i) * 2).toByte())
+      }
+    }
     log("getArrayBuffer", buffer, buffer)
     return buffer
   }

scripts/cxx-api/api-snapshots/ReactAndroidDebugCxx.api

@@ -2803,6 +2803,7 @@ class facebook::react::JByteBufferMutableBuffer : public facebook::jsi::MutableB
   public JByteBufferMutableBuffer(jni::alias_ref<jni::JByteBuffer> byteBuffer);
   public facebook::react::JByteBufferMutableBuffer& operator=(const facebook::react::JByteBufferMutableBuffer&) = delete;
   public facebook::react::JByteBufferMutableBuffer& operator=(facebook::react::JByteBufferMutableBuffer&&) = delete;
+  public jni::local_ref<jni::JByteBuffer> getJavaByteBuffer() const;
   public virtual size_t size() const override;
   public virtual uint8_t* data() override;
   public ~JByteBufferMutableBuffer() override;

scripts/cxx-api/api-snapshots/ReactAndroidNewarchCxx.api

@@ -2761,6 +2761,7 @@ class facebook::react::JByteBufferMutableBuffer : public facebook::jsi::MutableB
   public JByteBufferMutableBuffer(jni::alias_ref<jni::JByteBuffer> byteBuffer);
   public facebook::react::JByteBufferMutableBuffer& operator=(const facebook::react::JByteBufferMutableBuffer&) = delete;
   public facebook::react::JByteBufferMutableBuffer& operator=(facebook::react::JByteBufferMutableBuffer&&) = delete;
+  public jni::local_ref<jni::JByteBuffer> getJavaByteBuffer() const;
   public virtual size_t size() const override;
   public virtual uint8_t* data() override;
   public ~JByteBufferMutableBuffer() override;

scripts/cxx-api/api-snapshots/ReactAndroidReleaseCxx.api

@@ -2800,6 +2800,7 @@ class facebook::react::JByteBufferMutableBuffer : public facebook::jsi::MutableB
   public JByteBufferMutableBuffer(jni::alias_ref<jni::JByteBuffer> byteBuffer);
   public facebook::react::JByteBufferMutableBuffer& operator=(const facebook::react::JByteBufferMutableBuffer&) = delete;
   public facebook::react::JByteBufferMutableBuffer& operator=(facebook::react::JByteBufferMutableBuffer&&) = delete;
+  public jni::local_ref<jni::JByteBuffer> getJavaByteBuffer() const;
   public virtual size_t size() const override;
   public virtual uint8_t* data() override;
   public ~JByteBufferMutableBuffer() override;