Implement AudioGraph.flush()

When flushing the AudioGraph, the AudioMixer will be preserved but the
sources will be recreated. This is probably a bit less efficient but
makes the logic simpler. Indeed, if we had to keep the sources alive,
we would need to add a way to reconfigure them with a new timestamp for
seeking. We would also need to change the way sources are ended because
they are currently removed when they are ended. Also, it is acceptable
to have a small delay when seeking, which means that performance is less
critical than for playback.

PiperOrigin-RevId: 615775501

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

@@ -20,34 +20,39 @@ import static androidx.media3.common.audio.AudioProcessor.EMPTY_BUFFER;
 import static androidx.media3.common.util.Assertions.checkArgument;
 import static androidx.media3.common.util.Assertions.checkState;
 
-import android.util.SparseArray;
 import androidx.media3.common.C;
 import androidx.media3.common.Format;
 import androidx.media3.common.audio.AudioProcessingPipeline;
 import androidx.media3.common.audio.AudioProcessor;
 import androidx.media3.common.audio.AudioProcessor.AudioFormat;
 import androidx.media3.common.audio.AudioProcessor.UnhandledAudioFormatException;
+import androidx.media3.common.util.Log;
 import com.google.common.collect.ImmutableList;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
+import java.util.List;
 import java.util.Objects;
 
 /** Processes raw audio samples. */
 /* package */ final class AudioGraph {
-  private final AudioMixer mixer;
-  private final SparseArray<AudioGraphInput> inputs;
 
-  private AudioProcessingPipeline audioProcessingPipeline;
+  private static final String TAG = "AudioGraph";
+
+  private final List<InputInfo> inputInfos;
+  private final AudioMixer mixer;
+
   private AudioFormat mixerAudioFormat;
-  private int finishedInputs;
   private ByteBuffer mixerOutput;
+  private AudioProcessingPipeline audioProcessingPipeline;
+  private int finishedInputs;
 
   /** Creates an instance. */
   public AudioGraph(AudioMixer.Factory mixerFactory) {
+    inputInfos = new ArrayList<>();
     mixer = mixerFactory.create();
-    inputs = new SparseArray<>();
-    audioProcessingPipeline = new AudioProcessingPipeline(ImmutableList.of());
-    mixerOutput = EMPTY_BUFFER;
     mixerAudioFormat = AudioFormat.NOT_SET;
+    mixerOutput = EMPTY_BUFFER;
+    audioProcessingPipeline = new AudioProcessingPipeline(ImmutableList.of());
   }
 
   /** Returns whether an {@link AudioFormat} is valid as an input format. */
@@ -89,21 +94,26 @@ import java.util.Objects;
   public AudioGraphInput registerInput(EditedMediaItem editedMediaItem, Format format)
       throws ExportException {
     checkArgument(format.pcmEncoding != Format.NO_VALUE);
+    AudioGraphInput audioGraphInput;
+    int sourceId;
     try {
-      AudioGraphInput audioGraphInput =
+      audioGraphInput = new AudioGraphInput(mixerAudioFormat, editedMediaItem, format);
-          new AudioGraphInput(mixerAudioFormat, editedMediaItem, format);
 
       if (Objects.equals(mixerAudioFormat, AudioFormat.NOT_SET)) {
         // Mixer not configured, configure before doing anything else.
-        configureMixer(audioGraphInput.getOutputAudioFormat());
+        this.mixerAudioFormat = audioGraphInput.getOutputAudioFormat();
+        mixer.configure(mixerAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
+        audioProcessingPipeline.configure(mixerAudioFormat);
+        audioProcessingPipeline.flush();
       }
 
-      int sourceId = mixer.addSource(audioGraphInput.getOutputAudioFormat(), /* startTimeUs= */ 0);
+      sourceId = mixer.addSource(audioGraphInput.getOutputAudioFormat(), /* startTimeUs= */ 0);
-      inputs.append(sourceId, audioGraphInput);
-      return audioGraphInput;
     } catch (UnhandledAudioFormatException e) {
-      throw ExportException.createForAudioProcessing(e, "existingInputs=" + inputs.size());
+      throw ExportException.createForAudioProcessing(
+          e, "Error while registering input " + inputInfos.size());
     }
+    inputInfos.add(new InputInfo(audioGraphInput, sourceId));
+    return audioGraphInput;
   }
 
   /**
@@ -118,7 +128,8 @@ import java.util.Objects;
   /**
    * Returns a {@link ByteBuffer} containing output data between the position and limit.
    *
-   * <p>The same buffer is returned until it has been fully consumed ({@code position == limit}).
+   * <p>The same buffer is returned until it has been fully consumed ({@code position == limit}),
+   * unless the graph was {@linkplain #flush() flushed}.
    */
   public ByteBuffer getOutput() throws ExportException {
     if (!mixer.isEnded()) {
@@ -136,16 +147,36 @@ import java.util.Objects;
     return mixerOutput;
   }
 
+  /** Clears any pending data. */
+  public void flush() {
+    for (int i = 0; i < inputInfos.size(); i++) {
+      InputInfo inputInfo = inputInfos.get(i);
+      inputInfo.mixerSourceId = C.INDEX_UNSET;
+      inputInfo.audioGraphInput.flush();
+    }
+    mixer.reset();
+    try {
+      mixer.configure(mixerAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
+      addMixerSources();
+    } catch (UnhandledAudioFormatException e) {
+      // Should never happen because mixer has already been configured with the same formats.
+      Log.e(TAG, "Unexpected mixer configuration error");
+    }
+    mixerOutput = EMPTY_BUFFER;
+    audioProcessingPipeline.flush();
+    finishedInputs = 0;
+  }
+
   /**
    * Resets the graph, un-registering inputs and releasing any underlying resources.
    *
    * <p>Call {@link #registerInput(EditedMediaItem, Format)} to prepare the audio graph again.
    */
   public void reset() {
-    for (int i = 0; i < inputs.size(); i++) {
+    for (int i = 0; i < inputInfos.size(); i++) {
-      inputs.valueAt(i).release();
+      inputInfos.get(i).audioGraphInput.release();
     }
-    inputs.clear();
+    inputInfos.clear();
     mixer.reset();
     audioProcessingPipeline.reset();
 
@@ -162,24 +193,8 @@ import java.util.Objects;
     return isMixerEnded();
   }
 
-  /**
-   * Configures the mixer.
-   *
-   * <p>Must be called before {@linkplain #getOutput() accessing output}.
-   *
-   * @param mixerAudioFormat The {@link AudioFormat} requested for output from the mixer.
-   * @throws UnhandledAudioFormatException If the audio format is not supported by the {@link
-   *     AudioMixer}.
-   */
-  private void configureMixer(AudioFormat mixerAudioFormat) throws UnhandledAudioFormatException {
-    this.mixerAudioFormat = mixerAudioFormat;
-    mixer.configure(mixerAudioFormat, /* bufferSizeMs= */ C.LENGTH_UNSET, /* startTimeUs= */ 0);
-    audioProcessingPipeline.configure(mixerAudioFormat);
-    audioProcessingPipeline.flush();
-  }
-
   private boolean isMixerEnded() {
-    return !mixerOutput.hasRemaining() && finishedInputs >= inputs.size() && mixer.isEnded();
+    return !mixerOutput.hasRemaining() && finishedInputs >= inputInfos.size() && mixer.isEnded();
   }
 
   private void feedProcessingPipelineFromMixer() {
@@ -191,18 +206,21 @@ import java.util.Objects;
   }
 
   private void feedMixer() throws ExportException {
-    for (int i = 0; i < inputs.size(); i++) {
+    for (int i = 0; i < inputInfos.size(); i++) {
-      feedMixerFromInput(inputs.keyAt(i), inputs.valueAt(i));
+      feedMixerFromInput(inputInfos.get(i));
     }
   }
 
-  private void feedMixerFromInput(int sourceId, AudioGraphInput input) throws ExportException {
+  private void feedMixerFromInput(InputInfo inputInfo) throws ExportException {
+    int sourceId = inputInfo.mixerSourceId;
     if (!mixer.hasSource(sourceId)) {
       return;
     }
 
+    AudioGraphInput input = inputInfo.audioGraphInput;
     if (input.isEnded()) {
       mixer.removeSource(sourceId);
+      inputInfo.mixerSourceId = C.INDEX_UNSET;
       finishedInputs++;
       return;
     }
@@ -214,4 +232,22 @@ import java.util.Objects;
           e, "AudioGraphInput (sourceId=" + sourceId + ") reconfiguration");
     }
   }
+
+  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;
+
+    public InputInfo(AudioGraphInput audioGraphInput, int mixerSourceId) {
+      this.audioGraphInput = audioGraphInput;
+      this.mixerSourceId = mixerSourceId;
+    }
+  }
 }

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

@@ -15,6 +15,7 @@
  */
 package androidx.media3.transformer;
 
+import static androidx.media3.common.util.Assertions.checkState;
 import static androidx.media3.common.util.Util.getPcmFormat;
 import static com.google.common.truth.Truth.assertThat;
 import static org.junit.Assert.assertThrows;
@@ -23,6 +24,8 @@ import androidx.media3.common.C;
 import androidx.media3.common.MediaItem;
 import androidx.media3.common.audio.AudioProcessor.AudioFormat;
 import androidx.media3.common.audio.SonicAudioProcessor;
+import androidx.media3.decoder.DecoderInputBuffer;
+import androidx.media3.test.utils.TestUtil;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import com.google.common.collect.ImmutableList;
 import java.nio.ByteBuffer;
@@ -92,6 +95,16 @@ public class AudioGraphTest {
     assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(MONO_48000);
   }
 
+  @Test
+  public void getOutputAudioFormat_afterFlush_isSet() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    audioGraph.registerInput(FAKE_ITEM, getPcmFormat(MONO_48000));
+
+    audioGraph.flush();
+
+    assertThat(audioGraph.getOutputAudioFormat()).isEqualTo(MONO_48000);
+  }
+
   @Test
   public void registerInput_afterRegisterInput_doesNotChangeOutputFormat() throws Exception {
     AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
@@ -164,6 +177,108 @@ public class AudioGraphTest {
         () -> audioGraph.configure(ImmutableList.of(sonicAudioProcessor)));
   }
 
+  @Test
+  public void flush_withoutAudioProcessor_clearsPendingData() 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();
+    checkState(audioGraphInput.queueInputBuffer());
+    checkState(audioGraph.getOutput().hasRemaining());
+
+    audioGraph.flush();
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer()); // Queue EOS.
+    int bytesOutput = drainAudioGraph(audioGraph);
+
+    assertThat(bytesOutput).isEqualTo(0);
+  }
+
+  @Test
+  public void flush_withAudioProcessor_clearsPendingData() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    SonicAudioProcessor sonicAudioProcessor = new SonicAudioProcessor();
+    sonicAudioProcessor.setOutputSampleRateHz(48_000);
+    audioGraph.configure(ImmutableList.of(sonicAudioProcessor));
+    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();
+    checkState(audioGraphInput.queueInputBuffer());
+    checkState(audioGraph.getOutput().hasRemaining());
+
+    audioGraph.flush();
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer()); // Queue EOS.
+    int bytesOutput = drainAudioGraph(audioGraph);
+
+    assertThat(bytesOutput).isEqualTo(0);
+  }
+
+  @Test
+  public void isEnded_afterFlushAndWithoutAudioProcessor_isFalse() 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.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer()); // Queue EOS.
+    drainAudioGraph(audioGraph);
+    checkState(audioGraph.isEnded());
+
+    audioGraph.flush();
+
+    assertThat(audioGraph.isEnded()).isFalse();
+  }
+
+  @Test
+  public void isEnded_afterFlushAndWithAudioProcessor_isFalse() throws Exception {
+    AudioGraph audioGraph = new AudioGraph(new DefaultAudioMixer.Factory());
+    SonicAudioProcessor sonicAudioProcessor = new SonicAudioProcessor();
+    sonicAudioProcessor.setOutputSampleRateHz(48_000);
+    audioGraph.configure(ImmutableList.of(sonicAudioProcessor));
+    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.
+    audioGraphInput.getInputBuffer().setFlags(C.BUFFER_FLAG_END_OF_STREAM);
+    checkState(audioGraphInput.queueInputBuffer()); // Queue EOS.
+    drainAudioGraph(audioGraph);
+    checkState(audioGraph.isEnded());
+
+    audioGraph.flush();
+
+    assertThat(audioGraph.isEnded()).isFalse();
+  }
+
   /** Drains the graph and returns the number of bytes output. */
   private static int drainAudioGraph(AudioGraph audioGraph) throws ExportException {
     int bytesOutput = 0;