Use estimated bitrates in AdaptiveTrackSelection

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=220643861

library/core/src/main/java/com/google/android/exoplayer2/trackselection/AdaptiveTrackSelection.java

@@ -47,6 +47,8 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
     private final long minTimeBetweenBufferReevaluationMs;
     private final Clock clock;
 
+    private TrackBitrateEstimator trackBitrateEstimator;
+
     /** Creates an adaptive track selection factory with default parameters. */
     public Factory() {
       this(
@@ -199,6 +201,18 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
           bufferedFractionToLiveEdgeForQualityIncrease;
       this.minTimeBetweenBufferReevaluationMs = minTimeBetweenBufferReevaluationMs;
       this.clock = clock;
+      trackBitrateEstimator = TrackBitrateEstimator.DEFAULT;
+    }
+
+    /**
+     * Sets a TrackBitrateEstimator.
+     *
+     * <p>This method is experimental, and will be renamed or removed in a future release.
+     *
+     * @param trackBitrateEstimator A {@link TrackBitrateEstimator}.
+     */
+    public void experimental_setTrackBitrateEstimator(TrackBitrateEstimator trackBitrateEstimator) {
+      this.trackBitrateEstimator = trackBitrateEstimator;
     }
 
     @Override
@@ -207,17 +221,20 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
       if (this.bandwidthMeter != null) {
         bandwidthMeter = this.bandwidthMeter;
       }
-      return new AdaptiveTrackSelection(
+      AdaptiveTrackSelection adaptiveTrackSelection =
-          group,
+          new AdaptiveTrackSelection(
-          tracks,
+              group,
-          bandwidthMeter,
+              tracks,
-          minDurationForQualityIncreaseMs,
+              bandwidthMeter,
-          maxDurationForQualityDecreaseMs,
+              minDurationForQualityIncreaseMs,
-          minDurationToRetainAfterDiscardMs,
+              maxDurationForQualityDecreaseMs,
-          bandwidthFraction,
+              minDurationToRetainAfterDiscardMs,
-          bufferedFractionToLiveEdgeForQualityIncrease,
+              bandwidthFraction,
-          minTimeBetweenBufferReevaluationMs,
+              bufferedFractionToLiveEdgeForQualityIncrease,
-          clock);
+              minTimeBetweenBufferReevaluationMs,
+              clock);
+      adaptiveTrackSelection.experimental_setTrackBitrateEstimator(trackBitrateEstimator);
+      return adaptiveTrackSelection;
     }
   }
 
@@ -236,7 +253,11 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
   private final float bufferedFractionToLiveEdgeForQualityIncrease;
   private final long minTimeBetweenBufferReevaluationMs;
   private final Clock clock;
+  private final Format[] formats;
+  private final int[] formatBitrates;
+  private final int[] estimatedBitrates;
 
+  private TrackBitrateEstimator trackBitrateEstimator;
   private float playbackSpeed;
   private int selectedIndex;
   private int reason;
@@ -314,11 +335,32 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
     playbackSpeed = 1f;
     reason = C.SELECTION_REASON_INITIAL;
     lastBufferEvaluationMs = C.TIME_UNSET;
+    trackBitrateEstimator = TrackBitrateEstimator.DEFAULT;
+    formats = new Format[length];
+    formatBitrates = new int[length];
+    estimatedBitrates = new int[length];
+    for (int i = 0; i < length; i++) {
+      @SuppressWarnings("nullness:method.invocation.invalid")
+      Format format = getFormat(i);
+      formats[i] = format;
+      formatBitrates[i] = formats[i].bitrate;
+    }
     @SuppressWarnings("nullness:method.invocation.invalid")
-    int selectedIndex = determineIdealSelectedIndex(Long.MIN_VALUE);
+    int selectedIndex = determineIdealSelectedIndex(Long.MIN_VALUE, formatBitrates);
     this.selectedIndex = selectedIndex;
   }
 
+  /**
+   * Sets a TrackBitrateEstimator.
+   *
+   * <p>This method is experimental, and will be renamed or removed in a future release.
+   *
+   * @param trackBitrateEstimator A {@link TrackBitrateEstimator}.
+   */
+  public void experimental_setTrackBitrateEstimator(TrackBitrateEstimator trackBitrateEstimator) {
+    this.trackBitrateEstimator = trackBitrateEstimator;
+  }
+
   @Override
   public void enable() {
     lastBufferEvaluationMs = C.TIME_UNSET;
@@ -338,9 +380,11 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
       MediaChunkIterator[] mediaChunkIterators) {
     long nowMs = clock.elapsedRealtime();
 
+    trackBitrateEstimator.getBitrates(formats, queue, mediaChunkIterators, estimatedBitrates);
+
     // Stash the current selection, then make a new one.
     int currentSelectedIndex = selectedIndex;
-    selectedIndex = determineIdealSelectedIndex(nowMs);
+    selectedIndex = determineIdealSelectedIndex(nowMs, estimatedBitrates);
     if (selectedIndex == currentSelectedIndex) {
       return;
     }
@@ -402,7 +446,7 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
     if (playoutBufferedDurationBeforeLastChunkUs < minDurationToRetainAfterDiscardUs) {
       return queueSize;
     }
-    int idealSelectedIndex = determineIdealSelectedIndex(nowMs);
+    int idealSelectedIndex = determineIdealSelectedIndex(nowMs, formatBitrates);
     Format idealFormat = getFormat(idealSelectedIndex);
     // If the chunks contain video, discard from the first SD chunk beyond
     // minDurationToRetainAfterDiscardUs whose resolution and bitrate are both lower than the ideal
@@ -429,14 +473,14 @@ public class AdaptiveTrackSelection extends BaseTrackSelection {
    *
    * @param nowMs The current time in the timebase of {@link Clock#elapsedRealtime()}, or {@link
    *     Long#MIN_VALUE} to ignore blacklisting.
+   * @param bitrates Track bitrates.
    */
-  private int determineIdealSelectedIndex(long nowMs) {
+  private int determineIdealSelectedIndex(long nowMs, int[] bitrates) {
     long effectiveBitrate = (long) (bandwidthMeter.getBitrateEstimate() * bandwidthFraction);
     int lowestBitrateNonBlacklistedIndex = 0;
     for (int i = 0; i < length; i++) {
       if (nowMs == Long.MIN_VALUE || !isBlacklisted(i, nowMs)) {
-        Format format = getFormat(i);
+        if (Math.round(bitrates[i] * playbackSpeed) <= effectiveBitrate) {
-        if (Math.round(format.bitrate * playbackSpeed) <= effectiveBitrate) {
           return i;
         } else {
           lowestBitrateNonBlacklistedIndex = i;

library/core/src/test/java/com/google/android/exoplayer2/trackselection/AdaptiveTrackSelectionTest.java

@@ -16,6 +16,9 @@
 package com.google.android.exoplayer2.trackselection;
 
 import static com.google.common.truth.Truth.assertThat;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.argThat;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -32,11 +35,13 @@ import com.google.android.exoplayer2.testutil.FakeMediaChunk;
 import com.google.android.exoplayer2.upstream.BandwidthMeter;
 import com.google.android.exoplayer2.util.MimeTypes;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentMatcher;
 import org.mockito.Mock;
 import org.robolectric.RobolectricTestRunner;
 
@@ -220,6 +225,52 @@ public final class AdaptiveTrackSelectionTest {
     assertThat(adaptiveTrackSelection.getSelectionReason()).isEqualTo(C.SELECTION_REASON_ADAPTIVE);
   }
 
+  @Test
+  public void testUpdateSelectedTrackSwitchUpIfTrackBitrateEstimateIsLow() {
+    Format format1 = videoFormat(/* bitrate= */ 500, /* width= */ 320, /* height= */ 240);
+    Format format2 = videoFormat(/* bitrate= */ 1000, /* width= */ 640, /* height= */ 480);
+    Format format3 = videoFormat(/* bitrate= */ 2000, /* width= */ 960, /* height= */ 720);
+    TrackGroup trackGroup = new TrackGroup(format1, format2, format3);
+
+    // The second measurement onward returns 1500L, which isn't enough to switch up to format3 as
+    // the format bitrate is 2000.
+    when(mockBandwidthMeter.getBitrateEstimate()).thenReturn(1000L, 1500L);
+
+    // But TrackBitrateEstimator returns 1500 for 3rd track so it should switch up.
+    TrackBitrateEstimator estimator = mock(TrackBitrateEstimator.class);
+    when(estimator.getBitrates(any(), any(), any(), any())).thenReturn(new int[] {500, 1000, 1500});
+
+    adaptiveTrackSelection = adaptiveTrackSelection(trackGroup);
+    adaptiveTrackSelection.experimental_setTrackBitrateEstimator(estimator);
+
+    adaptiveTrackSelection.updateSelectedTrack(
+        /* playbackPositionUs= */ 0,
+        /* bufferedDurationUs= */ AdaptiveTrackSelection
+                .DEFAULT_MIN_DURATION_FOR_QUALITY_INCREASE_MS
+            * 1000,
+        /* availableDurationUs= */ C.TIME_UNSET,
+        /* queue= */ Collections.emptyList(),
+        /* mediaChunkIterators= */ THREE_EMPTY_MEDIA_CHUNK_ITERATORS);
+
+    ArgumentMatcher<Format[]> matcher =
+        new ArgumentMatcher<Format[]>() {
+          @Override
+          public boolean matches(Object argument) {
+            Format[] formats = (Format[]) argument;
+            return formats.length == 3
+                && Arrays.asList(formats).containsAll(Arrays.asList(format1, format2, format3));
+          }
+        };
+    verify(estimator)
+        .getBitrates(
+            argThat(matcher),
+            eq(Collections.emptyList()),
+            eq(THREE_EMPTY_MEDIA_CHUNK_ITERATORS),
+            any());
+    assertThat(adaptiveTrackSelection.getSelectedFormat()).isEqualTo(format3);
+    assertThat(adaptiveTrackSelection.getSelectionReason()).isEqualTo(C.SELECTION_REASON_ADAPTIVE);
+  }
+
   @Test
   public void testEvaluateQueueSizeReturnQueueSizeIfBandwidthIsNotImproved() {
     Format format1 = videoFormat(/* bitrate= */ 500, /* width= */ 320, /* height= */ 240);
@@ -322,17 +373,8 @@ public final class AdaptiveTrackSelectionTest {
   }
 
   private AdaptiveTrackSelection adaptiveTrackSelection(TrackGroup trackGroup) {
-    return new AdaptiveTrackSelection(
+    return adaptiveTrackSelectionWithMinDurationForQualityIncreaseMs(
-        trackGroup,
+        trackGroup, AdaptiveTrackSelection.DEFAULT_MIN_DURATION_FOR_QUALITY_INCREASE_MS);
-        selectedAllTracksInGroup(trackGroup),
-        mockBandwidthMeter,
-        AdaptiveTrackSelection.DEFAULT_MIN_DURATION_FOR_QUALITY_INCREASE_MS,
-        AdaptiveTrackSelection.DEFAULT_MAX_DURATION_FOR_QUALITY_DECREASE_MS,
-        AdaptiveTrackSelection.DEFAULT_MIN_DURATION_TO_RETAIN_AFTER_DISCARD_MS,
-        /* bandwidthFraction= */ 1.0f,
-        AdaptiveTrackSelection.DEFAULT_BUFFERED_FRACTION_TO_LIVE_EDGE_FOR_QUALITY_INCREASE,
-        AdaptiveTrackSelection.DEFAULT_MIN_TIME_BETWEEN_BUFFER_REEVALUTATION_MS,
-        fakeClock);
   }
 
   private AdaptiveTrackSelection adaptiveTrackSelectionWithMinDurationForQualityIncreaseMs(