Implement AudioGraphInput.flush

PiperOrigin-RevId: 613563855

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

@@ -183,13 +183,35 @@ import java.util.concurrent.atomic.AtomicReference;
     return true;
   }
 
+  /**
+   * Clears any pending input and output data.
+   *
+   * <p>Should only be called by the processing thread.
+   */
+  public void flush() {
+    pendingMediaItemChange.set(null);
+    processedFirstMediaItemChange = true;
+    if (!availableInputBuffers.isEmpty()) {
+      // Clear first available buffer in case the caller wrote data in the input buffer without
+      // queueing it.
+      clearAndAddToAvailableBuffers(availableInputBuffers.remove());
+    }
+    while (!pendingInputBuffers.isEmpty()) {
+      clearAndAddToAvailableBuffers(pendingInputBuffers.remove());
+    }
+    silentAudioGenerator.flush();
+    audioProcessingPipeline.flush();
+    currentInputBufferBeingOutput = null;
+    receivedEndOfStreamFromInput = false;
+    queueEndOfStreamAfterSilence = false;
+  }
+
   /**
    * Releases any underlying resources.
    *
    * <p>Should only be called by the processing thread.
    */
   public void release() {
-    // TODO(b/303029174): Impl flush(), reset() & decide if a separate release() is still needed.
     audioProcessingPipeline.reset();
   }
 

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

@@ -73,4 +73,10 @@ import java.util.concurrent.atomic.AtomicLong;
   public boolean hasRemaining() {
     return internalBuffer.hasRemaining() || remainingBytesToOutput.get() > 0;
   }
+
+  public void flush() {
+    remainingBytesToOutput.set(0);
+    internalBuffer.position(0);
+    internalBuffer.limit(0);
+  }
 }

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

@@ -107,6 +107,47 @@ public class AudioGraphInputTest {
     assertThat(audioGraphInput.getOutputAudioFormat()).isEqualTo(STEREO_44100);
   }
 
+  @Test
+  public void getOutputAudioFormat_afterFlush_isSet() throws Exception {
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(
+            /* requestedOutputAudioFormat= */ STEREO_44100,
+            /* editedMediaItem= */ FAKE_ITEM,
+            /* inputFormat= */ getPcmFormat(MONO_48000));
+
+    audioGraphInput.flush();
+
+    assertThat(audioGraphInput.getOutputAudioFormat()).isEqualTo(STEREO_44100);
+  }
+
+  @Test
+  public void getInputBuffer_afterFlush_returnsEmptyBuffer() 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());
+
+    // Fill input buffer.
+    DecoderInputBuffer inputBuffer = audioGraphInput.getInputBuffer();
+    inputBuffer.ensureSpaceForWrite(inputData.length);
+    inputBuffer.data.put(inputData).flip();
+
+    audioGraphInput.flush();
+
+    assertThat(audioGraphInput.getInputBuffer().data.remaining()).isEqualTo(0);
+  }
+
   @Test
   public void isEnded_whenInitialized_returnsFalse() throws Exception {
     AudioGraphInput audioGraphInput =
@@ -143,6 +184,35 @@ public class AudioGraphInputTest {
     assertThat(audioGraphInput.isEnded()).isTrue();
   }
 
+  @Test
+  public void isEnded_afterFlush_returnsFalse() throws Exception {
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(
+            /* requestedOutputAudioFormat= */ AudioFormat.NOT_SET,
+            /* editedMediaItem= */ FAKE_ITEM,
+            /* inputFormat= */ getPcmFormat(MONO_44100));
+
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ C.TIME_UNSET,
+        /* decodedFormat= */ getPcmFormat(MONO_44100),
+        /* isLast= */ false);
+
+    // Force the media item change to be processed.
+    checkState(!audioGraphInput.getOutput().hasRemaining());
+
+    // Queue EOS.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer());
+
+    drainAudioGraphInputUntilEnded(audioGraphInput);
+    checkState(audioGraphInput.isEnded());
+
+    audioGraphInput.flush();
+
+    assertThat(audioGraphInput.isEnded()).isFalse();
+  }
+
   @Test
   public void getOutput_withoutMediaItemChange_returnsEmptyBuffer() throws Exception {
     AudioGraphInput audioGraphInput =
@@ -224,6 +294,80 @@ public class AudioGraphInputTest {
     assertThat(bytesOutput).isEqualTo(expectedSampleCount * STEREO_44100.bytesPerFrame);
   }
 
+  @Test
+  public void getOutput_afterFlush_returnsEmptyBuffer() 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();
+    checkState(audioGraphInput.queueInputBuffer());
+
+    audioGraphInput.flush();
+
+    // Queue EOS.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer());
+
+    List<Byte> outputBytes = drainAudioGraphInputUntilEnded(audioGraphInput);
+    assertThat(outputBytes).isEmpty();
+  }
+
+  @Test
+  public void getOutput_afterFlushAndInput_returnsCorrectAmountOfBytes() 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();
+    checkState(audioGraphInput.queueInputBuffer());
+
+    audioGraphInput.flush();
+
+    // Queue inputData.
+    inputBuffer = audioGraphInput.getInputBuffer();
+    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());
+
+    List<Byte> outputBytes = drainAudioGraphInputUntilEnded(audioGraphInput);
+    assertThat(outputBytes).containsExactlyElementsIn(Bytes.asList(inputData));
+  }
+
   /** 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/SilentAudioGeneratorTest.java

@@ -89,6 +89,33 @@ public class SilentAudioGeneratorTest {
     assertThat(generator.getBuffer().remaining()).isEqualTo(960);
   }
 
+  @Test
+  public void addSilence_afterFlush_producesCorrectNumberOfBytes() {
+    SilentAudioGenerator generator =
+        new SilentAudioGenerator(
+            new AudioFormat(/* sampleRate= */ 88_200, /* channelCount= */ 6, C.ENCODING_PCM_16BIT));
+
+    generator.addSilence(/* durationUs= */ 3_000_000);
+    generator.flush();
+    generator.addSilence(/* durationUs= */ 1_500_000);
+    int bytesOutput = drainGenerator(generator);
+
+    // 88_200 * 12 * 1.5s = 1_587_600
+    assertThat(bytesOutput).isEqualTo(1_587_600);
+  }
+
+  @Test
+  public void hasRemaining_afterFlush_isFalse() {
+    SilentAudioGenerator generator =
+        new SilentAudioGenerator(
+            new AudioFormat(/* sampleRate= */ 88_200, /* channelCount= */ 6, C.ENCODING_PCM_16BIT));
+
+    generator.addSilence(/* durationUs= */ 3_000_000);
+    generator.flush();
+
+    assertThat(generator.hasRemaining()).isFalse();
+  }
+
   /** Drains the generator and returns the number of bytes output. */
   private static int drainGenerator(SilentAudioGenerator generator) {
     int bytesOutput = 0;