Fix a race condition in AudioGraphInput

AudioGraphInput.onMediaItemChanged is called on input thread. Pending media item changes are processed on processing thread, inside calls to getOutput(). This change allows multiple pending media item changes to be enqueued, and processed in sequence. PiperOrigin-RevId: 638995291
androidx · May 31, 2024 · d3fa332 · d3fa332
1 parent 4dd8360
commit d3fa332
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 14 deletions.
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/AudioGraphInput.java b/libraries/transformer/src/main/java/androidx/media3/transformer/AudioGraphInput.java
@@ -36,7 +36,6 @@
 import androidx.media3.common.audio.ChannelMixingMatrix;
 import androidx.media3.common.audio.SonicAudioProcessor;
 import androidx.media3.common.audio.SpeedChangingAudioProcessor;
-import androidx.media3.common.util.NullableType;
 import androidx.media3.common.util.Util;
 import androidx.media3.decoder.DecoderInputBuffer;
 import com.google.common.collect.ImmutableList;
@@ -45,7 +44,6 @@
 import java.util.Queue;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Processes a single sequential stream of PCM audio samples.
@@ -64,7 +62,7 @@
 
   private final Queue<DecoderInputBuffer> availableInputBuffers;
   private final Queue<DecoderInputBuffer> pendingInputBuffers;
-  private final AtomicReference<@NullableType MediaItemChange> pendingMediaItemChange;
+  private final Queue<MediaItemChange> pendingMediaItemChanges;
   private final AtomicLong startTimeUs;
 
   // silentAudioGenerator.audioFormat must match the current media item's input format.
@@ -103,7 +101,7 @@ public AudioGraphInput(
       availableInputBuffers.add(inputBuffer);
     }
     pendingInputBuffers = new ConcurrentLinkedQueue<>();
-    pendingMediaItemChange = new AtomicReference<>();
+    pendingMediaItemChanges = new ConcurrentLinkedQueue<>();
     silentAudioGenerator = new SilentAudioGenerator(inputAudioFormat);
     audioProcessingPipeline =
         configureProcessing(
@@ -138,7 +136,7 @@ public ByteBuffer getOutput() throws UnhandledAudioFormatException {
       return outputBuffer;
     }
 
-    if (!hasDataToOutput() && pendingMediaItemChange.get() != null) {
+    if (!hasDataToOutput() && !pendingMediaItemChanges.isEmpty()) {
       configureForPendingMediaItemChange();
     }
 
@@ -167,7 +165,7 @@ public void onMediaItemChanged(
       AudioFormat audioFormat = new AudioFormat(decodedFormat);
       checkState(isInputAudioFormatValid(audioFormat), /* errorMessage= */ audioFormat);
     }
-    pendingMediaItemChange.set(
+    pendingMediaItemChanges.add(
         new MediaItemChange(editedMediaItem, durationUs, decodedFormat, isLast));
   }
 
@@ -179,7 +177,7 @@ public void onMediaItemChanged(
   @Override
   @Nullable
   public DecoderInputBuffer getInputBuffer() {
-    if (inputBlocked || (pendingMediaItemChange.get() != null)) {
+    if (inputBlocked || !pendingMediaItemChanges.isEmpty()) {
       return null;
     }
     return availableInputBuffers.peek();
@@ -195,7 +193,7 @@ public boolean queueInputBuffer() {
     if (inputBlocked) {
       return false;
     }
-    checkState(pendingMediaItemChange.get() == null);
+    checkState(pendingMediaItemChanges.isEmpty());
     DecoderInputBuffer inputBuffer = availableInputBuffers.remove();
     pendingInputBuffers.add(inputBuffer);
     startTimeUs.compareAndSet(
@@ -235,7 +233,7 @@ public void unblockInput() {
    * <p>Should only be called if the input thread and processing thread are the same.
    */
   public void flush() {
-    pendingMediaItemChange.set(null);
+    pendingMediaItemChanges.clear();
     processedFirstMediaItemChange = true;
     if (!availableInputBuffers.isEmpty()) {
       // Clear first available buffer in case the caller wrote data in the input buffer without
@@ -279,7 +277,7 @@ public boolean isEnded() {
     if (hasDataToOutput()) {
       return false;
     }
-    if (pendingMediaItemChange.get() != null) {
+    if (!pendingMediaItemChanges.isEmpty()) {
       return false;
     }
     if (currentItemExpectedInputDurationUs != C.TIME_UNSET) {
@@ -323,7 +321,7 @@ private boolean feedProcessingPipelineFromInput() {
 
     @Nullable DecoderInputBuffer pendingInputBuffer = pendingInputBuffers.peek();
     if (pendingInputBuffer == null) {
-      if (pendingMediaItemChange.get() != null) {
+      if (!pendingMediaItemChanges.isEmpty()) {
         if (shouldAppendSilence()) {
           appendSilence();
           return true;
@@ -377,7 +375,7 @@ private ByteBuffer feedOutputFromInput() {
 
     @Nullable DecoderInputBuffer currentInputBuffer = pendingInputBuffers.poll();
     if (currentInputBuffer == null) {
-      if (pendingMediaItemChange.get() != null && shouldAppendSilence()) {
+      if (!pendingMediaItemChanges.isEmpty() && shouldAppendSilence()) {
         appendSilence();
       }
       return EMPTY_BUFFER;
@@ -439,7 +437,7 @@ private void clearAndAddToAvailableBuffers(DecoderInputBuffer inputBuffer) {
    * through {@link #getOutput()}.
    */
   private void configureForPendingMediaItemChange() throws UnhandledAudioFormatException {
-    MediaItemChange pendingChange = checkStateNotNull(pendingMediaItemChange.get());
+    MediaItemChange pendingChange = checkStateNotNull(pendingMediaItemChanges.poll());
 
     currentItemInputBytesRead = 0;
     isCurrentItemLast = pendingChange.isLast;
@@ -476,7 +474,6 @@ private void configureForPendingMediaItemChange() throws UnhandledAudioFormatExc
               /* requiredOutputAudioFormat= */ outputAudioFormat);
     }
     audioProcessingPipeline.flush();
-    pendingMediaItemChange.set(null);
     receivedEndOfStreamFromInput = false;
     processedFirstMediaItemChange = true;
   }

diff --git a/libraries/transformer/src/test/java/androidx/media3/transformer/AudioGraphInputTest.java b/libraries/transformer/src/test/java/androidx/media3/transformer/AudioGraphInputTest.java
@@ -481,6 +481,36 @@ public void getOutput_withSilentMediaItemChange_outputsCorrectAmountOfSilentByte
     assertThat(bytesOutput).isEqualTo(expectedSampleCount * STEREO_44100.bytesPerFrame);
   }
 
+  @Test
+  public void getOutput_withThreeSilentMediaItemChanges_outputsCorrectAmountOfSilentBytes()
+      throws Exception {
+    AudioGraphInput audioGraphInput =
+        new AudioGraphInput(
+            /* requestedOutputAudioFormat= */ AudioFormat.NOT_SET,
+            /* editedMediaItem= */ FAKE_ITEM,
+            /* inputFormat= */ getPcmFormat(STEREO_44100));
+
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ 200_000,
+        /* decodedFormat= */ null,
+        /* isLast= */ false);
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ 300_000,
+        /* decodedFormat= */ null,
+        /* isLast= */ false);
+    audioGraphInput.onMediaItemChanged(
+        /* editedMediaItem= */ FAKE_ITEM,
+        /* durationUs= */ 500_000,
+        /* decodedFormat= */ null,
+        /* isLast= */ true);
+
+    int bytesOutput = drainAudioGraphInputUntilEnded(audioGraphInput).size();
+    long expectedSampleCount = Util.durationUsToSampleCount(1_000_000, STEREO_44100.sampleRate);
+    assertThat(bytesOutput).isEqualTo(expectedSampleCount * STEREO_44100.bytesPerFrame);
+  }
+
   @Test
   public void getOutput_withSilentMediaItemAndEffectsChange_outputsCorrectAmountOfSilentBytes()
       throws Exception {