From 081baa03b9c59cce0b32e36cbbd37cca1f5a3cb0 Mon Sep 17 00:00:00 2001
From: sheenachhabra <sheenachhabra@google.com>
Date: Tue, 23 Jan 2024 09:18:25 -0800
Subject: [PATCH] Process all tracks before writing fragment in fragmented MP4

Earlier implementation processed each track (pending sample's buffer info)
individually when writing their corresponding "traf" box in a fragment.
The change involves processing all tracks before start writing "traf" boxes.

#minor-release

PiperOrigin-RevId: 600811093
(cherry picked from commit 4c1581a17542f56db9fd85e0c7a6894aef57be52)
---
 .../media3/muxer/FragmentedMp4Writer.java     | 53 +++++++++++++------
 1 file changed, 36 insertions(+), 17 deletions(-)

diff --git a/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java b/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
index f67b471aeb..0602728743 100644
--- a/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
+++ b/libraries/muxer/src/main/java/androidx/media3/muxer/FragmentedMp4Writer.java
@@ -27,6 +27,7 @@ import androidx.media3.common.Format;
 import androidx.media3.common.MimeTypes;
 import androidx.media3.common.util.Util;
 import androidx.media3.muxer.Mp4Muxer.TrackToken;
+import com.google.common.collect.ImmutableList;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -115,6 +116,16 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     }
   }
 
+  private static ImmutableList<ByteBuffer> createTrafBoxes(List<ProcessedTrackInfo> trackInfos) {
+    ImmutableList.Builder<ByteBuffer> trafBoxes = new ImmutableList.Builder<>();
+    for (int i = 0; i < trackInfos.size(); i++) {
+      ProcessedTrackInfo currentTrackInfo = trackInfos.get(i);
+      ByteBuffer trun = Boxes.trun(currentTrackInfo.pendingSamplesMetadata);
+      trafBoxes.add(Boxes.traf(Boxes.tfhd(currentTrackInfo.trackId), trun));
+    }
+    return trafBoxes.build();
+  }
+
   private void createHeader() throws IOException {
     output.position(0L);
     output.write(Boxes.ftyp());
@@ -147,8 +158,9 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
   }
 
   private void createFragment() throws IOException {
+    ImmutableList<ProcessedTrackInfo> trackInfos = processAllTracks();
     // Write moof box.
-    List<ByteBuffer> trafBoxes = createTrafBoxes();
+    ImmutableList<ByteBuffer> trafBoxes = createTrafBoxes(trackInfos);
     if (trafBoxes.isEmpty()) {
       return;
     }
@@ -159,19 +171,6 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     currentFragmentSequenceNumber++;
   }
 
-  private List<ByteBuffer> createTrafBoxes() {
-    List<ByteBuffer> trafBoxes = new ArrayList<>();
-    for (int i = 0; i < tracks.size(); i++) {
-      Track currentTrack = tracks.get(i);
-      if (!currentTrack.pendingSamplesBufferInfo.isEmpty()) {
-        List<SampleMetadata> samplesMetadata = processPendingSamplesBufferInfo(currentTrack);
-        ByteBuffer trun = Boxes.trun(samplesMetadata);
-        trafBoxes.add(Boxes.traf(Boxes.tfhd(/* trackId= */ i + 1), trun));
-      }
-    }
-    return trafBoxes;
-  }
-
   private void writeMdatBox() throws IOException {
     long mdatStartPosition = output.position();
     int mdatHeaderSize = 8; // 4 bytes (box size) + 4 bytes (box name)
@@ -210,7 +209,17 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
     output.position(currentPosition);
   }
 
-  private List<SampleMetadata> processPendingSamplesBufferInfo(Track track) {
+  private ImmutableList<ProcessedTrackInfo> processAllTracks() {
+    ImmutableList.Builder<ProcessedTrackInfo> trackInfos = new ImmutableList.Builder<>();
+    for (int i = 0; i < tracks.size(); i++) {
+      if (!tracks.get(i).pendingSamplesBufferInfo.isEmpty()) {
+        trackInfos.add(processTrack(/* trackId= */ i + 1, tracks.get(i)));
+      }
+    }
+    return trackInfos.build();
+  }
+
+  private ProcessedTrackInfo processTrack(int trackId, Track track) {
     List<BufferInfo> sampleBufferInfos = new ArrayList<>(track.pendingSamplesBufferInfo);
 
     List<Long> sampleDurations =
@@ -222,7 +231,7 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
             track.videoUnitTimebase(),
             Mp4Muxer.LAST_FRAME_DURATION_BEHAVIOR_DUPLICATE_PREV_DURATION);
 
-    List<SampleMetadata> pendingSamplesMetadata = new ArrayList<>(sampleBufferInfos.size());
+    ImmutableList.Builder<SampleMetadata> pendingSamplesMetadata = new ImmutableList.Builder<>();
     for (int i = 0; i < sampleBufferInfos.size(); i++) {
       pendingSamplesMetadata.add(
           new SampleMetadata(
@@ -233,6 +242,16 @@ import org.checkerframework.checker.nullness.qual.MonotonicNonNull;
 
     // Clear the queue.
     track.pendingSamplesBufferInfo.clear();
-    return pendingSamplesMetadata;
+    return new ProcessedTrackInfo(trackId, pendingSamplesMetadata.build());
+  }
+
+  private static class ProcessedTrackInfo {
+    public final int trackId;
+    public final ImmutableList<SampleMetadata> pendingSamplesMetadata;
+
+    public ProcessedTrackInfo(int trackId, ImmutableList<SampleMetadata> pendingSamplesMetadata) {
+      this.trackId = trackId;
+      this.pendingSamplesMetadata = pendingSamplesMetadata;
+    }
   }
 }