From 4f57c46560006f188d6c8cf64aa3d7ffc003368c Mon Sep 17 00:00:00 2001
From: hoangtc <hoangtc@google.com>
Date: Thu, 28 Jun 2018 06:01:34 -0700
Subject: [PATCH] Do not discard TS PES payload readers on Seek for TSExtractor
 in non-HLS mode.

Currently, we are discarding all TS PES payload readers on seek. This is to fix
a bug that makes PES stream readers read packets inconsistently after seeking (first
time, the readers ignore all packets before PMT, second seek time after PMT is
read, the readers will read all packets before PMT).
This works fine if we don't support seeking to arbitrary position within the TS
stream. However, this may lead to a problem if we want to support seeking, because
some TS streams only have track format at the begining of the stream, and reseting
the payload readers after seeking can lead to track format not available anymore.

This change reverts the original fix (do not discard TS PES payload readers on seek
anymore). Instead, in non-HLS mode, after we have processed all PMTs packets in
the strea, we will re-seek to the beginning and perform readin again to make sure
the result will be consistent across multiple seeks. This is a prepare step to
support seeking in TS streams.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=202466233
---
 .../exoplayer2/extractor/ts/TsExtractor.java  | 35 +++++++++++++++----
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/TsExtractor.java b/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/TsExtractor.java
index 8ec4a4a962..3fdbd4707a 100644
--- a/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/TsExtractor.java
+++ b/library/core/src/main/java/com/google/android/exoplayer2/extractor/ts/TsExtractor.java
@@ -117,6 +117,7 @@ public final class TsExtractor implements Extractor {
   private final TsPayloadReader.Factory payloadReaderFactory;
   private final SparseArray<TsPayloadReader> tsPayloadReaders; // Indexed by pid
   private final SparseBooleanArray trackIds;
+  private final SparseBooleanArray trackPids;
 
   // Accessed only by the loading thread.
   private ExtractorOutput output;
@@ -167,6 +168,7 @@ public final class TsExtractor implements Extractor {
     }
     tsPacketBuffer = new ParsableByteArray(new byte[BUFFER_SIZE], 0);
     trackIds = new SparseBooleanArray();
+    trackPids = new SparseBooleanArray();
     tsPayloadReaders = new SparseArray<>();
     continuityCounters = new SparseIntArray();
     resetPayloadReaders();
@@ -203,14 +205,16 @@ public final class TsExtractor implements Extractor {
 
   @Override
   public void seek(long position, long timeUs) {
+    Assertions.checkState(mode != MODE_HLS);
     int timestampAdjustersCount = timestampAdjusters.size();
     for (int i = 0; i < timestampAdjustersCount; i++) {
       timestampAdjusters.get(i).reset();
     }
     tsPacketBuffer.reset();
     continuityCounters.clear();
-    // Elementary stream readers' state should be cleared to get consistent behaviours when seeking.
-    resetPayloadReaders();
+    for (int i = 0; i < tsPayloadReaders.size(); i++) {
+      tsPayloadReaders.valueAt(i).seek();
+    }
     bytesSinceLastSync = 0;
   }
 
@@ -274,10 +278,21 @@ public final class TsExtractor implements Extractor {
     }
 
     // Read the payload.
-    tsPacketBuffer.setLimit(endOfPacket);
-    payloadReader.consume(tsPacketBuffer, payloadUnitStartIndicator);
-    tsPacketBuffer.setLimit(limit);
+    boolean wereTracksEnded = tracksEnded;
+    if (shouldConsumePacketPayload(pid)) {
+      tsPacketBuffer.setLimit(endOfPacket);
+      payloadReader.consume(tsPacketBuffer, payloadUnitStartIndicator);
+      tsPacketBuffer.setLimit(limit);
+    }
 
+    if (mode != MODE_HLS && !wereTracksEnded && tracksEnded) {
+      // We have read all tracks from all PMTs in this stream. Now seek to the beginning and read
+      // again to make sure we output all media, including any contained in packets prior to those
+      // containing the track information.
+      seek(/* position= */ 0, /* timeUs= */ 0);
+      seekPosition.position = 0;
+      return RESULT_SEEK;
+    }
     tsPacketBuffer.setPosition(endOfPacket);
     return RESULT_CONTINUE;
   }
@@ -345,6 +360,12 @@ public final class TsExtractor implements Extractor {
     return position;
   }
 
+  private boolean shouldConsumePacketPayload(int packetPid) {
+    return mode == MODE_HLS
+        || tracksEnded
+        || !trackPids.get(packetPid, /* valueIfKeyNotFound= */ false); // It's a PSI packet
+  }
+
   private void resetPayloadReaders() {
     trackIds.clear();
     tsPayloadReaders.clear();
@@ -511,14 +532,16 @@ public final class TsExtractor implements Extractor {
       int trackIdCount = trackIdToPidScratch.size();
       for (int i = 0; i < trackIdCount; i++) {
         int trackId = trackIdToPidScratch.keyAt(i);
+        int trackPid = trackIdToPidScratch.valueAt(i);
         trackIds.put(trackId, true);
+        trackPids.put(trackPid, true);
         TsPayloadReader reader = trackIdToReaderScratch.valueAt(i);
         if (reader != null) {
           if (reader != id3Reader) {
             reader.init(timestampAdjuster, output,
                 new TrackIdGenerator(programNumber, trackId, MAX_PID_PLUS_ONE));
           }
-          tsPayloadReaders.put(trackIdToPidScratch.valueAt(i), reader);
+          tsPayloadReaders.put(trackPid, reader);
         }
       }