Add methods required for seek to AudioGraph

Seeking will consist of the following steps:
- Block the AudioGraph input data.
- Flush the AudioGraph.
- Seek the ExoPlayers.
- Unblock the AudioGraph input data.

PiperOrigin-RevId: 617868124

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

@@ -42,6 +42,8 @@ import java.util.Objects;
   private final AudioMixer mixer;
 
   private AudioFormat mixerAudioFormat;
+  private long pendingStartTimeUs;
+  private int mixerSourcesToAdd;
   private ByteBuffer mixerOutput;
   private AudioProcessingPipeline audioProcessingPipeline;
   private int finishedInputs;
@@ -134,6 +136,12 @@ import java.util.Objects;
    * unless the graph was {@linkplain #flush() flushed}.
    */
   public ByteBuffer getOutput() throws ExportException {
+    if (mixerSourcesToAdd > 0) {
+      addMixerSources();
+      if (mixerSourcesToAdd > 0) {
+        return EMPTY_BUFFER;
+      }
+    }
     if (!mixer.isEnded()) {
       feedMixer();
     }
@@ -149,6 +157,28 @@ import java.util.Objects;
     return mixerOutput;
   }
 
+  /** Instructs the {@code AudioGraph} to not queue any input buffer. */
+  public void blockInput() {
+    for (int i = 0; i < inputInfos.size(); i++) {
+      inputInfos.get(i).audioGraphInput.blockInput();
+    }
+  }
+
+  /** Unblocks incoming data if {@linkplain #blockInput() blocked}. */
+  public void unblockInput() {
+    for (int i = 0; i < inputInfos.size(); i++) {
+      inputInfos.get(i).audioGraphInput.unblockInput();
+    }
+  }
+
+  /**
+   * Sets the start time of the audio streams that will enter the audio graph after the next calls
+   * to {@link #flush()}, in microseconds.
+   */
+  public void setPendingStartTimeUs(long startTimeUs) {
+    this.pendingStartTimeUs = startTimeUs;
+  }
+
   /** Clears any pending data. */
   public void flush() {
     for (int i = 0; i < inputInfos.size(); i++) {
@@ -158,12 +188,12 @@ import java.util.Objects;
     }
     mixer.reset();
     try {
-      mixer.configure(mixerAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
-      addMixerSources();
+      mixer.configure(mixerAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, pendingStartTimeUs);
     } catch (UnhandledAudioFormatException e) {
-      // Should never happen because mixer has already been configured with the same formats.
+      // Should never happen because mixer has already been configured with the same format.
       Log.e(TAG, "Unexpected mixer configuration error");
     }
+    mixerSourcesToAdd = inputInfos.size();
     mixerOutput = EMPTY_BUFFER;
     audioProcessingPipeline.flush();
     finishedInputs = 0;
@@ -207,6 +237,33 @@ import java.util.Objects;
     audioProcessingPipeline.queueInput(mixerOutput);
   }
 
+  private void addMixerSources() throws ExportException {
+    for (int i = 0; i < inputInfos.size(); i++) {
+      InputInfo inputInfo = inputInfos.get(i);
+      if (inputInfo.mixerSourceId != C.INDEX_UNSET) {
+        continue; // The source has already been added.
+      }
+      AudioGraphInput audioGraphInput = inputInfo.audioGraphInput;
+      try {
+        // Force processing input.
+        audioGraphInput.getOutput();
+        long sourceStartTimeUs = audioGraphInput.getStartTimeUs();
+        if (sourceStartTimeUs == C.TIME_UNSET) {
+          continue;
+        } else if (sourceStartTimeUs == C.TIME_END_OF_SOURCE) {
+          mixerSourcesToAdd--;
+          continue;
+        }
+        inputInfo.mixerSourceId =
+            mixer.addSource(audioGraphInput.getOutputAudioFormat(), sourceStartTimeUs);
+      } catch (UnhandledAudioFormatException e) {
+        throw ExportException.createForAudioProcessing(
+            e, "Unhandled format while adding source " + inputInfo.mixerSourceId);
+      }
+      mixerSourcesToAdd--;
+    }
+  }
+
   private void feedMixer() throws ExportException {
     for (int i = 0; i < inputInfos.size(); i++) {
       feedMixerFromInput(inputInfos.get(i));
@@ -235,14 +292,6 @@ import java.util.Objects;
     }
   }
 
-  private void addMixerSources() throws UnhandledAudioFormatException {
-    for (int i = 0; i < inputInfos.size(); i++) {
-      InputInfo inputInfo = inputInfos.get(i);
-      inputInfo.mixerSourceId =
-          mixer.addSource(inputInfo.audioGraphInput.getOutputAudioFormat(), /* startTimeUs= */ 0);
-    }
-  }
-
   private static final class InputInfo {
     public final AudioGraphInput audioGraphInput;
     public int mixerSourceId;

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

@@ -70,6 +70,8 @@ import java.util.concurrent.atomic.AtomicReference;
   private boolean processedFirstMediaItemChange;
   private boolean receivedEndOfStreamFromInput;
   private boolean queueEndOfStreamAfterSilence;
+  private long startTimeUs;
+  private boolean inputBlocked;
 
   /**
    * Creates an instance.
@@ -102,6 +104,7 @@ import java.util.concurrent.atomic.AtomicReference;
     // APP configuration not active until flush called. getOutputAudioFormat based on active config.
     audioProcessingPipeline.flush();
     outputAudioFormat = audioProcessingPipeline.getOutputAudioFormat();
+    startTimeUs = C.TIME_UNSET;
   }
 
   /** Returns the {@link AudioFormat} of {@linkplain #getOutput() output buffers}. */
@@ -164,7 +167,7 @@ import java.util.concurrent.atomic.AtomicReference;
   @Override
   @Nullable
   public DecoderInputBuffer getInputBuffer() {
-    if (pendingMediaItemChange.get() != null) {
+    if (inputBlocked || (pendingMediaItemChange.get() != null)) {
       return null;
     }
     return availableInputBuffers.peek();
@@ -177,16 +180,52 @@ import java.util.concurrent.atomic.AtomicReference;
    */
   @Override
   public boolean queueInputBuffer() {
+    if (inputBlocked) {
+      return false;
+    }
     checkState(pendingMediaItemChange.get() == null);
     DecoderInputBuffer inputBuffer = availableInputBuffers.remove();
     pendingInputBuffers.add(inputBuffer);
+    if (startTimeUs == C.TIME_UNSET) {
+      startTimeUs = inputBuffer.timeUs;
+    }
     return true;
   }
 
   /**
-   * Clears any pending input and output data.
+   * Returns the stream start time in microseconds, or {@link C#TIME_UNSET} if unknown.
    *
-   * <p>Should only be called by the processing thread.
+   * <p>Should only be called if the input thread and processing thread are the same.
+   */
+  public long getStartTimeUs() {
+    return startTimeUs;
+  }
+
+  /**
+   * Instructs the {@code AudioGraphInput} to not queue any input buffer.
+   *
+   * <p>Should only be called if the input thread and processing thread are the same.
+   */
+  public void blockInput() {
+    inputBlocked = true;
+  }
+
+  /**
+   * Unblocks incoming data if {@linkplain #blockInput() blocked}.
+   *
+   * <p>Should only be called if the input thread and processing thread are the same.
+   */
+  public void unblockInput() {
+    inputBlocked = false;
+  }
+
+  /**
+   * Clears any pending data.
+   *
+   * <p>If an {@linkplain #getInputBuffer() input buffer} has been retrieved without being queued,
+   * it shouldn't be used after calling this method.
+   *
+   * <p>Should only be called if the input thread and processing thread are the same.
    */
   public void flush() {
     pendingMediaItemChange.set(null);
@@ -204,6 +243,7 @@ import java.util.concurrent.atomic.AtomicReference;
     currentInputBufferBeingOutput = null;
     receivedEndOfStreamFromInput = false;
     queueEndOfStreamAfterSilence = false;
+    startTimeUs = C.TIME_UNSET;
   }
 
   /**

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

@@ -368,6 +368,64 @@ public class AudioGraphInputTest {
     assertThat(outputBytes).containsExactlyElementsIn(Bytes.asList(inputData));
   }
 
+  @Test
+  public void blockInput_blocksInputData() throws Exception {
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(
+            /* requestedOutputAudioFormat= */ AudioFormat.NOT_SET,
+            /* editedMediaItem= */ FAKE_ITEM,
+            /* inputFormat= */ getPcmFormat(STEREO_44100));
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+
+    // Force the media item change to be processed.
+    checkState(!audioGraphInput.getOutput().hasRemaining());
+
+    // Queue inputData.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+
+    audioGraphInput.blockInput();
+
+    assertThat(audioGraphInput.queueInputBuffer()).isFalse();
+    assertThat(audioGraphInput.getInputBuffer()).isNull();
+  }
+
+  @Test
+  public void unblockInput_unblocksInputData() throws Exception {
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(
+            /* requestedOutputAudioFormat= */ AudioFormat.NOT_SET,
+            /* editedMediaItem= */ FAKE_ITEM,
+            /* inputFormat= */ getPcmFormat(STEREO_44100));
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+
+    // Force the media item change to be processed.
+    checkState(!audioGraphInput.getOutput().hasRemaining());
+
+    // Queue inputData.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+
+    audioGraphInput.blockInput();
+    audioGraphInput.unblockInput();
+
+    assertThat(audioGraphInput.queueInputBuffer()).isTrue();
+  }
+
   /** Drains the graph and returns the bytes output. */
   private static List<Byte> drainAudioGraphInputUntilEnded(AudioGraphInput audioGraphInput)
       throws Exception {

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

@@ -186,6 +186,112 @@ public class AudioGraphTest {
         () -> audioGraph.configure(ImmutableList.of(sonicAudioProcessor)));
   }
 
+  @Test
+  public void blockInput_blocksInputData() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    audioGraph.configure(ImmutableList.of());
+    AudioGraphInput audioGraphInput =
+        audioGraph.registerInput(FAKE_ITEM, getPcmFormat(STEREO_44100));
+    audioGraphInput.onMediaItemChanged(
+        FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+    audioGraphInput.getOutput(); // Force the media item change to be processed.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+
+    audioGraph.blockInput();
+
+    assertThat(audioGraphInput.queueInputBuffer()).isFalse();
+  }
+
+  @Test
+  public void unblockInput_unblocksInputData() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    audioGraph.configure(ImmutableList.of());
+    AudioGraphInput audioGraphInput =
+        audioGraph.registerInput(FAKE_ITEM, getPcmFormat(STEREO_44100));
+    audioGraphInput.onMediaItemChanged(
+        FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+    audioGraphInput.getOutput(); // Force the media item change to be processed.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+    audioGraph.blockInput();
+
+    audioGraph.unblockInput();
+
+    assertThat(audioGraphInput.queueInputBuffer()).isTrue();
+  }
+
+  @Test
+  public void setPendingStartTimeUs_discardsPrecedingData() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    audioGraph.configure(ImmutableList.of());
+    AudioGraphInput audioGraphInput =
+        audioGraph.registerInput(FAKE_ITEM, getPcmFormat(STEREO_44100));
+    audioGraphInput.onMediaItemChanged(
+        FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+    audioGraphInput.getOutput(); // Force the media item change to be processed.
+
+    audioGraph.setPendingStartTimeUs(500_000);
+    audioGraph.flush();
+    // Queue input buffer with timestamp 0.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+    checkState(audioGraphInput.queueInputBuffer());
+    // Queue EOS.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer());
+    // Drain output.
+    int bytesOutput = drainAudioGraph(audioGraph);
+
+    assertThat(bytesOutput).isEqualTo(0);
+  }
+
+  @Test
+  public void setPendingStartTimeUs_doesNotDiscardFollowingData() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    audioGraph.configure(ImmutableList.of());
+    AudioGraphInput audioGraphInput =
+        audioGraph.registerInput(FAKE_ITEM, getPcmFormat(STEREO_44100));
+    audioGraphInput.onMediaItemChanged(
+        FAKE_ITEM,
+        /* durationUs= */ 1_000_000,
+        /* decodedFormat= */ getPcmFormat(STEREO_44100),
+        /* isLast= */ true);
+    audioGraphInput.getOutput(); // Force the media item change to be processed.
+
+    audioGraph.setPendingStartTimeUs(500_000);
+    audioGraph.flush();
+    // Queue input buffer with timestamp 600 ms.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    byte[] inputData = TestUtil.buildTestData(/* length= */ 100 * STEREO_44100.bytesPerFrame);
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+    inputBuffer.timeUs = 600_000;
+    checkState(audioGraphInput.queueInputBuffer());
+    // Queue EOS.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer());
+    // Drain output.
+    int bytesOutput = drainAudioGraph(audioGraph);
+
+    assertThat(bytesOutput).isGreaterThan(0);
+  }
+
   @Test
   public void flush_withoutAudioProcessor_clearsPendingData() throws Exception {
     AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());