Reduce memory allocations of AnnexBToAvccConverter

Add a new ByteBufferAllocator interface, with a simple
LinearByteBufferAllocator implementation that supports basic memory
allocations and reuse.

Add an AnnexBToAvccConverter.process override that takes a custom allocator

PiperOrigin-RevId: 726380184

libraries/muxer/src/main/java/androidx/media3/muxer/AnnexBToAvccConverter.java

@@ -28,32 +28,40 @@ import java.nio.ByteBuffer;
 public interface AnnexBToAvccConverter {
   /** Default implementation for {@link AnnexBToAvccConverter}. */
   AnnexBToAvccConverter DEFAULT =
-      (ByteBuffer inputBuffer) -> {
-        if (!inputBuffer.hasRemaining()) {
-          return inputBuffer;
+      new AnnexBToAvccConverter() {
+        @Override
+        public ByteBuffer process(ByteBuffer inputBuffer) {
+          return process(inputBuffer, ByteBufferAllocator.DEFAULT);
         }
 
-        ImmutableList<ByteBuffer> nalUnitList = AnnexBUtils.findNalUnits(inputBuffer);
+        @Override
+        public ByteBuffer process(ByteBuffer inputBuffer, ByteBufferAllocator byteBufferAllocator) {
+          if (!inputBuffer.hasRemaining()) {
+            return inputBuffer;
+          }
 
-        int totalBytesNeeded = 0;
+          ImmutableList<ByteBuffer> nalUnitList = AnnexBUtils.findNalUnits(inputBuffer);
 
-        for (int i = 0; i < nalUnitList.size(); i++) {
-          // 4 bytes to store NAL unit length.
-          totalBytesNeeded += 4 + nalUnitList.get(i).remaining();
+          int totalBytesNeeded = 0;
+
+          for (int i = 0; i < nalUnitList.size(); i++) {
+            // 4 bytes to store NAL unit length.
+            totalBytesNeeded += 4 + nalUnitList.get(i).remaining();
+          }
+
+          ByteBuffer outputBuffer = byteBufferAllocator.allocate(totalBytesNeeded);
+
+          for (int i = 0; i < nalUnitList.size(); i++) {
+            ByteBuffer currentNalUnit = nalUnitList.get(i);
+            int currentNalUnitLength = currentNalUnit.remaining();
+
+            // Rewrite NAL units with NAL unit length in place of start code.
+            outputBuffer.putInt(currentNalUnitLength);
+            outputBuffer.put(currentNalUnit);
+          }
+          outputBuffer.rewind();
+          return outputBuffer;
         }
-
-        ByteBuffer outputBuffer = ByteBuffer.allocate(totalBytesNeeded);
-
-        for (int i = 0; i < nalUnitList.size(); i++) {
-          ByteBuffer currentNalUnit = nalUnitList.get(i);
-          int currentNalUnitLength = currentNalUnit.remaining();
-
-          // Rewrite NAL units with NAL unit length in place of start code.
-          outputBuffer.putInt(currentNalUnitLength);
-          outputBuffer.put(currentNalUnit);
-        }
-        outputBuffer.rewind();
-        return outputBuffer;
       };
 
   /**
@@ -64,4 +72,16 @@ public interface AnnexBToAvccConverter {
    * @param inputBuffer The buffer to be converted.
    */
   ByteBuffer process(ByteBuffer inputBuffer);
+
+  /**
+   * Returns the processed {@link ByteBuffer}.
+   *
+   * <p>Expects a {@link ByteBuffer} input with a zero offset.
+   *
+   * @param inputBuffer The buffer to be converted.
+   * @param allocator An allocator for {@link ByteBuffer} instances that enables memory reuse.
+   */
+  default ByteBuffer process(ByteBuffer inputBuffer, ByteBufferAllocator allocator) {
+    return process(inputBuffer);
+  }
 }

libraries/muxer/src/main/java/androidx/media3/muxer/ByteBufferAllocator.java

@@ -0,0 +1,34 @@
+/*
+ * Copyright 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package androidx.media3.muxer;
+
+import androidx.media3.common.util.UnstableApi;
+import java.nio.ByteBuffer;
+
+/** A memory allocator for {@link ByteBuffer}. */
+@UnstableApi
+public interface ByteBufferAllocator {
+  /** Default implementation. */
+  ByteBufferAllocator DEFAULT = ByteBuffer::allocateDirect;
+
+  /**
+   * Allocates and returns a new {@link ByteBuffer}.
+   *
+   * @param capacity The new buffer's capacity, in bytes.
+   * @throws IllegalArgumentException If the {@code capacity} is a negative integer.
+   */
+  ByteBuffer allocate(int capacity);
+}

libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java

@@ -115,6 +115,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   private final boolean sampleCopyEnabled;
   private final @Mp4Muxer.LastSampleDurationBehavior int lastSampleDurationBehavior;
   private final List<Track> tracks;
+  private final LinearByteBufferAllocator linearByteBufferAllocator;
 
   private @MonotonicNonNull Track videoTrack;
   private int currentFragmentSequenceNumber;
@@ -151,6 +152,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     tracks = new ArrayList<>();
     minInputPresentationTimeUs = Long.MAX_VALUE;
     currentFragmentSequenceNumber = 1;
+    linearByteBufferAllocator = new LinearByteBufferAllocator(/* initialCapacity= */ 0);
   }
 
   public Track addTrack(int sortKey, Format format) {
@@ -325,6 +327,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
         outputChannel.write(currentTrackInfo.pendingSamplesByteBuffer.get(sampleIndex));
       }
     }
+    linearByteBufferAllocator.reset();
   }
 
   private ImmutableList<ProcessedTrackInfo> processAllTracks() {
@@ -346,7 +349,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     if (doesSampleContainAnnexBNalUnits(checkNotNull(track.format.sampleMimeType))) {
       while (!track.pendingSamplesByteBuffer.isEmpty()) {
         ByteBuffer currentSampleByteBuffer = track.pendingSamplesByteBuffer.removeFirst();
-        currentSampleByteBuffer = annexBToAvccConverter.process(currentSampleByteBuffer);
+        currentSampleByteBuffer =
+            annexBToAvccConverter.process(currentSampleByteBuffer, linearByteBufferAllocator);
         pendingSamplesByteBuffer.add(currentSampleByteBuffer);
         BufferInfo currentSampleBufferInfo = track.pendingSamplesBufferInfo.removeFirst();
         currentSampleBufferInfo.set(

libraries/muxer/src/main/java/androidx/media3/muxer/LinearByteBufferAllocator.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package androidx.media3.muxer;
+
+import static androidx.media3.common.util.Assertions.checkArgument;
+import static java.lang.Math.max;
+
+import java.nio.ByteBuffer;
+
+/** A simple linear allocator for {@link ByteBuffer} instances. */
+/* package */ final class LinearByteBufferAllocator implements ByteBufferAllocator {
+
+  private ByteBuffer memoryPool;
+
+  /**
+   * Creates a new instance.
+   *
+   * @param initialCapacity The initial capacity reserved by the linear allocator.
+   */
+  public LinearByteBufferAllocator(int initialCapacity) {
+    checkArgument(initialCapacity >= 0);
+    memoryPool = ByteBuffer.allocateDirect(initialCapacity);
+  }
+
+  @Override
+  public ByteBuffer allocate(int capacity) {
+    checkArgument(capacity >= 0);
+    if (memoryPool.remaining() < capacity) {
+      int newCapacity = max(capacity, memoryPool.capacity() * 2);
+      memoryPool = ByteBuffer.allocateDirect(newCapacity);
+    }
+    ByteBuffer outputBuffer = memoryPool.slice();
+    memoryPool.position(memoryPool.position() + capacity);
+    outputBuffer.limit(capacity);
+
+    return outputBuffer;
+  }
+
+  /** Frees all previously allocated memory and returns it to the allocator. */
+  public void reset() {
+    memoryPool.clear();
+  }
+}

libraries/muxer/src/main/java/androidx/media3/muxer/Mp4Writer.java

@@ -57,6 +57,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
   private final List<Track> tracks;
   private final List<Track> auxiliaryTracks;
   private final AtomicBoolean hasWrittenSamples;
+  private final LinearByteBufferAllocator linearByteBufferAllocator;
 
   // Stores location of the space reserved for the moov box at the beginning of the file (after ftyp
   // box)
@@ -106,6 +107,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
     canWriteMoovAtStart = attemptStreamableOutputEnabled;
     lastMoovWritten = Range.closed(0L, 0L);
     lastMoovWrittenAtSampleTimestampUs = 0L;
+    linearByteBufferAllocator = new LinearByteBufferAllocator(/* initialCapacity= */ 0);
   }
 
   /**
@@ -459,7 +461,8 @@ import java.util.concurrent.atomic.AtomicBoolean;
       // Convert the H.264/H.265 samples from Annex-B format (output by MediaCodec) to
       // Avcc format (required by MP4 container).
       if (doesSampleContainAnnexBNalUnits(checkNotNull(track.format.sampleMimeType))) {
-        currentSampleByteBuffer = annexBToAvccConverter.process(currentSampleByteBuffer);
+        currentSampleByteBuffer =
+            annexBToAvccConverter.process(currentSampleByteBuffer, linearByteBufferAllocator);
         currentSampleBufferInfo.set(
             currentSampleByteBuffer.position(),
             currentSampleByteBuffer.remaining(),
@@ -472,6 +475,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
       maybeExtendMdatAndRewriteMoov(currentSampleByteBuffer.remaining());
 
       mdatDataEnd += outputFileChannel.write(currentSampleByteBuffer, mdatDataEnd);
+      linearByteBufferAllocator.reset();
       track.writtenSamples.add(currentSampleBufferInfo);
     } while (!track.pendingSamplesBufferInfo.isEmpty());
     checkState(mdatDataEnd <= mdatEnd);