Add different modes in MuxerWrapper

For pause and resume feature we need to use same `MuxerWrapper`
for `remuxing processed video` and then to `process remaining video`.
In order to use same `MuxerWrapper` across `different Exports`
we need to preserve its state.

PiperOrigin-RevId: 564728396

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

@@ -21,9 +21,11 @@ import static androidx.media3.common.util.Assertions.checkNotNull;
 import static androidx.media3.common.util.Assertions.checkState;
 import static androidx.media3.common.util.Util.contains;
 import static java.lang.Math.max;
+import static java.lang.annotation.ElementType.TYPE_USE;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
 import android.util.SparseArray;
+import androidx.annotation.IntDef;
 import androidx.annotation.IntRange;
 import androidx.annotation.Nullable;
 import androidx.media3.common.C;
@@ -34,6 +36,10 @@ import androidx.media3.common.util.Util;
 import androidx.media3.effect.DebugTraceUtil;
 import com.google.common.collect.ImmutableList;
 import java.io.File;
+import java.lang.annotation.Documented;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 import java.nio.ByteBuffer;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
@@ -47,6 +53,26 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
  * <p>This wrapper can contain at most one video track and one audio track.
  */
 /* package */ final class MuxerWrapper {
+  /** Different modes for muxing. */
+  @Documented
+  @Retention(RetentionPolicy.SOURCE)
+  @Target(TYPE_USE)
+  @IntDef({MUXER_MODE_DEFAULT, MUXER_MODE_MUX_PARTIAL_VIDEO, MUXER_MODE_APPEND_VIDEO})
+  public @interface MuxerMode {}
+
+  /** The default muxer mode. */
+  public static final int MUXER_MODE_DEFAULT = 0;
+
+  /**
+   * Used for muxing a partial video. The video {@link TrackInfo} is kept the same when {@linkplain
+   * #changeToAppendVideoMode() transitioning} to {@link #MUXER_MODE_APPEND_VIDEO} after finishing
+   * muxing partial video. Only one video track can be {@link #addTrackFormat(Format) added} in this
+   * mode.
+   */
+  public static final int MUXER_MODE_MUX_PARTIAL_VIDEO = 1;
+
+  /** Used for appending the remaining video samples with the previously muxed partial video. */
+  public static final int MUXER_MODE_APPEND_VIDEO = 2;
 
   private static final String TIMER_THREAD_NAME = "Muxer:Timer";
   private static final String MUXER_TIMEOUT_ERROR_FORMAT_STRING =
@@ -83,21 +109,45 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
   private @MonotonicNonNull ScheduledFuture<?> abortScheduledFuture;
   private boolean isAborted;
   private @MonotonicNonNull Muxer muxer;
+  private @MuxerMode int muxerMode;
 
+  private volatile boolean muxedPartialVideo;
   private volatile int additionalRotationDegrees;
   private volatile int trackCount;
 
-  /** Creates an instance. */
+  /**
-  public MuxerWrapper(String outputPath, Muxer.Factory muxerFactory, Listener listener) {
+   * Creates an instance.
+   *
+   * @param outputPath The output file path to write the media data to.
+   * @param muxerFactory A {@link Muxer.Factory} to create a {@link Muxer}.
+   * @param listener A {@link MuxerWrapper.Listener}.
+   * @param muxerMode The {@link MuxerMode}. The initial mode must be {@link #MUXER_MODE_DEFAULT} or
+   *     {@link #MUXER_MODE_MUX_PARTIAL_VIDEO}.
+   */
+  public MuxerWrapper(
+      String outputPath, Muxer.Factory muxerFactory, Listener listener, @MuxerMode int muxerMode) {
     this.outputPath = outputPath;
     this.muxerFactory = muxerFactory;
     this.listener = listener;
-
+    checkArgument(muxerMode == MUXER_MODE_DEFAULT || muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO);
+    this.muxerMode = muxerMode;
     trackTypeToInfo = new SparseArray<>();
     previousTrackType = C.TRACK_TYPE_NONE;
     abortScheduledExecutorService = Util.newSingleThreadScheduledExecutor(TIMER_THREAD_NAME);
   }
 
+  /**
+   * Changes {@link MuxerMode} to {@link #MUXER_MODE_APPEND_VIDEO}.
+   *
+   * <p>This method must be called only after partial video is muxed using {@link
+   * #MUXER_MODE_MUX_PARTIAL_VIDEO}.
+   */
+  public void changeToAppendVideoMode() {
+    checkState(muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO);
+
+    muxerMode = MUXER_MODE_APPEND_VIDEO;
+  }
+
   /**
    * Sets the clockwise rotation to add to the {@linkplain #addTrackFormat(Format) video track's}
    * rotation, in degrees.
@@ -123,12 +173,23 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
    * <p>The track count must be set before any track format is {@linkplain #addTrackFormat(Format)
    * added}.
    *
+   * <p>When using muxer mode other than {@link #MUXER_MODE_DEFAULT}, the track count must be 1.
+   *
    * <p>Can be called from any thread.
    *
    * @throws IllegalStateException If a track format was {@linkplain #addTrackFormat(Format) added}
    *     before calling this method.
    */
   public void setTrackCount(@IntRange(from = 1) int trackCount) {
+    if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO || muxerMode == MUXER_MODE_APPEND_VIDEO) {
+      checkArgument(
+          trackCount == 1,
+          "Only one video track can be added in MUXER_MODE_MUX_PARTIAL_VIDEO and"
+              + " MUXER_MODE_APPEND_VIDEO");
+      if (muxerMode == MUXER_MODE_APPEND_VIDEO) {
+        return;
+      }
+    }
     checkState(
         trackTypeToInfo.size() == 0,
         "The track count cannot be changed after adding track formats.");
@@ -158,8 +219,11 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
    *
    * <p>{@link Muxer#addMetadata(Metadata)} is called if the {@link Format#metadata} is present.
    *
-   * @param format The {@link Format} to be added.
+   * @param format The {@link Format} to be added. In {@link #MUXER_MODE_APPEND_VIDEO} mode, the
-   * @throws IllegalArgumentException If the format is unsupported.
+   *     added {@link Format} must match the existing {@link Format} set when the muxer was in
+   *     {@link #MUXER_MODE_MUX_PARTIAL_VIDEO} mode.
+   * @throws IllegalArgumentException If the format is unsupported or if it does not match the
+   *     existing format in {@link #MUXER_MODE_APPEND_VIDEO} mode.
    * @throws IllegalStateException If the number of formats added exceeds the {@linkplain
    *     #setTrackCount track count}, if {@link #setTrackCount(int)} has not been called or if there
    *     is already a track of that {@link C.TrackType}.
@@ -167,6 +231,27 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
    *     the track.
    */
   public void addTrackFormat(Format format) throws Muxer.MuxerException {
+    if (muxerMode == MUXER_MODE_APPEND_VIDEO) {
+      checkState(contains(trackTypeToInfo, C.TRACK_TYPE_VIDEO));
+      TrackInfo videoTrackInfo = trackTypeToInfo.get(C.TRACK_TYPE_VIDEO);
+
+      // Ensure that video formats are the same. Some fields like codecs, averageBitrate, framerate,
+      // etc, don't match exactly in the Extractor output format and the Encoder output
+      // format but these fields can be ignored.
+      Format existingFormat = videoTrackInfo.format;
+      checkArgument(Util.areEqual(existingFormat.sampleMimeType, format.sampleMimeType));
+      checkArgument(existingFormat.width == format.width);
+      checkArgument(existingFormat.height == format.height);
+      checkArgument(existingFormat.initializationDataEquals(format));
+      checkArgument(Util.areEqual(existingFormat.colorInfo, format.colorInfo));
+
+      checkNotNull(muxer);
+      resetAbortTimer();
+      return;
+    } else if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO) {
+      checkArgument(MimeTypes.isVideo(format.sampleMimeType));
+    }
+
     int trackCount = this.trackCount;
     checkState(trackCount > 0, "The track count should be set before the formats are added.");
     checkState(trackTypeToInfo.size() < trackCount, "All track formats have already been added.");
@@ -275,25 +360,44 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
       DebugTraceUtil.logEvent(DebugTraceUtil.EVENT_MUXER_TRACK_ENDED_AUDIO, trackInfo.timeUs);
     }
 
-    trackTypeToInfo.delete(trackType);
+    if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO) {
-    if (trackTypeToInfo.size() == 0) {
+      muxedPartialVideo = true;
-      abortScheduledExecutorService.shutdownNow();
+    } else {
-      if (!isEnded) {
+      trackTypeToInfo.delete(trackType);
+      if (trackTypeToInfo.size() == 0) {
         isEnded = true;
-        listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes());
       }
     }
+
+    if (muxedPartialVideo) {
+      listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes());
+      if (abortScheduledFuture != null) {
+        abortScheduledFuture.cancel(/* mayInterruptIfRunning= */ false);
+      }
+      return;
+    }
+
+    if (isEnded) {
+      listener.onEnded(Util.usToMs(maxEndedTrackTimeUs), getCurrentOutputSizeBytes());
+      abortScheduledExecutorService.shutdownNow();
+    }
   }
 
-  /** Returns whether all the tracks are {@linkplain #endTrack(int) ended}. */
+  /**
+   * Returns whether all the tracks are {@linkplain #endTrack(int) ended} or a partial video is
+   * completely muxed using {@link #MUXER_MODE_MUX_PARTIAL_VIDEO}.
+   */
   public boolean isEnded() {
-    return isEnded;
+    return isEnded || (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO && muxedPartialVideo);
   }
 
   /**
    * Finishes writing the output and releases any resources associated with muxing.
    *
-   * <p>The muxer cannot be used anymore once this method has been called.
+   * <p>When this method is called in {@link #MUXER_MODE_MUX_PARTIAL_VIDEO} mode, the resources are
+   * not released and the {@link MuxerWrapper} can be reused after {@link #changeToAppendVideoMode()
+   * changing mode} to {@link #MUXER_MODE_APPEND_VIDEO} mode. In all other modes the {@link
+   * MuxerWrapper} cannot be used anymore once this method has been called.
    *
    * @param forCancellation Whether the reason for releasing the resources is the transformation
    *     cancellation.
@@ -301,6 +405,9 @@ import org.checkerframework.checker.nullness.qual.RequiresNonNull;
    *     and {@code forCancellation} is false.
    */
   public void release(boolean forCancellation) throws Muxer.MuxerException {
+    if (muxerMode == MUXER_MODE_MUX_PARTIAL_VIDEO && !forCancellation) {
+      return;
+    }
     isReady = false;
     abortScheduledExecutorService.shutdownNow();
     if (muxer != null) {

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

@@ -832,7 +832,9 @@ public final class Transformer {
   public void start(Composition composition, String path) {
     ComponentListener componentListener = new ComponentListener(composition);
     startInternal(
-        composition, new MuxerWrapper(path, muxerFactory, componentListener), componentListener);
+        composition,
+        new MuxerWrapper(path, muxerFactory, componentListener, MuxerWrapper.MUXER_MODE_DEFAULT),
+        componentListener);
   }
 
   /**

libraries/transformer/src/test/java/androidx/media3/transformer/MuxerWrapperTest.java

@@ -0,0 +1,210 @@
+/*
+ * Copyright 2023 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
+ *
+ *      http://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.transformer;
+
+import static androidx.media3.common.MimeTypes.VIDEO_H264;
+import static androidx.media3.transformer.MuxerWrapper.MUXER_MODE_MUX_PARTIAL_VIDEO;
+import static com.google.common.truth.Truth.assertThat;
+import static org.junit.Assert.assertThrows;
+
+import androidx.media3.common.C;
+import androidx.media3.common.ColorInfo;
+import androidx.media3.common.Format;
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import com.google.common.collect.ImmutableList;
+import java.io.IOException;
+import java.nio.ByteBuffer;
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+
+/** Unit tests for {@link MuxerWrapper}. */
+@RunWith(AndroidJUnit4.class)
+public class MuxerWrapperTest {
+  @Rule public final TemporaryFolder temporaryFolder = new TemporaryFolder();
+
+  private static final Format FAKE_VIDEO_TRACK_FORMAT =
+      new Format.Builder()
+          .setSampleMimeType(VIDEO_H264)
+          .setWidth(1080)
+          .setHeight(720)
+          .setInitializationData(ImmutableList.of(new byte[] {1, 2, 3, 4}))
+          .setColorInfo(ColorInfo.SDR_BT709_LIMITED)
+          .build();
+  private static final Format FAKE_AUDIO_TRACK_FORMAT =
+      new Format.Builder()
+          .setSampleMimeType("audio/mp4a-latm")
+          .setSampleRate(40000)
+          .setChannelCount(2)
+          .build();
+
+  private static final ByteBuffer FAKE_SAMPLE = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});
+
+  private String outputFilePath;
+
+  @Before
+  public void setUp() throws IOException {
+    outputFilePath = temporaryFolder.newFile("output.mp4").getPath();
+  }
+
+  @Test
+  public void changeToAppendVideoMode_afterDefaultMode_throws() {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MuxerWrapper.MUXER_MODE_DEFAULT);
+
+    assertThrows(IllegalStateException.class, muxerWrapper::changeToAppendVideoMode);
+  }
+
+  @Test
+  public void setTrackCount_toTwoInMuxPartialVideoMode_throws() {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+
+    assertThrows(IllegalArgumentException.class, () -> muxerWrapper.setTrackCount(2));
+  }
+
+  @Test
+  public void setTrackCount_toTwoInAppendVideoMode_throws() throws Exception {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+    muxerWrapper.writeSample(
+        C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0);
+    muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO);
+    muxerWrapper.changeToAppendVideoMode();
+
+    assertThrows(IllegalArgumentException.class, () -> muxerWrapper.setTrackCount(2));
+  }
+
+  @Test
+  public void addTrackFormat_withAudioFormatInMuxPartialVideoMode_throws() {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+
+    assertThrows(
+        IllegalArgumentException.class, () -> muxerWrapper.addTrackFormat(FAKE_AUDIO_TRACK_FORMAT));
+  }
+
+  @Test
+  public void addTrackFormat_withSameVideoFormatInAppendVideoMode_doesNotThrow() throws Exception {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+    muxerWrapper.writeSample(
+        C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0);
+    muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO);
+    muxerWrapper.changeToAppendVideoMode();
+    muxerWrapper.setTrackCount(1);
+
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+  }
+
+  @Test
+  public void addTrackFormat_withDifferentVideoFormatInAppendVideoMode_throws() throws Exception {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+    muxerWrapper.writeSample(
+        C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0);
+    muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO);
+    muxerWrapper.changeToAppendVideoMode();
+    muxerWrapper.setTrackCount(1);
+    Format differentVideoFormat = FAKE_VIDEO_TRACK_FORMAT.buildUpon().setHeight(5000).build();
+
+    assertThrows(
+        IllegalArgumentException.class, () -> muxerWrapper.addTrackFormat(differentVideoFormat));
+  }
+
+  @Test
+  public void isEnded_afterPartialVideoMuxed_returnsTrue() throws Exception {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+    muxerWrapper.writeSample(
+        C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0);
+    muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO);
+
+    assertThat(muxerWrapper.isEnded()).isTrue();
+  }
+
+  @Test
+  public void isEnded_afterStartingAppendVideo_returnsFalse() throws Exception {
+    MuxerWrapper muxerWrapper =
+        new MuxerWrapper(
+            outputFilePath,
+            new DefaultMuxer.Factory(),
+            new NoOpMuxerListenerImpl(),
+            MUXER_MODE_MUX_PARTIAL_VIDEO);
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+    muxerWrapper.writeSample(
+        C.TRACK_TYPE_VIDEO, FAKE_SAMPLE, /* isKeyFrame= */ true, /* presentationTimeUs= */ 0);
+    muxerWrapper.endTrack(C.TRACK_TYPE_VIDEO);
+    muxerWrapper.changeToAppendVideoMode();
+    muxerWrapper.setTrackCount(1);
+    muxerWrapper.addTrackFormat(FAKE_VIDEO_TRACK_FORMAT);
+
+    assertThat(muxerWrapper.isEnded()).isFalse();
+  }
+
+  private static final class NoOpMuxerListenerImpl implements MuxerWrapper.Listener {
+
+    @Override
+    public void onTrackEnded(
+        @C.TrackType int trackType, Format format, int averageBitrate, int sampleCount) {}
+
+    @Override
+    public void onEnded(long durationMs, long fileSizeBytes) {}
+
+    @Override
+    public void onError(ExportException exportException) {}
+  }
+}