Add release callback and reference count to java EncodedImage class

Callback set by HardwareVideoEncoder, and wired to the codec's
releaseOutputBuffer. Intention is to move call of this method to the
destructor of a corresponding C++ class in a followup cl, and
eliminate an allocation and memcpy in the process.

Bug: webrtc:9378
Change-Id: I578480b63b68e6ac7a96cdde36379b3c50f05c3f
Reviewed-on: https://webrtc-review.googlesource.com/c/src/+/142160
Commit-Queue: Niels Moller <nisse@webrtc.org>
Reviewed-by: Sami Kalliomäki <sakal@webrtc.org>
Cr-Commit-Position: refs/heads/master@{#29283}
diff --git a/sdk/android/api/org/webrtc/EncodedImage.java b/sdk/android/api/org/webrtc/EncodedImage.java
index f8de3d2..682d9c4 100644
--- a/sdk/android/api/org/webrtc/EncodedImage.java
+++ b/sdk/android/api/org/webrtc/EncodedImage.java
@@ -18,7 +18,7 @@
  * An encoded frame from a video stream. Used as an input for decoders and as an output for
  * encoders.
  */
-public class EncodedImage {
+public class EncodedImage implements RefCounted {
   // Must be kept in sync with common_types.h FrameType.
   public enum FrameType {
     EmptyFrame(0),
@@ -46,6 +46,8 @@
     }
   }
 
+  private final RefCountDelegate refCountDelegate;
+  private final boolean supportsRetain;
   public final ByteBuffer buffer;
   public final int encodedWidth;
   public final int encodedHeight;
@@ -56,8 +58,31 @@
   public final boolean completeFrame;
   public final @Nullable Integer qp;
 
+  // TODO(bugs.webrtc.org/9378): Use retain and release from jni code.
+  @Override
+  public void retain() {
+    refCountDelegate.retain();
+  }
+
+  @Override
+  public void release() {
+    refCountDelegate.release();
+  }
+
+  // A false return value means that the encoder expects that the buffer is no longer used after
+  // VideoEncoder.Callback.onEncodedFrame returns.
+  boolean maybeRetain() {
+    if (supportsRetain) {
+      retain();
+      return true;
+    } else {
+      return false;
+    }
+  }
+
   @CalledByNative
-  private EncodedImage(ByteBuffer buffer, int encodedWidth, int encodedHeight, long captureTimeNs,
+  private EncodedImage(ByteBuffer buffer, boolean supportsRetain,
+      @Nullable Runnable releaseCallback, int encodedWidth, int encodedHeight, long captureTimeNs,
       FrameType frameType, int rotation, boolean completeFrame, @Nullable Integer qp) {
     this.buffer = buffer;
     this.encodedWidth = encodedWidth;
@@ -68,6 +93,8 @@
     this.rotation = rotation;
     this.completeFrame = completeFrame;
     this.qp = qp;
+    this.supportsRetain = supportsRetain;
+    this.refCountDelegate = new RefCountDelegate(releaseCallback);
   }
 
   @CalledByNative
@@ -116,6 +143,8 @@
 
   public static class Builder {
     private ByteBuffer buffer;
+    private boolean supportsRetain;
+    private @Nullable Runnable releaseCallback;
     private int encodedWidth;
     private int encodedHeight;
     private long captureTimeNs;
@@ -126,8 +155,18 @@
 
     private Builder() {}
 
+    @Deprecated
     public Builder setBuffer(ByteBuffer buffer) {
       this.buffer = buffer;
+      this.releaseCallback = null;
+      this.supportsRetain = false;
+      return this;
+    }
+
+    public Builder setBuffer(ByteBuffer buffer, @Nullable Runnable releaseCallback) {
+      this.buffer = buffer;
+      this.releaseCallback = releaseCallback;
+      this.supportsRetain = true;
       return this;
     }
 
@@ -173,8 +212,8 @@
     }
 
     public EncodedImage createEncodedImage() {
-      return new EncodedImage(buffer, encodedWidth, encodedHeight, captureTimeNs, frameType,
-          rotation, completeFrame, qp);
+      return new EncodedImage(buffer, supportsRetain, releaseCallback, encodedWidth, encodedHeight,
+          captureTimeNs, frameType, rotation, completeFrame, qp);
     }
   }
 }
diff --git a/sdk/android/api/org/webrtc/VideoEncoder.java b/sdk/android/api/org/webrtc/VideoEncoder.java
index 79b5bf3..cfc131f 100644
--- a/sdk/android/api/org/webrtc/VideoEncoder.java
+++ b/sdk/android/api/org/webrtc/VideoEncoder.java
@@ -238,8 +238,14 @@
 
   public interface Callback {
     /**
-     * Call to return an encoded frame. It is safe to assume the byte buffer held by |frame| is not
-     * accessed after the call to this method returns.
+     * Old encoders assume that the byte buffer held by |frame| is not accessed after the call to
+     * this method returns. If the pipeline downstream needs to hold on to the buffer, it then has
+     * to make its own copy. We want to move to a model where no copying is needed, and instead use
+     * retain()/release() to signal to the encoder when it is safe to reuse the buffer.
+     *
+     * Over the transition, implementations of this class should use the maybeRetain() method if
+     * they want to keep a reference to the buffer, and fall back to copying if that method returns
+     * false.
      */
     void onEncodedFrame(EncodedImage frame, CodecSpecificInfo info);
   }
diff --git a/sdk/android/src/java/org/webrtc/HardwareVideoEncoder.java b/sdk/android/src/java/org/webrtc/HardwareVideoEncoder.java
index c5620bc..d14e764 100644
--- a/sdk/android/src/java/org/webrtc/HardwareVideoEncoder.java
+++ b/sdk/android/src/java/org/webrtc/HardwareVideoEncoder.java
@@ -56,6 +56,52 @@
   private static final int MEDIA_CODEC_RELEASE_TIMEOUT_MS = 5000;
   private static final int DEQUEUE_OUTPUT_BUFFER_TIMEOUT_US = 100000;
 
+  /**
+   * Keeps track of the number of output buffers that have been passed down the pipeline and not yet
+   * released. We need to wait for this to go down to zero before operations invalidating the output
+   * buffers, i.e., stop() and getOutputBuffers().
+   */
+  private static class BusyCount {
+    private final Object countLock = new Object();
+    private int count;
+
+    public void increment() {
+      synchronized (countLock) {
+        count++;
+      }
+    }
+
+    // This method may be called on an arbitrary thread.
+    public void decrement() {
+      synchronized (countLock) {
+        count--;
+        if (count == 0) {
+          countLock.notifyAll();
+        }
+      }
+    }
+
+    // The increment and waitForZero methods are called on the same thread (deliverEncodedImage,
+    // running on the output thread). Hence, after waitForZero returns, the count will stay zero
+    // until the same thread calls increment.
+    public void waitForZero() {
+      boolean wasInterrupted = false;
+      synchronized (countLock) {
+        while (count > 0) {
+          try {
+            countLock.wait();
+          } catch (InterruptedException e) {
+            Logging.e(TAG, "Interrupted while waiting on busy count", e);
+            wasInterrupted = true;
+          }
+        }
+      }
+
+      if (wasInterrupted) {
+        Thread.currentThread().interrupt();
+      }
+    }
+  }
   // --- Initialized on construction.
   private final MediaCodecWrapperFactory mediaCodecWrapperFactory;
   private final String codecName;
@@ -81,6 +127,7 @@
 
   private final ThreadChecker encodeThreadChecker = new ThreadChecker();
   private final ThreadChecker outputThreadChecker = new ThreadChecker();
+  private final BusyCount outputBuffersBusyCount = new BusyCount();
 
   // --- Set on initialize and immutable until release.
   private Callback callback;
@@ -492,6 +539,7 @@
       int index = codec.dequeueOutputBuffer(info, DEQUEUE_OUTPUT_BUFFER_TIMEOUT_US);
       if (index < 0) {
         if (index == MediaCodec.INFO_OUTPUT_BUFFERS_CHANGED) {
+          outputBuffersBusyCount.waitForZero();
           outputBuffers = codec.getOutputBuffers();
         }
         return;
@@ -535,12 +583,21 @@
             ? EncodedImage.FrameType.VideoFrameKey
             : EncodedImage.FrameType.VideoFrameDelta;
 
+        outputBuffersBusyCount.increment();
         EncodedImage.Builder builder = outputBuilders.poll();
-        builder.setBuffer(frameBuffer).setFrameType(frameType);
+        EncodedImage encodedImage = builder
+                                        .setBuffer(frameBuffer,
+                                            () -> {
+                                              codec.releaseOutputBuffer(index, false);
+                                              outputBuffersBusyCount.decrement();
+                                            })
+                                        .setFrameType(frameType)
+                                        .createEncodedImage();
         // TODO(mellem):  Set codec-specific info.
-        callback.onEncodedFrame(builder.createEncodedImage(), new CodecSpecificInfo());
+        callback.onEncodedFrame(encodedImage, new CodecSpecificInfo());
+        // Note that the callback may have retained the image.
+        encodedImage.release();
       }
-      codec.releaseOutputBuffer(index, false);
     } catch (IllegalStateException e) {
       Logging.e(TAG, "deliverOutput failed", e);
     }
@@ -549,6 +606,7 @@
   private void releaseCodecOnOutputThread() {
     outputThreadChecker.checkIsOnValidThread();
     Logging.d(TAG, "Releasing MediaCodec on output thread");
+    outputBuffersBusyCount.waitForZero();
     try {
       codec.stop();
     } catch (Exception e) {
diff --git a/sdk/android/src/jni/encoded_image.cc b/sdk/android/src/jni/encoded_image.cc
index 12dc5fe..0cc0a5b 100644
--- a/sdk/android/src/jni/encoded_image.cc
+++ b/sdk/android/src/jni/encoded_image.cc
@@ -34,8 +34,12 @@
   ScopedJavaLocalRef<jobject> qp;
   if (image.qp_ != -1)
     qp = NativeToJavaInteger(jni, image.qp_);
+  // TODO(bugs.webrtc.org/9378): Keep a reference to the C++ EncodedImage data,
+  // and use the releaseCallback to manage lifetime.
   return Java_EncodedImage_Constructor(
-      jni, buffer, static_cast<int>(image._encodedWidth),
+      jni, buffer, /*supportsRetain=*/true,
+      /*releaseCallback=*/ScopedJavaGlobalRef<jobject>(nullptr),
+      static_cast<int>(image._encodedWidth),
       static_cast<int>(image._encodedHeight),
       image.capture_time_ms_ * rtc::kNumNanosecsPerMillisec, frame_type,
       static_cast<jint>(image.rotation_), image._completeFrame, qp);