Fix HlsTrackStreamWrapper chunk queue updating

Also fix possible repeat-preparation for HlsTrackStreamWrapper
and ExtractorSampleSource.
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=126640380

library/src/main/java/com/google/android/exoplayer/extractor/DefaultTrackOutput.java

@@ -111,6 +111,15 @@ public final class DefaultTrackOutput implements TrackOutput {
     }
   }
 
+  /**
+   * Sets a source identifier for subsequent samples.
+   *
+   * @param sourceId The source identifier.
+   */
+  public void sourceId(int sourceId) {
+    infoQueue.sourceId(sourceId);
+  }
+
   /**
    * Indicates that samples subsequently queued to the buffer should be spliced into those already
    * queued.
@@ -187,6 +196,16 @@ public final class DefaultTrackOutput implements TrackOutput {
     return infoQueue.getReadIndex();
   }
 
+  /**
+   * Peeks the source id of the next sample, or the current upstream source id if the buffer is
+   * empty.
+   *
+   * @return The source id.
+   */
+  public int peekSourceId() {
+    return infoQueue.peekSourceId();
+  }
+
   /**
    * Returns the upstream {@link Format} in which samples are being queued.
    */
@@ -581,6 +600,7 @@ public final class DefaultTrackOutput implements TrackOutput {
 
     private int capacity;
 
+    private int[] sourceIds;
     private long[] offsets;
     private int[] sizes;
     private int[] flags;
@@ -596,9 +616,11 @@ public final class DefaultTrackOutput implements TrackOutput {
     private long largestDequeuedTimestampUs;
     private long largestQueuedTimestampUs;
     private Format upstreamFormat;
+    private int upstreamSourceId;
 
     public InfoQueue() {
       capacity = SAMPLE_CAPACITY_INCREMENT;
+      sourceIds = new int[capacity];
       offsets = new long[capacity];
       timesUs = new long[capacity];
       flags = new int[capacity];
@@ -664,6 +686,10 @@ public final class DefaultTrackOutput implements TrackOutput {
       return offsets[relativeWriteIndex];
     }
 
+    public void sourceId(int sourceId) {
+      upstreamSourceId = sourceId;
+    }
+
     // Called by the consuming thread.
 
     /**
@@ -673,6 +699,14 @@ public final class DefaultTrackOutput implements TrackOutput {
       return absoluteReadIndex;
     }
 
+    /**
+     * Peeks the source id of the next sample, or the current upstream source id if the queue is
+     * empty.
+     */
+    public int peekSourceId() {
+      return queueSize == 0 ? upstreamSourceId : sourceIds[relativeReadIndex];
+    }
+
     /**
      * Returns whether the queue is empty.
      */
@@ -815,11 +849,13 @@ public final class DefaultTrackOutput implements TrackOutput {
       flags[relativeWriteIndex] = sampleFlags;
       encryptionKeys[relativeWriteIndex] = encryptionKey;
       formats[relativeWriteIndex] = upstreamFormat;
+      sourceIds[relativeWriteIndex] = upstreamSourceId;
       // Increment the write index.
       queueSize++;
       if (queueSize == capacity) {
         // Increase the capacity.
         int newCapacity = capacity + SAMPLE_CAPACITY_INCREMENT;
+        int[] newSourceIds = new int[newCapacity];
         long[] newOffsets = new long[newCapacity];
         long[] newTimesUs = new long[newCapacity];
         int[] newFlags = new int[newCapacity];
@@ -833,6 +869,7 @@ public final class DefaultTrackOutput implements TrackOutput {
         System.arraycopy(sizes, relativeReadIndex, newSizes, 0, beforeWrap);
         System.arraycopy(encryptionKeys, relativeReadIndex, newEncryptionKeys, 0, beforeWrap);
         System.arraycopy(formats, relativeReadIndex, newFormats, 0, beforeWrap);
+        System.arraycopy(sourceIds, relativeReadIndex, newSourceIds, 0, beforeWrap);
         int afterWrap = relativeReadIndex;
         System.arraycopy(offsets, 0, newOffsets, beforeWrap, afterWrap);
         System.arraycopy(timesUs, 0, newTimesUs, beforeWrap, afterWrap);
@@ -840,12 +877,14 @@ public final class DefaultTrackOutput implements TrackOutput {
         System.arraycopy(sizes, 0, newSizes, beforeWrap, afterWrap);
         System.arraycopy(encryptionKeys, 0, newEncryptionKeys, beforeWrap, afterWrap);
         System.arraycopy(formats, 0, newFormats, beforeWrap, afterWrap);
+        System.arraycopy(sourceIds, 0, newSourceIds, beforeWrap, afterWrap);
         offsets = newOffsets;
         timesUs = newTimesUs;
         flags = newFlags;
         sizes = newSizes;
         encryptionKeys = newEncryptionKeys;
         formats = newFormats;
+        sourceIds = newSourceIds;
         relativeReadIndex = 0;
         relativeWriteIndex = capacity;
         queueSize = capacity;

library/src/main/java/com/google/android/exoplayer/extractor/ExtractorSampleSource.java

@@ -544,7 +544,7 @@ public final class ExtractorSampleSource implements SampleSource, ExtractorOutpu
   // Internal methods.
 
   private void maybeFinishPrepare() {
-    if (seekMap == null || !tracksBuilt) {
+    if (prepared || seekMap == null || !tracksBuilt) {
       return;
     }
     for (DefaultTrackOutput sampleQueue : sampleQueues) {

library/src/main/java/com/google/android/exoplayer/hls/HlsMediaChunk.java

@@ -25,12 +25,20 @@ import com.google.android.exoplayer.upstream.DataSpec;
 import com.google.android.exoplayer.util.Util;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * An HLS {@link MediaChunk}.
  */
 /* package */ final class HlsMediaChunk extends MediaChunk {
 
+  private static final AtomicInteger UID_SOURCE = new AtomicInteger();
+
+  /**
+   * A unique identifier for the chunk.
+   */
+  public final int uid;
+
   /**
    * The discontinuity sequence number of the chunk.
    */
@@ -80,6 +88,7 @@ import java.io.IOException;
     this.shouldSpliceIn = shouldSpliceIn;
     // Note: this.dataSource and dataSource may be different.
     this.isEncrypted = this.dataSource instanceof Aes128DataSource;
+    uid = UID_SOURCE.getAndIncrement();
   }
 
   /**
@@ -89,9 +98,7 @@ import java.io.IOException;
    * @param output The output that will receive the loaded samples.
    */
   public void init(HlsTrackStreamWrapper output) {
-    if (shouldSpliceIn) {
+    output.init(uid, shouldSpliceIn);
-      output.splice();
-    }
     if (extractorNeedsInit) {
       extractor.init(output);
     }

library/src/main/java/com/google/android/exoplayer/hls/HlsTrackStreamWrapper.java

@@ -84,6 +84,7 @@ import java.util.List;
   private boolean prepared;
   private int enabledTrackCount;
   private Format downstreamFormat;
+  private int upstreamChunkUid;
 
   // Tracks are complicated in HLS. See documentation of buildTracks for details.
   // Indexed by track (as exposed by this source).
@@ -254,10 +255,9 @@ import java.util.List;
       return TrackStream.NOTHING_READ;
     }
 
-    // TODO[REFACTOR]: Restore this.
+    while (mediaChunks.size() > 1 && finishedReadingChunk(mediaChunks.getFirst())) {
-    // while (mediaChunks.size() > 1 && mediaChunks.get(1).startTimeUs <= downstreamPositionUs) {
+      mediaChunks.removeFirst();
-    //   mediaChunks.removeFirst();
+    }
-    // }
     HlsMediaChunk currentChunk = mediaChunks.getFirst();
     Format format = currentChunk.format;
     if (!format.equals(downstreamFormat)) {
@@ -271,6 +271,16 @@ import java.util.List;
         lastSeekPositionUs);
   }
 
+  private boolean finishedReadingChunk(HlsMediaChunk chunk) {
+    int chunkUid = chunk.uid;
+    for (int i = 0; i < sampleQueues.size(); i++) {
+      if (groupEnabledStates[i] && sampleQueues.valueAt(i).peekSourceId() == chunkUid) {
+        return false;
+      }
+    }
+    return true;
+  }
+
   // SequenceableLoader implementation
 
   @Override
@@ -379,11 +389,21 @@ import java.util.List;
   // Called by the consuming thread, but only when there is no loading thread.
 
   /**
-   * Indicates to all track outputs that they should splice in subsequently queued samples.
+   * Initializes the wrapper for loading a chunk.
+   *
+   * @param chunkUid The chunk's uid.
+   * @param shouldSpliceIn Whether the samples parsed from the chunk should be spliced into any
+   *     samples already queued to the wrapper.
    */
-  public void splice() {
+  public void init(int chunkUid, boolean shouldSpliceIn) {
+    upstreamChunkUid = chunkUid;
     for (int i = 0; i < sampleQueues.size(); i++) {
-      sampleQueues.valueAt(i).splice();
+      sampleQueues.valueAt(i).sourceId(chunkUid);
+    }
+    if (shouldSpliceIn) {
+      for (int i = 0; i < sampleQueues.size(); i++) {
+        sampleQueues.valueAt(i).splice();
+      }
     }
   }
 
@@ -396,6 +416,7 @@ import java.util.List;
     }
     DefaultTrackOutput trackOutput = new DefaultTrackOutput(allocator);
     trackOutput.setUpstreamFormatChangeListener(this);
+    trackOutput.sourceId(upstreamChunkUid);
     sampleQueues.put(id, trackOutput);
     return trackOutput;
   }
@@ -421,7 +442,7 @@ import java.util.List;
   // Internal methods.
 
   private void maybeFinishPrepare() {
-    if (!sampleQueuesBuilt) {
+    if (prepared || !sampleQueuesBuilt) {
       return;
     }
     int sampleQueueCount = sampleQueues.size();