Add API to disable sample copy in Mp4Muxer

PiperOrigin-RevId: 627002635

libraries/muxer/src/androidTest/java/androidx/media3/muxer/FragmentedMp4MuxerEndToEndAndroidTest.java

@@ -66,7 +66,7 @@ public class FragmentedMp4MuxerEndToEndAndroidTest {
     @Nullable Muxer fragmentedMp4Muxer = null;
 
     try {
-      fragmentedMp4Muxer = new FragmentedMp4Muxer(checkNotNull(outputStream));
+      fragmentedMp4Muxer = new FragmentedMp4Muxer.Builder(checkNotNull(outputStream)).build();
       fragmentedMp4Muxer.addMetadata(
           new Mp4TimestampData(
               /* creationTimestampSeconds= */ 100_000_000L,
@@ -93,7 +93,7 @@ public class FragmentedMp4MuxerEndToEndAndroidTest {
     @Nullable Muxer fragmentedMp4Muxer = null;
 
     try {
-      fragmentedMp4Muxer = new FragmentedMp4Muxer(checkNotNull(outputStream));
+      fragmentedMp4Muxer = new FragmentedMp4Muxer.Builder(checkNotNull(outputStream)).build();
       fragmentedMp4Muxer.addMetadata(
           new Mp4TimestampData(
               /* creationTimestampSeconds= */ 100_000_000L,

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

@@ -18,7 +18,7 @@ package androidx.media3.muxer;
 import static androidx.media3.common.util.Assertions.checkArgument;
 import static androidx.media3.common.util.Assertions.checkNotNull;
 
-import android.media.MediaCodec;
+import android.media.MediaCodec.BufferInfo;
 import androidx.media3.common.Format;
 import androidx.media3.common.Metadata;
 import androidx.media3.common.util.UnstableApi;
@@ -27,6 +27,7 @@ import androidx.media3.container.Mp4LocationData;
 import androidx.media3.container.Mp4OrientationData;
 import androidx.media3.container.Mp4TimestampData;
 import androidx.media3.container.XmpData;
+import com.google.errorprone.annotations.CanIgnoreReturnValue;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -43,8 +44,8 @@ import java.nio.ByteBuffer;
  * <ul>
  *   <li>Add tracks using {@link #addTrack(Format)} which will return a {@link Mp4Muxer.TrackToken}.
  *   <li>Use the associated {@link Mp4Muxer.TrackToken} when {@linkplain
- *       #writeSampleData(Mp4Muxer.TrackToken, ByteBuffer, MediaCodec.BufferInfo) writing samples}
+ *       #writeSampleData(Mp4Muxer.TrackToken, ByteBuffer, BufferInfo) writing samples} for that
- *       for that track.
+ *       track.
  *   <li>{@link #close} the muxer when all data has been written.
  * </ul>
  *
@@ -54,7 +55,7 @@ import java.nio.ByteBuffer;
  *   <li>All tracks must be added before writing any samples.
  *   <li>The caller is responsible for ensuring that samples of different track types are well
  *       interleaved by calling {@link #writeSampleData(Mp4Muxer.TrackToken, ByteBuffer,
- *       MediaCodec.BufferInfo)} in an order that interleaves samples from different tracks.
+ *       BufferInfo)} in an order that interleaves samples from different tracks.
  * </ul>
  */
 @UnstableApi
@@ -62,23 +63,65 @@ public final class FragmentedMp4Muxer implements Muxer {
   /** The default fragment duration. */
   public static final long DEFAULT_FRAGMENT_DURATION_MS = 2_000;
 
-  private final FragmentedMp4Writer fragmentedMp4Writer;
+  /** A builder for {@link FragmentedMp4Muxer} instances. */
-  private final MetadataCollector metadataCollector;
+  public static final class Builder {
+    private final FileOutputStream fileOutputStream;
 
-  /** Creates an instance with {@link #DEFAULT_FRAGMENT_DURATION_MS}. */
+    private long fragmentDurationMs;
-  public FragmentedMp4Muxer(FileOutputStream fileOutputStream) {
+    private boolean sampleCopyEnabled;
-    this(fileOutputStream, DEFAULT_FRAGMENT_DURATION_MS);
+
+    /**
+     * Creates a {@link Builder} instance with default values.
+     *
+     * @param fileOutputStream The {@link FileOutputStream} to write the media data to.
+     */
+    public Builder(FileOutputStream fileOutputStream) {
+      this.fileOutputStream = fileOutputStream;
+      fragmentDurationMs = DEFAULT_FRAGMENT_DURATION_MS;
+      sampleCopyEnabled = true;
     }
 
     /**
-   * Creates an instance.
+     * Sets the fragment duration (in milliseconds).
      *
-   * @param fileOutputStream The {@link FileOutputStream} to write the media data to.
+     * <p>The muxer will attempt to create fragments of the given duration but the actual duration
-   * @param fragmentDurationMs The fragment duration (in milliseconds). The muxer will attempt to
+     * might be greater depending upon the frequency of sync samples.
-   *     create fragments of the given duration but the actual duration might be greater depending
+     *
-   *     upon the frequency of sync samples.
+     * <p>The default value is {@link #DEFAULT_FRAGMENT_DURATION_MS}.
      */
-  public FragmentedMp4Muxer(FileOutputStream fileOutputStream, long fragmentDurationMs) {
+    @CanIgnoreReturnValue
+    public Builder setFragmentDurationMs(long fragmentDurationMs) {
+      this.fragmentDurationMs = fragmentDurationMs;
+      return this;
+    }
+
+    /**
+     * Sets whether to enable the sample copy.
+     *
+     * <p>If the sample copy is enabled, {@link #writeSampleData(TrackToken, ByteBuffer,
+     * BufferInfo)} copies the input {@link ByteBuffer} and {@link BufferInfo} before it returns, so
+     * it is safe to reuse them immediately. Otherwise, the muxer takes ownership of the {@link
+     * ByteBuffer} and the {@link BufferInfo} and the caller must not modify them.
+     *
+     * <p>The default value is {@code true}.
+     */
+    @CanIgnoreReturnValue
+    public Builder setSampleCopyEnabled(boolean enabled) {
+      this.sampleCopyEnabled = enabled;
+      return this;
+    }
+
+    /** Builds a {@link FragmentedMp4Muxer} instance. */
+    public FragmentedMp4Muxer build() {
+      return new FragmentedMp4Muxer(fileOutputStream, fragmentDurationMs, sampleCopyEnabled);
+    }
+  }
+
+  private final FragmentedMp4Writer fragmentedMp4Writer;
+  private final MetadataCollector metadataCollector;
+
+  private FragmentedMp4Muxer(
+      FileOutputStream fileOutputStream, long fragmentDurationMs, boolean sampleCopyEnabled) {
     checkNotNull(fileOutputStream);
     metadataCollector = new MetadataCollector();
     Mp4MoovStructure moovStructure =
@@ -86,7 +129,11 @@ public final class FragmentedMp4Muxer implements Muxer {
             metadataCollector, Mp4Muxer.LAST_FRAME_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION);
     fragmentedMp4Writer =
         new FragmentedMp4Writer(
-            fileOutputStream, moovStructure, AnnexBToAvccConverter.DEFAULT, fragmentDurationMs);
+            fileOutputStream,
+            moovStructure,
+            AnnexBToAvccConverter.DEFAULT,
+            fragmentDurationMs,
+            sampleCopyEnabled);
   }
 
   @Override
@@ -97,19 +144,22 @@ public final class FragmentedMp4Muxer implements Muxer {
   /**
    * {@inheritDoc}
    *
-   * <p>The samples are cached and are written in batches so the caller must not change the {@link
+   * <p>Samples are written to the disk in batches. If {@link Builder#setSampleCopyEnabled(boolean)
-   * ByteBuffer} and the {@link MediaCodec.BufferInfo} after calling this method.
+   * sample copying} is disabled, the {@code byteBuffer} and the {@code bufferInfo} must not be
+   * modified after calling this method. Otherwise, they are copied and it is safe to modify them
+   * after this method returns.
    *
    * <p>Note: Out of order B-frames are currently not supported.
    *
    * @param trackToken The {@link TrackToken} for which this sample is being written.
-   * @param byteBuffer The encoded sample.
+   * @param byteBuffer The encoded sample. The muxer takes ownership of the buffer if {@link
-   * @param bufferInfo The {@link MediaCodec.BufferInfo} related to this sample.
+   *     Builder#setSampleCopyEnabled(boolean) sample copying} is disabled. Otherwise, the position
+   *     of the buffer is updated but the caller retains ownership.
+   * @param bufferInfo The {@link BufferInfo} related to this sample.
    * @throws IOException If there is any error while writing data to the disk.
    */
   @Override
-  public void writeSampleData(
+  public void writeSampleData(TrackToken trackToken, ByteBuffer byteBuffer, BufferInfo bufferInfo)
-      TrackToken trackToken, ByteBuffer byteBuffer, MediaCodec.BufferInfo bufferInfo)
       throws IOException {
     fragmentedMp4Writer.writeSampleData(trackToken, byteBuffer, bufferInfo);
   }

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

@@ -66,6 +66,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   private final AnnexBToAvccConverter annexBToAvccConverter;
   private final List<Track> tracks;
   private final long fragmentDurationUs;
+  private final boolean sampleCopyEnabled;
 
   private @MonotonicNonNull Track videoTrack;
   private int currentFragmentSequenceNumber;
@@ -73,15 +74,28 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   private long minInputPresentationTimeUs;
   private long maxTrackDurationUs;
 
+  /**
+   * Creates an instance.
+   *
+   * @param outputStream The {@link FileOutputStream} to write the data to.
+   * @param moovGenerator An {@link Mp4MoovStructure} instance to generate the moov box.
+   * @param annexBToAvccConverter The {@link AnnexBToAvccConverter} to be used to convert H.264 and
+   *     H.265 NAL units from the Annex-B format (using start codes to delineate NAL units) to the
+   *     AVCC format (which uses length prefixes).
+   * @param fragmentDurationMs The fragment duration (in milliseconds).
+   * @param sampleCopyEnabled Whether sample copying is enabled.
+   */
   public FragmentedMp4Writer(
       FileOutputStream outputStream,
       Mp4MoovStructure moovGenerator,
       AnnexBToAvccConverter annexBToAvccConverter,
-      long fragmentDurationMs) {
+      long fragmentDurationMs,
+      boolean sampleCopyEnabled) {
     this.outputStream = outputStream;
     this.output = outputStream.getChannel();
     this.moovGenerator = moovGenerator;
     this.annexBToAvccConverter = annexBToAvccConverter;
+    this.sampleCopyEnabled = sampleCopyEnabled;
     tracks = new ArrayList<>();
     this.fragmentDurationUs = fragmentDurationMs * 1_000;
     minInputPresentationTimeUs = Long.MAX_VALUE;
@@ -89,7 +103,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   }
 
   public TrackToken addTrack(int sortKey, Format format) {
-    Track track = new Track(format);
+    Track track = new Track(format, sampleCopyEnabled);
     tracks.add(track);
     if (MimeTypes.isVideo(format.sampleMimeType)) {
       videoTrack = track;

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

@@ -94,6 +94,7 @@ public final class Mp4Muxer implements Muxer {
 
     private @LastFrameDurationBehavior int lastFrameDurationBehavior;
     @Nullable private AnnexBToAvccConverter annexBToAvccConverter;
+    private boolean sampleCopyEnabled;
 
     /**
      * Creates a {@link Builder} instance with default values.
@@ -103,6 +104,7 @@ public final class Mp4Muxer implements Muxer {
     public Builder(FileOutputStream fileOutputStream) {
       this.fileOutputStream = checkNotNull(fileOutputStream);
       lastFrameDurationBehavior = LAST_FRAME_DURATION_BEHAVIOR_INSERT_SHORT_FRAME;
+      sampleCopyEnabled = true;
     }
 
     /**
@@ -130,14 +132,33 @@ public final class Mp4Muxer implements Muxer {
       return this;
     }
 
+    /**
+     * Sets whether to enable the sample copy.
+     *
+     * <p>If the sample copy is enabled, {@link #writeSampleData(TrackToken, ByteBuffer,
+     * BufferInfo)} copies the input {@link ByteBuffer} and {@link BufferInfo} before it returns, so
+     * it is safe to reuse them immediately. Otherwise, the muxer takes ownership of the {@link
+     * ByteBuffer} and the {@link BufferInfo} and the caller must not modify them.
+     *
+     * <p>The default value is {@code true}.
+     */
+    @CanIgnoreReturnValue
+    public Mp4Muxer.Builder setSampleCopyEnabled(boolean enabled) {
+      this.sampleCopyEnabled = enabled;
+      return this;
+    }
+
     /** Builds an {@link Mp4Muxer} instance. */
     public Mp4Muxer build() {
       MetadataCollector metadataCollector = new MetadataCollector();
       Mp4MoovStructure moovStructure =
           new Mp4MoovStructure(metadataCollector, lastFrameDurationBehavior);
-      AnnexBToAvccConverter avccConverter =
+      Mp4Writer mp4Writer =
-          annexBToAvccConverter == null ? AnnexBToAvccConverter.DEFAULT : annexBToAvccConverter;
+          new Mp4Writer(
-      Mp4Writer mp4Writer = new Mp4Writer(fileOutputStream, moovStructure, avccConverter);
+              fileOutputStream,
+              moovStructure,
+              annexBToAvccConverter == null ? AnnexBToAvccConverter.DEFAULT : annexBToAvccConverter,
+              sampleCopyEnabled);
 
       return new Mp4Muxer(mp4Writer, metadataCollector);
     }
@@ -188,13 +209,17 @@ public final class Mp4Muxer implements Muxer {
   /**
    * {@inheritDoc}
    *
-   * <p>The samples are cached and are written in batches so the caller must not change the {@link
+   * <p>Samples are written to the disk in batches. If {@link Builder#setSampleCopyEnabled(boolean)
-   * ByteBuffer} and the {@link BufferInfo} after calling this method.
+   * sample copying} is disabled, the {@code byteBuffer} and the {@code bufferInfo} must not be
+   * modified after calling this method. Otherwise, they are copied and it is safe to modify them
+   * after this method returns.
    *
    * <p>Note: Out of order B-frames are currently not supported.
    *
    * @param trackToken The {@link TrackToken} for which this sample is being written.
-   * @param byteBuffer The encoded sample.
+   * @param byteBuffer The encoded sample. The muxer takes ownership of the buffer if {@link
+   *     Builder#setSampleCopyEnabled(boolean) sample copying} is disabled. Otherwise, the position
+   *     of the buffer is updated but the caller retains ownership.
    * @param bufferInfo The {@link BufferInfo} related to this sample.
    * @throws IOException If there is any error while writing data to the disk.
    */

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

@@ -46,6 +46,7 @@ import java.util.concurrent.atomic.AtomicBoolean;
   private final AnnexBToAvccConverter annexBToAvccConverter;
   private final List<Track> tracks;
   private final AtomicBoolean hasWrittenSamples;
+  private final boolean sampleCopyEnabled;
 
   private long mdatStart;
   private long mdatEnd;
@@ -62,22 +63,25 @@ import java.util.concurrent.atomic.AtomicBoolean;
    * @param annexBToAvccConverter The {@link AnnexBToAvccConverter} to be used to convert H.264 and
    *     H.265 NAL units from the Annex-B format (using start codes to delineate NAL units) to the
    *     AVCC format (which uses length prefixes).
+   * @param sampleCopyEnabled Whether sample copying is enabled.
    */
   public Mp4Writer(
       FileOutputStream outputStream,
       Mp4MoovStructure moovGenerator,
-      AnnexBToAvccConverter annexBToAvccConverter) {
+      AnnexBToAvccConverter annexBToAvccConverter,
+      boolean sampleCopyEnabled) {
     this.outputStream = outputStream;
     this.output = outputStream.getChannel();
     this.moovGenerator = moovGenerator;
     this.annexBToAvccConverter = annexBToAvccConverter;
+    this.sampleCopyEnabled = sampleCopyEnabled;
     tracks = new ArrayList<>();
     hasWrittenSamples = new AtomicBoolean(false);
     lastMoovWritten = Range.closed(0L, 0L);
   }
 
   public TrackToken addTrack(int sortKey, Format format) {
-    Track track = new Track(format, sortKey);
+    Track track = new Track(format, sortKey, sampleCopyEnabled);
     tracks.add(track);
     Collections.sort(tracks, (a, b) -> Integer.compare(a.sortKey, b.sortKey));
     return track;

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

@@ -41,11 +41,12 @@ import java.util.List;
   public final Deque<ByteBuffer> pendingSamplesByteBuffer;
   public boolean hadKeyframe;
 
+  private final boolean sampleCopyEnabled;
   private long lastSamplePresentationTimeUs;
 
   /** Creates an instance with {@code sortKey} set to 1. */
-  public Track(Format format) {
+  public Track(Format format, boolean sampleCopyEnabled) {
-    this(format, /* sortKey= */ 1);
+    this(format, /* sortKey= */ 1, sampleCopyEnabled);
   }
 
   /**
@@ -53,10 +54,12 @@ import java.util.List;
    *
    * @param format The {@link Format} for the track.
    * @param sortKey The key used for sorting the track list.
+   * @param sampleCopyEnabled Whether sample copying is enabled.
    */
-  public Track(Format format, int sortKey) {
+  public Track(Format format, int sortKey, boolean sampleCopyEnabled) {
     this.format = format;
     this.sortKey = sortKey;
+    this.sampleCopyEnabled = sampleCopyEnabled;
     writtenSamples = new ArrayList<>();
     writtenChunkOffsets = new ArrayList<>();
     writtenChunkSampleCounts = new ArrayList<>();
@@ -84,9 +87,26 @@ import java.util.List;
       return;
     }
 
-    pendingSamplesBufferInfo.addLast(bufferInfo);
+    ByteBuffer byteBufferToAdd = byteBuffer;
-    pendingSamplesByteBuffer.addLast(byteBuffer);
+    BufferInfo bufferInfoToAdd = bufferInfo;
-    lastSamplePresentationTimeUs = bufferInfo.presentationTimeUs;
+
+    if (sampleCopyEnabled) {
+      // Copy sample data and release the original buffer.
+      byteBufferToAdd = ByteBuffer.allocateDirect(byteBuffer.remaining());
+      byteBufferToAdd.put(byteBuffer);
+      byteBufferToAdd.rewind();
+
+      bufferInfoToAdd = new BufferInfo();
+      bufferInfoToAdd.set(
+          /* newOffset= */ byteBufferToAdd.position(),
+          /* newSize= */ byteBufferToAdd.remaining(),
+          bufferInfo.presentationTimeUs,
+          bufferInfo.flags);
+    }
+
+    pendingSamplesBufferInfo.addLast(bufferInfoToAdd);
+    pendingSamplesByteBuffer.addLast(byteBufferToAdd);
+    lastSamplePresentationTimeUs = bufferInfoToAdd.presentationTimeUs;
   }
 
   @Override

libraries/transformer/src/main/java/androidx/media3/transformer/InAppMuxer.java

@@ -140,8 +140,10 @@ public final class InAppMuxer implements Muxer {
       androidx.media3.muxer.Muxer muxer =
           outputFragmentedMp4
               ? fragmentDurationMs != C.TIME_UNSET
-                  ? new FragmentedMp4Muxer(outputStream, fragmentDurationMs)
+                  ? new FragmentedMp4Muxer.Builder(outputStream)
-                  : new FragmentedMp4Muxer(outputStream)
+                      .setFragmentDurationMs(fragmentDurationMs)
+                      .build()
+                  : new FragmentedMp4Muxer.Builder(outputStream).build()
               : new Mp4Muxer.Builder(outputStream).build();
       return new InAppMuxer(muxer, metadataProvider);
     }
@@ -194,19 +196,7 @@ public final class InAppMuxer implements Muxer {
         data.position(), size, presentationTimeUs, TransformerUtil.getMediaCodecFlags(flags));
 
     try {
-      // Copy sample data and release the original buffer.
+      muxer.writeSampleData(trackTokenList.get(trackIndex), data, bufferInfo);
-      ByteBuffer byteBufferCopy = ByteBuffer.allocateDirect(data.remaining());
-      byteBufferCopy.put(data);
-      byteBufferCopy.rewind();
-
-      BufferInfo bufferInfoCopy = new BufferInfo();
-      bufferInfoCopy.set(
-          /* newOffset= */ byteBufferCopy.position(),
-          /* newSize= */ byteBufferCopy.remaining(),
-          bufferInfo.presentationTimeUs,
-          bufferInfo.flags);
-
-      muxer.writeSampleData(trackTokenList.get(trackIndex), byteBufferCopy, bufferInfoCopy);
     } catch (IOException e) {
       throw new MuxerException(
           "Failed to write sample for trackIndex="