Internal change

PiperOrigin-RevId: 495236576

mediapipe/java/com/google/mediapipe/framework/GraphTextureFrame.java

@@ -36,6 +36,7 @@ public class GraphTextureFrame implements TextureFrame {
   // when calling getTextureName().
   private final boolean deferredSync;
   private final Set<Long> activeConsumerContextHandleSet = new HashSet<>();
+  private int refCount = 1;
 
   GraphTextureFrame(long nativeHandle, long timestamp) {
     this(nativeHandle, timestamp, false);
@@ -94,6 +95,17 @@ public class GraphTextureFrame implements TextureFrame {
     return timestamp;
   }
 
+  @Override
+  public boolean supportsRetain() {
+    return true;
+  }
+
+  @Override
+  public synchronized void retain() {
+    // TODO: check that refCount is > 0 and handle is not 0.
+    refCount++;
+  }
+
   /**
    * Releases a reference to the underlying buffer.
    *
@@ -121,22 +133,32 @@ public class GraphTextureFrame implements TextureFrame {
    * currently cannot create a GlSyncToken, so they cannot call this method.
    */
   @Override
-  public void release(GlSyncToken consumerSyncToken) {
-    if (nativeBufferHandle != 0) {
-      long token = consumerSyncToken == null ? 0 : consumerSyncToken.nativeToken();
-      nativeReleaseBuffer(nativeBufferHandle, token);
-      nativeBufferHandle = 0;
-    } else if (consumerSyncToken != null) {
+  public synchronized void release(GlSyncToken consumerSyncToken) {
+    if (nativeBufferHandle == 0) {
+      if (consumerSyncToken != null) {
         logger.atWarning().log("release with sync token, but handle is 0");
       }
+      return;
+    }
+
     if (consumerSyncToken != null) {
+      long token = consumerSyncToken.nativeToken();
+      nativeDidRead(nativeBufferHandle, token);
+      // We should remove the token's context from activeConsumerContextHandleSet here, but for now
+      // we do it in the release(void) overload.
       consumerSyncToken.release();
     }
+
+    refCount--;
+    if (refCount <= 0) {
+      nativeReleaseBuffer(nativeBufferHandle);
+      nativeBufferHandle = 0;
+    }
   }
 
   @Override
   protected void finalize() throws Throwable {
-    if (nativeBufferHandle != 0) {
+    if (refCount >= 0 || nativeBufferHandle != 0) {
       logger.atWarning().log("release was not called before finalize");
     }
     if (!activeConsumerContextHandleSet.isEmpty()) {
@@ -144,7 +166,7 @@ public class GraphTextureFrame implements TextureFrame {
     }
   }
 
-  private native void nativeReleaseBuffer(long nativeHandle, long consumerSyncToken);
+  private native void nativeReleaseBuffer(long nativeHandle);
 
   private native int nativeGetTextureName(long nativeHandle);
   private native int nativeGetWidth(long nativeHandle);
@@ -155,4 +177,6 @@ public class GraphTextureFrame implements TextureFrame {
   private native long nativeCreateSyncTokenForCurrentExternalContext(long nativeHandle);
 
   private native long nativeGetCurrentExternalContextHandle();
+
+  private native void nativeDidRead(long nativeHandle, long consumerSyncToken);
 }

mediapipe/java/com/google/mediapipe/framework/TextureFrame.java

@@ -59,4 +59,18 @@ public interface TextureFrame extends TextureReleaseCallback {
    */
   @Override
   void release(GlSyncToken syncToken);
+
+  /**
+   * If this method returns true, this object supports the retain method, and can be used with
+   * multiple consumers. Call retain for each additional consumer beyond the first; each consumer
+   * should call release.
+   */
+  default boolean supportsRetain() {
+    return false;
+  }
+
+  /** Increments the reference count. Only available with some implementations of TextureFrame. */
+  default void retain() {
+    throw new UnsupportedOperationException();
+  }
 }

mediapipe/java/com/google/mediapipe/framework/jni/graph_texture_frame_jni.cc

@@ -22,14 +22,9 @@
 using mediapipe::GlTextureBufferSharedPtr;
 
 JNIEXPORT void JNICALL GRAPH_TEXTURE_FRAME_METHOD(nativeReleaseBuffer)(
-    JNIEnv* env, jobject thiz, jlong nativeHandle, jlong consumerSyncToken) {
+    JNIEnv* env, jobject thiz, jlong nativeHandle) {
   GlTextureBufferSharedPtr* buffer =
       reinterpret_cast<GlTextureBufferSharedPtr*>(nativeHandle);
-  if (consumerSyncToken) {
-    mediapipe::GlSyncToken& token =
-        *reinterpret_cast<mediapipe::GlSyncToken*>(consumerSyncToken);
-    (*buffer)->DidRead(token);
-  }
   delete buffer;
 }
 
@@ -91,3 +86,12 @@ JNIEXPORT jlong JNICALL GRAPH_TEXTURE_FRAME_METHOD(
   return reinterpret_cast<jlong>(
       mediapipe::GlContext::GetCurrentNativeContext());
 }
+
+JNIEXPORT void JNICALL GRAPH_TEXTURE_FRAME_METHOD(nativeDidRead)(
+    JNIEnv* env, jobject thiz, jlong nativeHandle, jlong consumerSyncToken) {
+  GlTextureBufferSharedPtr* buffer =
+      reinterpret_cast<GlTextureBufferSharedPtr*>(nativeHandle);
+  mediapipe::GlSyncToken& token =
+      *reinterpret_cast<mediapipe::GlSyncToken*>(consumerSyncToken);
+  (*buffer)->DidRead(token);
+}

mediapipe/java/com/google/mediapipe/framework/jni/graph_texture_frame_jni.h

@@ -26,7 +26,7 @@ extern "C" {
 
 // Releases a native mediapipe::GpuBuffer.
 JNIEXPORT void JNICALL GRAPH_TEXTURE_FRAME_METHOD(nativeReleaseBuffer)(
-    JNIEnv* env, jobject thiz, jlong nativeHandle, jlong consumerSyncToken);
+    JNIEnv* env, jobject thiz, jlong nativeHandle);
 
 JNIEXPORT jint JNICALL GRAPH_TEXTURE_FRAME_METHOD(nativeGetTextureName)(
     JNIEnv* env, jobject thiz, jlong nativeHandle);
@@ -44,6 +44,9 @@ JNIEXPORT jlong JNICALL GRAPH_TEXTURE_FRAME_METHOD(
     nativeCreateSyncTokenForCurrentExternalContext)(JNIEnv* env, jobject thiz,
                                                     jlong nativeHandle);
 
+JNIEXPORT void JNICALL GRAPH_TEXTURE_FRAME_METHOD(nativeDidRead)(
+    JNIEnv* env, jobject thiz, jlong nativeHandle, jlong consumerSyncToken);
+
 JNIEXPORT jlong JNICALL GRAPH_TEXTURE_FRAME_METHOD(
     nativeGetCurrentExternalContextHandle)(JNIEnv* env, jobject thiz);