Add support for setting last sample duration in Mp4Muxer

PiperOrigin-RevId: 669340763

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

@@ -142,7 +142,8 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
               track.writtenSamples,
               minInputPtsUs,
               track.videoUnitTimebase(),
-              lastSampleDurationBehavior);
+              lastSampleDurationBehavior,
+              track.endOfStreamTimestampUs);
 
       long trackDurationInTrackUnitsVu = 0;
       for (int j = 0; j < sampleDurationsVu.size(); j++) {
@@ -772,14 +773,15 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
    *     {@code samplesInfo} list.
    * @param videoUnitTimescale The timescale of the track.
    * @param lastSampleDurationBehavior The behaviour for the last sample duration.
+   * @param endOfStreamTimestampUs The timestamp (in microseconds) of the end of stream sample.
    * @return A list of all the sample durations.
    */
-  // TODO: b/280084657 - Add support for setting last sample duration.
   public static List<Integer> convertPresentationTimestampsToDurationsVu(
       List<BufferInfo> samplesInfo,
       long firstSamplePresentationTimeUs,
       int videoUnitTimescale,
-      @Mp4Muxer.LastSampleDurationBehavior int lastSampleDurationBehavior) {
+      @Mp4Muxer.LastSampleDurationBehavior int lastSampleDurationBehavior,
+      long endOfStreamTimestampUs) {
     List<Long> presentationTimestampsUs = new ArrayList<>(samplesInfo.size());
     List<Integer> durationsVu = new ArrayList<>(samplesInfo.size());
 
@@ -816,7 +818,19 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
       currentSampleTimeUs = nextSampleTimeUs;
     }
 
-    durationsVu.add(getLastSampleDurationVu(durationsVu, lastSampleDurationBehavior));
+    long lastSampleDurationVuFromEndOfStream = 0;
+    if (endOfStreamTimestampUs != C.TIME_UNSET) {
+      lastSampleDurationVuFromEndOfStream =
+          vuFromUs(endOfStreamTimestampUs, videoUnitTimescale)
+              - vuFromUs(currentSampleTimeUs, videoUnitTimescale);
+      checkState(
+          lastSampleDurationVuFromEndOfStream <= Integer.MAX_VALUE,
+          "Only 32-bit sample duration is allowed");
+    }
+
+    durationsVu.add(
+        getLastSampleDurationVu(
+            durationsVu, lastSampleDurationBehavior, (int) lastSampleDurationVuFromEndOfStream));
 
     return durationsVu;
   }
@@ -1221,13 +1235,11 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
     return timestampVu * 1_000_000L / videoUnitTimebase;
   }
 
-  /**
+  /** Returns the duration of the last sample (in video units). */
-   * Returns the duration of the last sample (in video units) based on previous sample durations and
-   * the {@code lastSampleDurationBehavior}.
-   */
   private static int getLastSampleDurationVu(
       List<Integer> sampleDurationsExceptLast,
-      @Mp4Muxer.LastSampleDurationBehavior int lastSampleDurationBehavior) {
+      @Mp4Muxer.LastSampleDurationBehavior int lastSampleDurationBehavior,
+      int lastSampleDurationVuFromEndOfStream) {
     switch (lastSampleDurationBehavior) {
       case Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION:
         // For a track having less than 3 samples, duplicating the last frame duration will
@@ -1238,6 +1250,8 @@ import org.checkerframework.checker.nullness.qual.PolyNull;
       case Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE:
         // Keep the last sample duration as short as possible.
         return 0;
+      case Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG:
+        return lastSampleDurationVuFromEndOfStream;
       default:
         throw new IllegalArgumentException(
             "Unexpected value for the last frame duration behavior " + lastSampleDurationBehavior);

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

@@ -333,7 +333,8 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
                 ? minInputPresentationTimeUs
                 : pendingSamplesBufferInfo.get(0).presentationTimeUs,
             track.videoUnitTimebase(),
-            Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION);
+            Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION,
+            track.endOfStreamTimestampUs);
 
     List<Integer> sampleCompositionTimeOffsets =
         Boxes.calculateSampleCompositionTimeOffsets(

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

@@ -27,6 +27,7 @@ import static androidx.media3.muxer.MuxerUtil.isMetadataSupported;
 import static androidx.media3.muxer.MuxerUtil.populateEditableVideoTracksMetadata;
 import static java.lang.annotation.ElementType.TYPE_USE;
 
+import android.media.MediaCodec;
 import android.media.MediaCodec.BufferInfo;
 import androidx.annotation.IntDef;
 import androidx.annotation.Nullable;
@@ -140,17 +141,18 @@ public final class Mp4Muxer implements Muxer {
     }
   }
 
-  /** Behavior for the last sample duration. */
+  /** Behavior for the duration of the last sample. */
   @Documented
   @Retention(RetentionPolicy.SOURCE)
   @Target(TYPE_USE)
   @IntDef({
     LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
     LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION,
+    LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG
   })
   public @interface LastSampleDurationBehavior {}
 
-  /** Insert a zero-length last sample. */
+  /** The duration of the last sample is set to 0. */
   public static final int LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE = 0;
 
   /**
@@ -159,6 +161,23 @@ public final class Mp4Muxer implements Muxer {
    */
   public static final int LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION = 1;
 
+  /**
+   * Use the {@link MediaCodec#BUFFER_FLAG_END_OF_STREAM end of stream sample} to set the duration
+   * of the last sample.
+   *
+   * <p>After {@linkplain #writeSampleData writing} all the samples for a track, the app must
+   * {@linkplain #writeSampleData write} an empty sample with flag {@link
+   * MediaCodec#BUFFER_FLAG_END_OF_STREAM}. The timestamp of this sample should be equal to the
+   * desired track duration.
+   *
+   * <p>Once a sample with flag {@link MediaCodec#BUFFER_FLAG_END_OF_STREAM} is {@linkplain
+   * #writeSampleData written}, no more samples can be written for that track.
+   *
+   * <p>If no explicit {@link MediaCodec#BUFFER_FLAG_END_OF_STREAM} sample is passed, then the
+   * duration of the last sample will be set to 0.
+   */
+  public static final int LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG = 2;
+
   /** The specific MP4 file format. */
   @Documented
   @Retention(RetentionPolicy.SOURCE)

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

@@ -15,8 +15,11 @@
  */
 package androidx.media3.muxer;
 
+import static androidx.media3.common.util.Assertions.checkArgument;
+
 import android.media.MediaCodec;
 import android.media.MediaCodec.BufferInfo;
+import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.MimeTypes;
 import androidx.media3.muxer.Muxer.TrackToken;
@@ -36,6 +39,7 @@ import java.util.List;
   public final Deque<BufferInfo> pendingSamplesBufferInfo;
   public final Deque<ByteBuffer> pendingSamplesByteBuffer;
   public boolean hadKeyframe;
+  public long endOfStreamTimestampUs;
 
   private final boolean sampleCopyEnabled;
 
@@ -60,11 +64,19 @@ import java.util.List;
     writtenChunkSampleCounts = new ArrayList<>();
     pendingSamplesBufferInfo = new ArrayDeque<>();
     pendingSamplesByteBuffer = new ArrayDeque<>();
+    endOfStreamTimestampUs = C.TIME_UNSET;
   }
 
   public void writeSampleData(ByteBuffer byteBuffer, BufferInfo bufferInfo) {
+    checkArgument(
+        endOfStreamTimestampUs == C.TIME_UNSET,
+        "Samples can not be written after writing a sample with"
+            + " MediaCodec.BUFFER_FLAG_END_OF_STREAM flag");
     //  Skip empty samples.
     if (bufferInfo.size == 0 || byteBuffer.remaining() == 0) {
+      if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) {
+        endOfStreamTimestampUs = bufferInfo.presentationTimeUs;
+      }
       return;
     }
 

libraries/muxer/src/test/java/androidx/media3/muxer/BoxesTest.java

@@ -477,7 +477,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     assertThat(durationsVu).containsExactly(0);
   }
@@ -493,7 +494,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     assertThat(durationsVu).containsExactly(0);
   }
@@ -509,7 +511,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     assertThat(durationsVu).containsExactly(3_000, 5_000, 0);
   }
@@ -525,7 +528,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION);
+            LAST_SAMPLE_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION,
+            C.TIME_UNSET);
 
     assertThat(durationsVu).containsExactly(3_000, 5_000, 5_000);
   }
@@ -541,11 +545,29 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     assertThat(durationsVu).containsExactly(100, 100, 800, 100, 0);
   }
 
+  @Test
+  public void
+      convertPresentationTimestampsToDurationsVu_withLastSampleDurationBehaviorUsingEndOfStreamFlag_returnsExpectedDurations() {
+    List<MediaCodec.BufferInfo> sampleBufferInfos =
+        createBufferInfoListWithSamplePresentationTimestamps(0L, 1_000L, 2_000L, 3_000L, 4_000L);
+
+    List<Integer> durationsVu =
+        Boxes.convertPresentationTimestampsToDurationsVu(
+            sampleBufferInfos,
+            /* firstSamplePresentationTimeUs= */ 0L,
+            VU_TIMEBASE,
+            Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG,
+            /* endOfStreamTimestampUs= */ 10_000);
+
+    assertThat(durationsVu).containsExactly(100, 100, 100, 100, 600);
+  }
+
   @Test
   public void createSttsBox_withSingleSampleDuration_matchesExpected() throws IOException {
     ImmutableList<Integer> sampleDurations = ImmutableList.of(500);
@@ -595,7 +617,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     ByteBuffer cttsBox = Boxes.ctts(sampleBufferInfos, durationsVu, VU_TIMEBASE);
 
@@ -612,7 +635,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     ByteBuffer cttsBox = Boxes.ctts(sampleBufferInfos, durationsVu, VU_TIMEBASE);
 
@@ -631,7 +655,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 0L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     ByteBuffer cttsBox = Boxes.ctts(sampleBufferInfos, durationsVu, VU_TIMEBASE);
 
@@ -651,7 +676,8 @@ public class BoxesTest {
             sampleBufferInfos,
             /* firstSamplePresentationTimeUs= */ 23698215060L,
             VU_TIMEBASE,
-            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE);
+            LAST_SAMPLE_DURATION_BEHAVIOR_INSERT_SHORT_SAMPLE,
+            C.TIME_UNSET);
 
     ByteBuffer cttsBox = Boxes.ctts(sampleBufferInfos, durationsVu, VU_TIMEBASE);
 

libraries/muxer/src/test/java/androidx/media3/muxer/Mp4MuxerEndToEndTest.java

@@ -22,6 +22,7 @@ import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
 
 import android.content.Context;
+import android.media.MediaCodec;
 import android.media.MediaCodec.BufferInfo;
 import android.util.Pair;
 import androidx.media3.common.C;
@@ -673,6 +674,93 @@ public class Mp4MuxerEndToEndTest {
             "mp4_with_editable_video_tracks_when_editable_track_samples_interleaved.mp4"));
   }
 
+  @Test
+  public void
+      createMp4File_withLastSampleDurationBehaviorUsingEndOfStreamFlag_writesSamplesWithCorrectDurations()
+          throws Exception {
+    String outputFilePath = temporaryFolder.newFile().getPath();
+    Mp4Muxer mp4Muxer =
+        new Mp4Muxer.Builder(new FileOutputStream(outputFilePath))
+            .setLastSampleDurationBehavior(
+                Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG)
+            .build();
+    mp4Muxer.addMetadataEntry(
+        new Mp4TimestampData(
+            /* creationTimestampSeconds= */ 100_000_000L,
+            /* modificationTimestampSeconds= */ 500_000_000L));
+    Pair<ByteBuffer, BufferInfo> sample1 = getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 0L);
+    Pair<ByteBuffer, BufferInfo> sample2 =
+        getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 100L);
+    Pair<ByteBuffer, BufferInfo> sample3 =
+        getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 200L);
+    Pair<ByteBuffer, BufferInfo> sample4 =
+        getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 300L);
+
+    long expectedDurationUs = 1_000L;
+    try {
+      TrackToken track = mp4Muxer.addTrack(FAKE_VIDEO_FORMAT);
+      mp4Muxer.writeSampleData(track, sample1.first, sample1.second);
+      mp4Muxer.writeSampleData(track, sample2.first, sample2.second);
+      mp4Muxer.writeSampleData(track, sample3.first, sample3.second);
+      mp4Muxer.writeSampleData(track, sample4.first, sample4.second);
+      // Write end of stream sample.
+      BufferInfo endOfStreamBufferInfo = new BufferInfo();
+      endOfStreamBufferInfo.set(
+          /* newOffset= */ 0,
+          /* newSize= */ 0,
+          /* newTimeUs= */ expectedDurationUs,
+          /* newFlags= */ MediaCodec.BUFFER_FLAG_END_OF_STREAM);
+      mp4Muxer.writeSampleData(track, ByteBuffer.allocate(0), endOfStreamBufferInfo);
+    } finally {
+      mp4Muxer.close();
+    }
+
+    FakeExtractorOutput fakeExtractorOutput =
+        TestUtil.extractAllSamplesFromFilePath(
+            new Mp4Extractor(new DefaultSubtitleParserFactory()), outputFilePath);
+    fakeExtractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO).assertSampleCount(4);
+    assertThat(fakeExtractorOutput.seekMap.getDurationUs()).isEqualTo(expectedDurationUs);
+  }
+
+  @Test
+  public void
+      createMp4File_withLastSampleDurationBehaviorUsingEndOfStreamFlagButNoEndOfStreamSample_outputsDurationEqualsToLastSampleTimestamp()
+          throws Exception {
+    String outputFilePath = temporaryFolder.newFile().getPath();
+    Mp4Muxer mp4Muxer =
+        new Mp4Muxer.Builder(new FileOutputStream(outputFilePath))
+            .setLastSampleDurationBehavior(
+                Mp4Muxer.LAST_SAMPLE_DURATION_BEHAVIOR_USING_END_OF_STREAM_FLAG)
+            .build();
+    mp4Muxer.addMetadataEntry(
+        new Mp4TimestampData(
+            /* creationTimestampSeconds= */ 100_000_000L,
+            /* modificationTimestampSeconds= */ 500_000_000L));
+    Pair<ByteBuffer, BufferInfo> sample1 = getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 0L);
+    Pair<ByteBuffer, BufferInfo> sample2 =
+        getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 100L);
+    Pair<ByteBuffer, BufferInfo> sample3 =
+        getFakeSampleAndSampleInfo(/* presentationTimeUs= */ 200L);
+    long lastSampleTimestampUs = 300L;
+    Pair<ByteBuffer, BufferInfo> sample4 = getFakeSampleAndSampleInfo(lastSampleTimestampUs);
+
+    try {
+      TrackToken track = mp4Muxer.addTrack(FAKE_VIDEO_FORMAT);
+      mp4Muxer.writeSampleData(track, sample1.first, sample1.second);
+      mp4Muxer.writeSampleData(track, sample2.first, sample2.second);
+      mp4Muxer.writeSampleData(track, sample3.first, sample3.second);
+      mp4Muxer.writeSampleData(track, sample4.first, sample4.second);
+    } finally {
+      mp4Muxer.close();
+    }
+
+    FakeExtractorOutput fakeExtractorOutput =
+        TestUtil.extractAllSamplesFromFilePath(
+            new Mp4Extractor(new DefaultSubtitleParserFactory()), outputFilePath);
+    fakeExtractorOutput.track(/* id= */ 0, C.TRACK_TYPE_VIDEO).assertSampleCount(4);
+    assertThat(fakeExtractorOutput.seekMap.getDurationUs()).isEqualTo(lastSampleTimestampUs);
+  }
+
   private static void writeFakeSamples(Mp4Muxer muxer, TrackToken trackToken, int sampleCount)
       throws Muxer.MuxerException {
     for (int i = 0; i < sampleCount; i++) {