ffmpeg: reallocate output buffer dynamically

With FFmpeg we can't determine size of output buffer ahead of time for all codecs,
so we need to reallocate it when needed instead of simply failing.

Author: equeim

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,18 @@ public ByteBuffer init(long timeUs, int size) {
     return data;
   }
 
+  public ByteBuffer grow(int newSize) {
+    Assertions.checkNotNull(data);
+    final ByteBuffer newData = ByteBuffer.allocateDirect(newSize).order(ByteOrder.nativeOrder());
+    final int restorePosition = data.position();
+    data.position(0);
+    newData.put(data);
+    newData.position(restorePosition);
+    newData.limit(newSize);
+    data = newData;
+    return newData;
+  }
+
   @Override
   public void clear() {
     super.clear();

libraries/decoder_ffmpeg/src/main/java/androidx/media3/decoder/ffmpeg/FfmpegAudioDecoder.java

@@ -15,6 +15,7 @@
  */
 package androidx.media3.decoder.ffmpeg;
 
+import androidx.annotation.Keep;
 import androidx.annotation.Nullable;
 import androidx.media3.common.C;
 import androidx.media3.common.Format;
@@ -27,22 +28,23 @@
 import androidx.media3.decoder.SimpleDecoderOutputBuffer;
 import java.nio.ByteBuffer;
 import java.util.List;
+import java.util.function.Function;
+import java.util.function.IntConsumer;
 
 /** FFmpeg audio decoder. */
 /* 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 +66,7 @@ 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) {
@@ -107,8 +109,10 @@ 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);
+    outputBuffer.init(inputBuffer.timeUs, outputBufferSize);
+    Assertions.checkNotNull(outputBuffer.data);
+
+    int result = ffmpegDecode(nativeContext, inputData, inputSize, outputBuffer, outputBuffer.data, outputBufferSize);
     if (result == AUDIO_DECODER_ERROR_OTHER) {
       return new FfmpegDecoderException("Error decoding (see logcat).");
     } else if (result == AUDIO_DECODER_ERROR_INVALID_DATA) {
@@ -135,11 +139,19 @@ protected FfmpegDecoderException decode(
       }
       hasOutputFormat = true;
     }
-    outputData.position(0);
-    outputData.limit(result);
+    outputBuffer.data.position(0);
+    outputBuffer.data.limit(result);
     return null;
   }
 
+  // Called from native code
+  @Keep
+  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 +233,7 @@ 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);
 

libraries/decoder_ffmpeg/src/main/jni/ffmpeg_jni.cc

@@ -67,6 +67,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 +83,21 @@ 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 +117,17 @@ 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 +159,12 @@ 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 +183,17 @@ 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 +295,7 @@ 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 +351,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).",
+      LOGE("Output buffer size (%d) too small for output data (%d), reallocating buffer.",
            outputSize, outSize + bufferOutSize);
-      av_frame_free(&frame);
-      return AUDIO_DECODER_ERROR_INVALID_DATA;
+      outputSize = outSize + bufferOutSize;
+      outputBuffer = growBuffer(outputSize);
+      if (outputBuffer) {
+        LOGE("Reallocated output buffer.");
+      } else {
+        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);