Keep selections for preloaded adaptive sample streams

When an adaptive source has been preloaded, the track selection
of the player should possibly not override the preloaded selection
to avoid discarding preloaded data.

PiperOrigin-RevId: 591256334

RELEASENOTES.md

@@ -40,6 +40,7 @@
         negative values for `bufferedDurationUs` from chunk sources, resulting
         in an `IllegalArgumentException`
         ([#888](https://github.com/androidx/media/issues/888)).
+    *   Support adaptive media sources with `PreloadMediaSource`.
 *   Transformer:
     *   Add support for flattening H.265/HEVC SEF slow motion videos.
     *   Increase transmuxing speed, especially for 'remove video' edits.

libraries/exoplayer/src/main/java/androidx/media3/exoplayer/source/preload/PreloadMediaPeriod.java

@@ -16,8 +16,10 @@
 package androidx.media3.exoplayer.source.preload;
 
 import static androidx.media3.common.util.Assertions.checkNotNull;
+import static androidx.media3.common.util.Assertions.checkState;
 
 import androidx.annotation.Nullable;
+import androidx.media3.common.C;
 import androidx.media3.common.util.NullableType;
 import androidx.media3.exoplayer.LoadingInfo;
 import androidx.media3.exoplayer.SeekParameters;
@@ -27,7 +29,7 @@ import androidx.media3.exoplayer.source.TrackGroupArray;
 import androidx.media3.exoplayer.trackselection.ExoTrackSelection;
 import androidx.media3.exoplayer.trackselection.TrackSelectorResult;
 import java.io.IOException;
-import java.util.Arrays;
+import java.util.Objects;
 
 /** A {@link MediaPeriod} that has data preloaded before playback. */
 /* package */ final class PreloadMediaPeriod implements MediaPeriod {
@@ -105,32 +107,108 @@ import java.util.Arrays;
       @NullableType SampleStream[] streams,
       boolean[] streamResetFlags,
       long positionUs) {
-    long trackSelectionPositionUs;
-    if (preloadTrackSelectionHolder != null
-        && Arrays.equals(selections, preloadTrackSelectionHolder.trackSelectorResult.selections)
-        && positionUs == preloadTrackSelectionHolder.trackSelectionPositionUs) {
-      trackSelectionPositionUs = preloadTrackSelectionHolder.trackSelectionPositionUs;
-      System.arraycopy(
-          preloadTrackSelectionHolder.streams,
-          0,
-          streams,
-          0,
-          preloadTrackSelectionHolder.streams.length);
-      System.arraycopy(
-          preloadTrackSelectionHolder.streamResetFlags,
-          0,
-          streamResetFlags,
-          0,
-          preloadTrackSelectionHolder.streamResetFlags.length);
-    } else {
+    if (preloadTrackSelectionHolder == null) {
+      // No preload track selection was done.
+      return mediaPeriod.selectTracks(
+          selections, mayRetainStreamFlags, streams, streamResetFlags, positionUs);
+    }
+    checkState(streams.length == preloadTrackSelectionHolder.streams.length);
+    if (positionUs != preloadTrackSelectionHolder.trackSelectionPositionUs) {
+      // Position changed. Copy formerly preloaded sample streams to the track selection properties
+      // to make sure we give the period the chance to release discarded sample streams.
+      for (int i = 0; i < preloadTrackSelectionHolder.streams.length; i++) {
+        if (preloadTrackSelectionHolder.streams[i] != null) {
+          streams[i] = preloadTrackSelectionHolder.streams[i];
+          mayRetainStreamFlags[i] = false;
+        }
+      }
+      preloadTrackSelectionHolder = null;
+      return mediaPeriod.selectTracks(
+          selections, mayRetainStreamFlags, streams, streamResetFlags, positionUs);
+    }
+    PreloadTrackSelectionHolder holder = checkNotNull(preloadTrackSelectionHolder);
+    long trackSelectionPositionUs = holder.trackSelectionPositionUs;
+    boolean[] preloadStreamResetFlags = holder.streamResetFlags;
+    if (maybeUpdatePreloadTrackSelectionHolderForReselection(selections, holder)) {
+      // Preload can only be partially reused. Select tracks again attempting to retain preloads.
+      preloadStreamResetFlags = new boolean[preloadStreamResetFlags.length];
       trackSelectionPositionUs =
           mediaPeriod.selectTracks(
-              selections, mayRetainStreamFlags, streams, streamResetFlags, positionUs);
+              holder.trackSelectorResult.selections,
+              holder.mayRetainStreamFlags,
+              holder.streams,
+              preloadStreamResetFlags,
+              holder.trackSelectionPositionUs);
+      // Make sure to set the reset flags for streams we preloaded.
+      for (int i = 0; i < holder.mayRetainStreamFlags.length; i++) {
+        if (holder.mayRetainStreamFlags[i]) {
+          preloadStreamResetFlags[i] = true;
+        }
+      }
     }
+    System.arraycopy(holder.streams, 0, streams, 0, holder.streams.length);
+    System.arraycopy(
+        preloadStreamResetFlags, 0, streamResetFlags, 0, preloadStreamResetFlags.length);
     preloadTrackSelectionHolder = null;
     return trackSelectionPositionUs;
   }
 
+  private static boolean maybeUpdatePreloadTrackSelectionHolderForReselection(
+      @NullableType ExoTrackSelection[] selections,
+      PreloadTrackSelectionHolder preloadTrackSelectionHolder) {
+    @NullableType
+    ExoTrackSelection[] preloadSelections =
+        checkNotNull(preloadTrackSelectionHolder).trackSelectorResult.selections;
+    boolean needsReselection = false;
+    for (int i = 0; i < selections.length; i++) {
+      ExoTrackSelection selection = selections[i];
+      ExoTrackSelection preloadSelection = preloadSelections[i];
+      if (selection == null && preloadSelection == null) {
+        continue;
+      }
+      preloadTrackSelectionHolder.mayRetainStreamFlags[i] = false;
+      if (selection == null) {
+        // Preloaded track got disabled. Discard preloaded stream.
+        preloadTrackSelectionHolder.trackSelectorResult.selections[i] = null;
+        needsReselection = true;
+      } else if (preloadSelection == null) {
+        // Enabled track not preloaded. Update selection.
+        preloadTrackSelectionHolder.trackSelectorResult.selections[i] = selection;
+        needsReselection = true;
+      } else if (!isSameAdaptionSet(selection, preloadSelection)) {
+        // Adaption set has changed. Discard preloaded stream.
+        preloadTrackSelectionHolder.trackSelectorResult.selections[i] = selection;
+        needsReselection = true;
+      } else if (selection.getTrackGroup().type == C.TRACK_TYPE_VIDEO
+          || selection.getTrackGroup().type == C.TRACK_TYPE_AUDIO
+          || selection.getSelectedIndexInTrackGroup()
+              == preloadSelection.getSelectedIndexInTrackGroup()) {
+        // The selection in a audio or video track has changed or it hasn't changed. Set the retain
+        // flag in case we reselect with a partially preloaded streams set.
+        preloadTrackSelectionHolder.mayRetainStreamFlags[i] = true;
+      } else {
+        // The selection in a non-audio or non-video track has changed. Discard preloaded stream.
+        preloadTrackSelectionHolder.trackSelectorResult.selections[i] = selection;
+        needsReselection = true;
+      }
+    }
+    return needsReselection;
+  }
+
+  private static boolean isSameAdaptionSet(
+      ExoTrackSelection selection, ExoTrackSelection preloadSelection) {
+    if (!Objects.equals(selection.getTrackGroup(), preloadSelection.getTrackGroup())
+        || selection.length() != preloadSelection.length()) {
+      return false;
+    }
+    for (int i = 0; i < selection.length(); i++) {
+      if (selection.getIndexInTrackGroup(i) != preloadSelection.getIndexInTrackGroup(i)) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   /* package */ long selectTracksForPreloading(
       TrackSelectorResult trackSelectorResult, long positionUs) {
     @NullableType ExoTrackSelection[] selections = trackSelectorResult.selections;
@@ -154,6 +232,7 @@ import java.util.Arrays;
     preloadTrackSelectionHolder =
         new PreloadTrackSelectionHolder(
             trackSelectorResult,
+            mayRetainStreamFlags,
             preloadedSampleStreams,
             preloadedStreamResetFlags,
             trackSelectionPositionUs);
@@ -207,16 +286,19 @@ import java.util.Arrays;
 
   private static class PreloadTrackSelectionHolder {
     public final TrackSelectorResult trackSelectorResult;
+    public final boolean[] mayRetainStreamFlags;
     public final @NullableType SampleStream[] streams;
     public final boolean[] streamResetFlags;
     public final long trackSelectionPositionUs;
 
     public PreloadTrackSelectionHolder(
         TrackSelectorResult trackSelectorResult,
+        boolean[] mayRetainStreamFlags,
         @NullableType SampleStream[] streams,
         boolean[] streamResetFlags,
         long trackSelectionPositionUs) {
       this.trackSelectorResult = trackSelectorResult;
+      this.mayRetainStreamFlags = mayRetainStreamFlags;
       this.streams = streams;
       this.streamResetFlags = streamResetFlags;
       this.trackSelectionPositionUs = trackSelectionPositionUs;

libraries/test_utils/src/main/java/androidx/media3/test/utils/FakeTrackSelection.java

@@ -25,6 +25,7 @@ import androidx.media3.common.util.UnstableApi;
 import androidx.media3.exoplayer.source.chunk.MediaChunk;
 import androidx.media3.exoplayer.source.chunk.MediaChunkIterator;
 import androidx.media3.exoplayer.trackselection.ExoTrackSelection;
+import com.google.common.base.Objects;
 import java.util.List;
 
 /**
@@ -35,13 +36,19 @@ import java.util.List;
 public final class FakeTrackSelection implements ExoTrackSelection {
 
   private final TrackGroup rendererTrackGroup;
+  private final int selectedIndex;
 
   public int enableCount;
   public int releaseCount;
   public boolean isEnabled;
 
   public FakeTrackSelection(TrackGroup rendererTrackGroup) {
+    this(rendererTrackGroup, /* selectedIndex= */ 0);
+  }
+
+  public FakeTrackSelection(TrackGroup rendererTrackGroup, int selectedIndex) {
     this.rendererTrackGroup = rendererTrackGroup;
+    this.selectedIndex = selectedIndex;
   }
 
   // TrackSelection implementation.
@@ -68,18 +75,23 @@ public final class FakeTrackSelection implements ExoTrackSelection {
 
   @Override
   public int getIndexInTrackGroup(int index) {
-    return 0;
+    return index;
   }
 
   @Override
   public int indexOf(Format format) {
     assertThat(isEnabled).isTrue();
-    return 0;
+    for (int i = 0; i < rendererTrackGroup.length; i++) {
+      if (rendererTrackGroup.getFormat(i).equals(format)) {
+        return i;
+      }
+    }
+    return -1;
   }
 
   @Override
   public int indexOf(int indexInTrackGroup) {
-    return 0;
+    return indexInTrackGroup;
   }
 
   // ExoTrackSelection specific methods.
@@ -102,17 +114,17 @@ public final class FakeTrackSelection implements ExoTrackSelection {
 
   @Override
   public Format getSelectedFormat() {
-    return rendererTrackGroup.getFormat(0);
+    return rendererTrackGroup.getFormat(selectedIndex);
   }
 
   @Override
   public int getSelectedIndexInTrackGroup() {
-    return 0;
+    return selectedIndex;
   }
 
   @Override
   public int getSelectedIndex() {
-    return 0;
+    return selectedIndex;
   }
 
   @Override
@@ -158,4 +170,26 @@ public final class FakeTrackSelection implements ExoTrackSelection {
     assertThat(isEnabled).isTrue();
     return false;
   }
+
+  @Override
+  public boolean equals(@Nullable Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (!(o instanceof FakeTrackSelection)) {
+      return false;
+    }
+    FakeTrackSelection that = (FakeTrackSelection) o;
+    return enableCount == that.enableCount
+        && releaseCount == that.releaseCount
+        && isEnabled == that.isEnabled
+        && selectedIndex == that.selectedIndex
+        && Objects.equal(rendererTrackGroup, that.rendererTrackGroup);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(
+        rendererTrackGroup, enableCount, releaseCount, isEnabled, selectedIndex);
+  }
 }