Send empty pusi flagged data, on end of input, from TS extractor to readers; add end of input flag to readers to flush accumulated data from last PES

libraries/extractor/src/main/java/androidx/media3/extractor/ts/Ac3Reader.java

@@ -157,7 +157,7 @@ public final class Ac3Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/Ac4Reader.java

@@ -159,7 +159,7 @@ public final class Ac4Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/AdtsReader.java

@@ -192,7 +192,7 @@ public final class AdtsReader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/DtsReader.java

@@ -131,7 +131,7 @@ public final class DtsReader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/DvbSubtitleReader.java

@@ -87,7 +87,7 @@ public final class DvbSubtitleReader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     if (writingSample) {
       if (sampleTimeUs != C.TIME_UNSET) {
         for (TrackOutput output : outputs) {

libraries/extractor/src/main/java/androidx/media3/extractor/ts/ElementaryStreamReader.java

@@ -54,5 +54,5 @@ public interface ElementaryStreamReader {
   void consume(ParsableByteArray data) throws ParserException;
 
   /** Called when a packet ends. */
-  void packetFinished();
+  void packetFinished(boolean isEndOfInput);
 }

libraries/extractor/src/main/java/androidx/media3/extractor/ts/H262Reader.java

@@ -217,7 +217,7 @@ public final class H262Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/H263Reader.java

@@ -217,7 +217,7 @@ public final class H263Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/H264Reader.java

@@ -168,8 +168,10 @@ public final class H264Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
+    if (isEndOfInput) {
+      sampleReader.end(totalBytesWritten);
+    }
   }
 
   @RequiresNonNull("sampleReader")
@@ -500,6 +502,12 @@ public final class H264Reader implements ElementaryStreamReader {
       output.sampleMetadata(sampleTimeUs, flags, size, offset, null);
     }
 
+    public void end(long position) {
+      // Output a final sample with the nal units currently held
+      nalUnitStartPosition = position;
+      outputSample(0);
+    }
+
     private static final class SliceHeaderData {
 
       private static final int SLICE_TYPE_I = 2;

libraries/extractor/src/main/java/androidx/media3/extractor/ts/H265Reader.java

@@ -173,8 +173,10 @@ public final class H265Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
+    if (isEndOfInput) {
+      sampleReader.end(totalBytesWritten);
+    }
   }
 
   @RequiresNonNull("sampleReader")
@@ -376,6 +378,12 @@ public final class H265Reader implements ElementaryStreamReader {
       output.sampleMetadata(sampleTimeUs, flags, size, offset, null);
     }
 
+    public void end(long position) {
+      // Output a final sample with the nal units currently held
+      nalUnitPosition = position;
+      outputSample(0);
+    }
+
     /** Returns whether a NAL unit type is one that occurs before any VCL NAL units in a sample. */
     private static boolean isPrefixNalUnit(int nalUnitType) {
       return (VPS_NUT <= nalUnitType && nalUnitType <= AUD_NUT) || nalUnitType == PREFIX_SEI_NUT;

libraries/extractor/src/main/java/androidx/media3/extractor/ts/Id3Reader.java

@@ -121,7 +121,7 @@ public final class Id3Reader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     Assertions.checkStateNotNull(output); // Asserts that createTracks has been called.
     if (!writingSample || sampleSize == 0 || sampleBytesRead != sampleSize) {
       return;

libraries/extractor/src/main/java/androidx/media3/extractor/ts/LatmReader.java

@@ -151,7 +151,7 @@ public final class LatmReader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/MpegAudioReader.java

@@ -118,7 +118,7 @@ public final class MpegAudioReader implements ElementaryStreamReader {
   }
 
   @Override
-  public void packetFinished() {
+  public void packetFinished(boolean isEndOfInput) {
     // Do nothing.
   }
 

libraries/extractor/src/main/java/androidx/media3/extractor/ts/PesReader.java

@@ -107,7 +107,8 @@ public final class PesReader implements TsPayloadReader {
             Log.w(TAG, "Unexpected start indicator: expected " + payloadSize + " more bytes");
           }
           // Either way, notify the reader that it has now finished.
-          reader.packetFinished();
+          boolean isEndOfInput = (data.limit() == 0);
+          reader.packetFinished(isEndOfInput);
           break;
         default:
           throw new IllegalStateException();
@@ -147,7 +148,8 @@ public final class PesReader implements TsPayloadReader {
           if (payloadSize != C.LENGTH_UNSET) {
             payloadSize -= readLength;
             if (payloadSize == 0) {
-              reader.packetFinished();
+              // There are bytes left in data, see above, so this is not the end of input
+              reader.packetFinished(false);
               setState(STATE_READING_HEADER);
             }
           }

libraries/extractor/src/main/java/androidx/media3/extractor/ts/PsExtractor.java

@@ -357,7 +357,7 @@ public final class PsExtractor implements Extractor {
       pesPayloadReader.packetStarted(timeUs, TsPayloadReader.FLAG_DATA_ALIGNMENT_INDICATOR);
       pesPayloadReader.consume(data);
       // We always have complete PES packets with program stream.
-      pesPayloadReader.packetFinished();
+      pesPayloadReader.packetFinished(false);
     }
 
     private void parseHeader() {

libraries/extractor/src/main/java/androidx/media3/extractor/ts/TsExtractor.java

@@ -323,7 +323,14 @@ public final class TsExtractor implements Extractor {
     }
 
     if (!fillBufferWithAtLeastOnePacket(input)) {
-      return RESULT_END_OF_INPUT;
+      if (mode != MODE_HLS) {
+        // Send a dummy pusi to allow for packetFinished to be triggered on the last unit
+        for (int i = 0; i < tsPayloadReaders.size(); i++) {
+          TsPayloadReader payloadReader = tsPayloadReaders.valueAt(i);
+          int pid = tsPayloadReaders.keyAt(i);
+          payloadReader.consume(new ParsableByteArray(), FLAG_PAYLOAD_UNIT_START_INDICATOR);
+        }
+      }      return RESULT_END_OF_INPUT;
     }
 
     int endOfPacket = findEndOfFirstTsPacketInBuffer();

libraries/extractor/src/test/java/androidx/media3/extractor/ts/TsExtractorTest.java

@@ -287,7 +287,7 @@ public final class TsExtractorTest {
     public void consume(ParsableByteArray data) {}
 
     @Override
-    public void packetFinished() {
+    public void packetFinished(boolean isEndOfInput) {
       packetsRead++;
     }