From 04a388999832cf7cf4c99867b54d84bbe6816497 Mon Sep 17 00:00:00 2001
From: kimvde <kimvde@google.com>
Date: Thu, 28 Mar 2024 01:19:19 -0700
Subject: [PATCH] Reset maxPositionOfRemovedSources in DefaultAudioMixer

If the mixer is reset without resetting maxPositionOfRemovedSources and
then reused, the value of maxPositionOfRemovedSources can be outdated,
leading to an incorrect number of bytes being output by the mixer.

PiperOrigin-RevId: 619832502
---
 RELEASENOTES.md                               |  2 ++
 .../media3/transformer/DefaultAudioMixer.java |  3 +++
 .../transformer/DefaultAudioMixerTest.java    | 22 +++++++++++++++++++
 3 files changed, 27 insertions(+)

diff --git a/RELEASENOTES.md b/RELEASENOTES.md
index b8abb69049..f7206c6a8b 100644
--- a/RELEASENOTES.md
+++ b/RELEASENOTES.md
@@ -44,6 +44,8 @@
     *   Add support for changing between SDR and HDR input media in a sequence.
     *   Add support for composition-level audio effects.
     *   Add support for transcoding Ultra HDR images into HDR videos.
+    *   Fix issue where the `DefaultAudioMixer` doesn’t output the correct
+        amount of bytes after being reset and reused.
 *   Track Selection:
     *   `DefaultTrackSelector`: Prefer video tracks with a 'reasonable' frame
         rate (>=10fps) over those with a lower or unset frame rate. This ensures
diff --git a/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultAudioMixer.java b/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultAudioMixer.java
index 61ec69a03c..4dda143854 100644
--- a/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultAudioMixer.java
+++ b/libraries/transformer/src/main/java/androidx/media3/transformer/DefaultAudioMixer.java
@@ -84,6 +84,7 @@ public final class DefaultAudioMixer implements AudioMixer {
   // TODO(b/290002438, b/276734854): Improve buffer management & determine best default size.
   private static final int DEFAULT_BUFFER_SIZE_MS = 500;
 
+  private final boolean outputSilenceWithNoSources;
   private final boolean clipFloatOutput;
   private final SparseArray<SourceInfo> sources;
   private int nextSourceId;
@@ -108,6 +109,7 @@ public final class DefaultAudioMixer implements AudioMixer {
   private long maxPositionOfRemovedSources;
 
   private DefaultAudioMixer(boolean outputSilenceWithNoSources, boolean clipFloatOutput) {
+    this.outputSilenceWithNoSources = outputSilenceWithNoSources;
     this.clipFloatOutput = clipFloatOutput;
     sources = new SparseArray<>();
     outputAudioFormat = AudioFormat.NOT_SET;
@@ -313,6 +315,7 @@ public final class DefaultAudioMixer implements AudioMixer {
     inputLimit = C.LENGTH_UNSET;
     outputPosition = 0;
     endPosition = Long.MAX_VALUE;
+    maxPositionOfRemovedSources = outputSilenceWithNoSources ? Long.MAX_VALUE : 0;
   }
 
   private void checkStateIsConfigured() {
diff --git a/libraries/transformer/src/test/java/androidx/media3/transformer/DefaultAudioMixerTest.java b/libraries/transformer/src/test/java/androidx/media3/transformer/DefaultAudioMixerTest.java
index d10a27ff29..d247a13801 100644
--- a/libraries/transformer/src/test/java/androidx/media3/transformer/DefaultAudioMixerTest.java
+++ b/libraries/transformer/src/test/java/androidx/media3/transformer/DefaultAudioMixerTest.java
@@ -388,6 +388,28 @@ public final class DefaultAudioMixerTest {
     assertThat(mixer.isEnded()).isTrue();
   }
 
+  @Test
+  public void reset_afterRemoveSource_whenNotOutputSilenceWithNoSources_resetsSourceEndPosition()
+      throws Exception {
+    assumeFalse(outputSilenceWithNoSources);
+    mixer.configure(AUDIO_FORMAT_STEREO_PCM_FLOAT, /* bufferSizeMs= */ 3, /* startTimeUs= */ 0);
+    int sourceId = mixer.addSource(AUDIO_FORMAT_STEREO_PCM_FLOAT, /* startTimeUs= */ 0);
+    ByteBuffer sourceBuffer = createByteBuffer(new float[] {0.1f, 0.2f, 0.3f, 0.4f, 0.5f, 0.6f});
+    mixer.queueInput(sourceId, sourceBuffer);
+    assertThat(sourceBuffer.remaining()).isEqualTo(0);
+    mixer.removeSource(sourceId);
+
+    mixer.reset();
+    mixer.configure(AUDIO_FORMAT_STEREO_PCM_FLOAT, /* bufferSizeMs= */ 3, /* startTimeUs= */ 0);
+    sourceId = mixer.addSource(AUDIO_FORMAT_STEREO_PCM_FLOAT, /* startTimeUs= */ 0);
+    sourceBuffer = createByteBuffer(new float[] {0.1f, 0.2f});
+    mixer.queueInput(sourceId, sourceBuffer);
+    assertThat(sourceBuffer.remaining()).isEqualTo(0);
+    mixer.removeSource(sourceId);
+
+    assertThat(createFloatArray(mixer.getOutput())).isEqualTo(new float[] {0.1f, 0.2f});
+  }
+
   @Test
   public void input_whileIsEnded_isNotConsumed() throws Exception {
     mixer.configure(