Merge pull request #746 from equeim:ffmpeg-output-buffer-reallocation

copybara-github · copybara-github · commit 4f99000ba159 · 2023-10-31T09:59:35.000-07:00
PiperOrigin-RevId: 578217064
diff --git a/libraries/decoder/src/main/java/androidx/media3/decoder/SimpleDecoderOutputBuffer.java b/libraries/decoder/src/main/java/androidx/media3/decoder/SimpleDecoderOutputBuffer.java
@@ -16,6 +16,7 @@
 package androidx.media3.decoder;
 
 import androidx.annotation.Nullable;
+import androidx.media3.common.util.Assertions;
 import androidx.media3.common.util.UnstableApi;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
@@ -49,6 +50,27 @@ public ByteBuffer init(long timeUs, int size) {
     return data;
   }
 
+  /**
+   * Grows the buffer to a new size.
+   *
+   * <p>Existing data is copied to the new buffer, and {@link ByteBuffer#position} is preserved.
+   *
+   * @param newSize New size of the buffer.
+   * @return The {@link #data} buffer, for convenience.
+   */
+  public ByteBuffer grow(int newSize) {
+    ByteBuffer oldData = Assertions.checkNotNull(this.data);
+    Assertions.checkArgument(newSize >= oldData.limit());
+    ByteBuffer newData = ByteBuffer.allocateDirect(newSize).order(ByteOrder.nativeOrder());
+    int restorePosition = oldData.position();
+    oldData.position(0);
+    newData.put(oldData);
+    newData.position(restorePosition);
+    newData.limit(newSize);
+    this.data = newData;
+    return newData;
+  }
+
   @Override
   public void clear() {
     super.clear();
diff --git a/libraries/decoder_ffmpeg/proguard-rules.txt b/libraries/decoder_ffmpeg/proguard-rules.txt
@@ -4,3 +4,8 @@
 -keepclasseswithmembernames class * {
     native <methods>;
 }
+
+# This method is called from native code
+-keep class androidx.media3.decoder.ffmpeg.FfmpegAudioDecoder {
+  private java.nio.ByteBuffer growOutputBuffer(androidx.media3.decoder.SimpleDecoderOutputBuffer, int);
+}
diff --git a/libraries/decoder_ffmpeg/src/main/java/androidx/media3/decoder/ffmpeg/FfmpegAudioDecoder.java b/libraries/decoder_ffmpeg/src/main/java/androidx/media3/decoder/ffmpeg/FfmpegAudioDecoder.java
@@ -32,17 +32,16 @@
 /* package */ final class FfmpegAudioDecoder
     extends SimpleDecoder<DecoderInputBuffer, SimpleDecoderOutputBuffer, FfmpegDecoderException> {
 
-  // Output buffer sizes when decoding PCM mu-law streams, which is the maximum FFmpeg outputs.
-  private static final int OUTPUT_BUFFER_SIZE_16BIT = 65536;
-  private static final int OUTPUT_BUFFER_SIZE_32BIT = OUTPUT_BUFFER_SIZE_16BIT * 2;
+  private static final int INITIAL_OUTPUT_BUFFER_SIZE_16BIT = 65535;
+  private static final int INITIAL_OUTPUT_BUFFER_SIZE_32BIT = INITIAL_OUTPUT_BUFFER_SIZE_16BIT * 2;
 
   private static final int AUDIO_DECODER_ERROR_INVALID_DATA = -1;
   private static final int AUDIO_DECODER_ERROR_OTHER = -2;
 
   private final String codecName;
   @Nullable private final byte[] extraData;
   private final @C.PcmEncoding int encoding;
-  private final int outputBufferSize;
+  private int outputBufferSize;
 
   private long nativeContext; // May be reassigned on resetting the codec.
   private boolean hasOutputFormat;
@@ -64,7 +63,8 @@ public FfmpegAudioDecoder(
     codecName = Assertions.checkNotNull(FfmpegLibrary.getCodecName(format.sampleMimeType));
     extraData = getExtraData(format.sampleMimeType, format.initializationData);
     encoding = outputFloat ? C.ENCODING_PCM_FLOAT : C.ENCODING_PCM_16BIT;
-    outputBufferSize = outputFloat ? OUTPUT_BUFFER_SIZE_32BIT : OUTPUT_BUFFER_SIZE_16BIT;
+    outputBufferSize =
+        outputFloat ? INITIAL_OUTPUT_BUFFER_SIZE_32BIT : INITIAL_OUTPUT_BUFFER_SIZE_16BIT;
     nativeContext =
         ffmpegInitialize(codecName, extraData, outputFloat, format.sampleRate, format.channelCount);
     if (nativeContext == 0) {
@@ -108,7 +108,9 @@ protected FfmpegDecoderException decode(
     ByteBuffer inputData = Util.castNonNull(inputBuffer.data);
     int inputSize = inputData.limit();
     ByteBuffer outputData = outputBuffer.init(inputBuffer.timeUs, outputBufferSize);
-    int result = ffmpegDecode(nativeContext, inputData, inputSize, outputData, outputBufferSize);
+    int result =
+        ffmpegDecode(
+            nativeContext, inputData, inputSize, outputBuffer, outputData, outputBufferSize);
     if (result == AUDIO_DECODER_ERROR_OTHER) {
       return new FfmpegDecoderException("Error decoding (see logcat).");
     } else if (result == AUDIO_DECODER_ERROR_INVALID_DATA) {
@@ -140,6 +142,14 @@ protected FfmpegDecoderException decode(
     return null;
   }
 
+  // Called from native code
+  @SuppressWarnings("unused")
+  private ByteBuffer growOutputBuffer(SimpleDecoderOutputBuffer outputBuffer, int requiredSize) {
+    // Use it for new buffer so that hopefully we won't need to reallocate again
+    outputBufferSize = requiredSize;
+    return outputBuffer.grow(requiredSize);
+  }
+
   @Override
   public void release() {
     super.release();
@@ -221,7 +231,12 @@ private native long ffmpegInitialize(
       int rawChannelCount);
 
   private native int ffmpegDecode(
-      long context, ByteBuffer inputData, int inputSize, ByteBuffer outputData, int outputSize);
+      long context,
+      ByteBuffer inputData,
+      int inputSize,
+      SimpleDecoderOutputBuffer decoderOutputBuffer,
+      ByteBuffer outputData,
+      int outputSize);
 
   private native int ffmpegGetChannelCount(long context);
 
diff --git a/libraries/decoder_ffmpeg/src/main/jni/ffmpeg_jni.cc b/libraries/decoder_ffmpeg/src/main/jni/ffmpeg_jni.cc
@@ -35,6 +35,8 @@ extern "C" {
 #define LOG_TAG "ffmpeg_jni"
 #define LOGE(...) \
   ((void)__android_log_print(ANDROID_LOG_ERROR, LOG_TAG, __VA_ARGS__))
+#define LOGD(...) \
+  ((void)__android_log_print(ANDROID_LOG_DEBUG, LOG_TAG, __VA_ARGS__))
 
 #define LIBRARY_FUNC(RETURN_TYPE, NAME, ...)                                   \
   extern "C" {                                                                 \
@@ -67,6 +69,8 @@ static const AVSampleFormat OUTPUT_FORMAT_PCM_FLOAT = AV_SAMPLE_FMT_FLT;
 static const int AUDIO_DECODER_ERROR_INVALID_DATA = -1;
 static const int AUDIO_DECODER_ERROR_OTHER = -2;
 
+static jmethodID growOutputBufferMethod;
+
 /**
  * Returns the AVCodec with the specified name, or NULL if it is not available.
  */
@@ -81,13 +85,22 @@ AVCodecContext *createContext(JNIEnv *env, AVCodec *codec, jbyteArray extraData,
                               jboolean outputFloat, jint rawSampleRate,
                               jint rawChannelCount);
 
+struct GrowOutputBufferCallback {
+  uint8_t *operator()(int requiredSize) const;
+
+  JNIEnv *env;
+  jobject thiz;
+  jobject decoderOutputBuffer;
+};
+
 /**
  * Decodes the packet into the output buffer, returning the number of bytes
  * written, or a negative AUDIO_DECODER_ERROR constant value in the case of an
  * error.
  */
 int decodePacket(AVCodecContext *context, AVPacket *packet,
-                 uint8_t *outputBuffer, int outputSize);
+                 uint8_t *outputBuffer, int outputSize,
+                 GrowOutputBufferCallback growBuffer);
 
 /**
  * Transforms ffmpeg AVERROR into a negative AUDIO_DECODER_ERROR constant value.
@@ -107,6 +120,21 @@ void releaseContext(AVCodecContext *context);
 jint JNI_OnLoad(JavaVM *vm, void *reserved) {
   JNIEnv *env;
   if (vm->GetEnv(reinterpret_cast<void **>(&env), JNI_VERSION_1_6) != JNI_OK) {
+    LOGE("JNI_OnLoad: GetEnv failed");
+    return -1;
+  }
+  jclass clazz =
+      env->FindClass("androidx/media3/decoder/ffmpeg/FfmpegAudioDecoder");
+  if (!clazz) {
+    LOGE("JNI_OnLoad: FindClass failed");
+    return -1;
+  }
+  growOutputBufferMethod =
+      env->GetMethodID(clazz, "growOutputBuffer",
+                       "(Landroidx/media3/decoder/"
+                       "SimpleDecoderOutputBuffer;I)Ljava/nio/ByteBuffer;");
+  if (!growOutputBufferMethod) {
+    LOGE("JNI_OnLoad: GetMethodID failed");
     return -1;
   }
   avcodec_register_all();
@@ -138,12 +166,13 @@ AUDIO_DECODER_FUNC(jlong, ffmpegInitialize, jstring codecName,
 }
 
 AUDIO_DECODER_FUNC(jint, ffmpegDecode, jlong context, jobject inputData,
-                   jint inputSize, jobject outputData, jint outputSize) {
+                   jint inputSize, jobject decoderOutputBuffer,
+                   jobject outputData, jint outputSize) {
   if (!context) {
     LOGE("Context must be non-NULL.");
     return -1;
   }
-  if (!inputData || !outputData) {
+  if (!inputData || !decoderOutputBuffer || !outputData) {
     LOGE("Input and output buffers must be non-NULL.");
     return -1;
   }
@@ -162,7 +191,19 @@ AUDIO_DECODER_FUNC(jint, ffmpegDecode, jlong context, jobject inputData,
   packet.data = inputBuffer;
   packet.size = inputSize;
   return decodePacket((AVCodecContext *)context, &packet, outputBuffer,
-                      outputSize);
+                      outputSize,
+                      GrowOutputBufferCallback{env, thiz, decoderOutputBuffer});
+}
+
+uint8_t *GrowOutputBufferCallback::operator()(int requiredSize) const {
+  jobject newOutputData = env->CallObjectMethod(
+      thiz, growOutputBufferMethod, decoderOutputBuffer, requiredSize);
+  if (env->ExceptionCheck()) {
+    LOGE("growOutputBuffer() failed");
+    env->ExceptionDescribe();
+    return nullptr;
+  }
+  return static_cast<uint8_t *>(env->GetDirectBufferAddress(newOutputData));
 }
 
 AUDIO_DECODER_FUNC(jint, ffmpegGetChannelCount, jlong context) {
@@ -264,7 +305,8 @@ AVCodecContext *createContext(JNIEnv *env, AVCodec *codec, jbyteArray extraData,
 }
 
 int decodePacket(AVCodecContext *context, AVPacket *packet,
-                 uint8_t *outputBuffer, int outputSize) {
+                 uint8_t *outputBuffer, int outputSize,
+                 GrowOutputBufferCallback growBuffer) {
   int result = 0;
   // Queue input data.
   result = avcodec_send_packet(context, packet);
@@ -320,15 +362,22 @@ int decodePacket(AVCodecContext *context, AVPacket *packet,
       }
       context->opaque = resampleContext;
     }
-    int inSampleSize = av_get_bytes_per_sample(sampleFormat);
+
     int outSampleSize = av_get_bytes_per_sample(context->request_sample_fmt);
     int outSamples = swr_get_out_samples(resampleContext, sampleCount);
     int bufferOutSize = outSampleSize * channelCount * outSamples;
     if (outSize + bufferOutSize > outputSize) {
-      LOGE("Output buffer size (%d) too small for output data (%d).",
-           outputSize, outSize + bufferOutSize);
-      av_frame_free(&frame);
-      return AUDIO_DECODER_ERROR_INVALID_DATA;
+      LOGD(
+          "Output buffer size (%d) too small for output data (%d), "
+          "reallocating buffer.",
+          outputSize, outSize + bufferOutSize);
+      outputSize = outSize + bufferOutSize;
+      outputBuffer = growBuffer(outputSize);
+      if (!outputBuffer) {
+        LOGE("Failed to reallocate output buffer.");
+        av_frame_free(&frame);
+        return AUDIO_DECODER_ERROR_OTHER;
+      }
     }
     result = swr_convert(resampleContext, &outputBuffer, bufferOutSize,
                          (const uint8_t **)frame->data, frame->nb_samples);

Original file line number	Diff line number	Diff line change
`@@ -4,3 +4,8 @@`
`4`	`4`	`-keepclasseswithmembernames class * {`
`5`	`5`	`native <methods>;`
`6`	`6`	`}`
	`7`	`+`
	`8`	`+# This method is called from native code`
	`9`	`+-keep class androidx.media3.decoder.ffmpeg.FfmpegAudioDecoder {`
	`10`	`+ private java.nio.ByteBuffer growOutputBuffer(androidx.media3.decoder.SimpleDecoderOutputBuffer, int);`
	`11`	`+}`