Fix issue where end of stream wasn't signaled if last pes had length field set

RELEASENOTES.md

@@ -80,6 +80,9 @@
     *   MP4: Prevent infinite loops and out-of-bounds reads when parsing empty
         `ilst` metadata tag items
         ([#3191](https://github.com/androidx/media/pull/3191)).
+    *   MPEG-TS: Ensure the last frame is rendered for streams where the last
+        PES packet has a known length
+        ([#3206](https://github.com/androidx/media/pull/3206)).
 *   Inspector:
 *   Audio:
     *   Update `MediaCodecAudioRenderer` to extract the spatial channelMask from

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

@@ -167,11 +167,6 @@ public void consume(ParsableByteArray data) {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Continues a read from the provided {@code source} into a given {@code target}. It's assumed
    * that the data should be written into {@code target} starting from an offset of zero.

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

@@ -169,11 +169,6 @@ public void consume(ParsableByteArray data) {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Continues a read from the provided {@code source} into a given {@code target}. It's assumed
    * that the data should be written into {@code target} starting from an offset of zero.

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

@@ -202,11 +202,6 @@ public void consume(ParsableByteArray data) throws ParserException {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Returns the duration in microseconds per sample, or {@link C#TIME_UNSET} if the sample duration
    * is not available.

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

@@ -293,7 +293,7 @@ public void consume(ParsableByteArray data) throws ParserException {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void packetFinished() {
     if (state == STATE_CHECKING_FOR_EXTSS_AFTER_CORE) {
       checkNotNull(output);
       if (coreFormatPendingEmit) {

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

@@ -90,7 +90,7 @@ public void packetStarted(long pesTimeUs, @TsPayloadReader.Flags int flags) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void packetFinished() {
     if (writingSample) {
       // packetStarted method must be called before reading sample.
       checkState(sampleTimeUs != C.TIME_UNSET);

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

@@ -31,8 +31,9 @@
  *   <li>{@link #seek()} (optional, to reset the state)
  *   <li>{@link #packetStarted(long, int)} (to signal the start of a new packet)
  *   <li>{@link #consume(ParsableByteArray)} (zero or more times, to provide packet data)
- *   <li>{@link #packetFinished(boolean)} (to signal the end of the current packet)
+ *   <li>{@link #packetFinished()} (to signal the end of the current PES packet)
  *   <li>Repeat steps 3-5 for subsequent packets
+ *   <li>{@link #endOfInputReached()} (called when there are no more packets)
  * </ol>
  */
 @UnstableApi
@@ -66,6 +67,18 @@ public interface ElementaryStreamReader {
    */
   void consume(ParsableByteArray data) throws ParserException;
 
-  /** Called when a packet ends. */
-  void packetFinished(boolean isEndOfInput);
+  /**
+   * Called when a PES packet ends. For PES packets with an unset length, this is called when the
+   * next packet starts or the stream ends.
+   */
+  default void packetFinished() {
+    // Do nothing.
+  }
+
+  /**
+   * Called when there will be no further packets (unless a seek occurs) because the stream ended.
+   */
+  default void endOfInputReached() {
+    // Do nothing.
+  }
 }

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

@@ -220,10 +220,10 @@ public void consume(ParsableByteArray data) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void endOfInputReached() {
     // Asserts that createTracks has been called.
     checkNotNull(output);
-    if (isEndOfInput) {
+    if (sampleTimeUs != C.TIME_UNSET) {
       @C.BufferFlags int flags = sampleIsKeyframe ? C.BUFFER_FLAG_KEY_FRAME : 0;
       int size = (int) (totalBytesWritten - samplePosition);
       output.sampleMetadata(sampleTimeUs, flags, size, /* offset= */ 0, /* cryptoData= */ null);

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

@@ -221,13 +221,11 @@ public void consume(ParsableByteArray data) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void endOfInputReached() {
     // Assert that createTracks has been called.
     checkNotNull(sampleReader);
-    if (isEndOfInput) {
-      sampleReader.onDataEnd(totalBytesWritten, /* bytesWrittenPastPosition= */ 0, hasOutputFormat);
-      sampleReader.reset();
-    }
+    sampleReader.onDataEnd(totalBytesWritten, /* bytesWrittenPastPosition= */ 0, hasOutputFormat);
+    sampleReader.reset();
   }
 
   /**

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

@@ -183,15 +183,13 @@ public void consume(ParsableByteArray data) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void endOfInputReached() {
     assertTracksCreated();
-    if (isEndOfInput) {
-      seiReader.flush();
-      // Simulate end of current NAL unit and start an AUD one to trigger output of current sample
-      endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
-      startNalUnit(totalBytesWritten, NalUnitUtil.H264_NAL_UNIT_TYPE_AUD, pesTimeUs);
-      endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
-    }
+    seiReader.flush();
+    // Simulate end of current NAL unit and start an AUD one to trigger output of current sample
+    endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
+    startNalUnit(totalBytesWritten, NalUnitUtil.H264_NAL_UNIT_TYPE_AUD, pesTimeUs);
+    endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
   }
 
   @RequiresNonNull("sampleReader")

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

@@ -173,15 +173,13 @@ public void consume(ParsableByteArray data) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void endOfInputReached() {
     assertTracksCreated();
-    if (isEndOfInput) {
-      seiReader.flush();
-      // Simulate end of current NAL unit and start an unspecified one to trigger output of current
-      // sample
-      endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
-      startNalUnit(totalBytesWritten, 0, NalUnitUtil.H265_NAL_UNIT_TYPE_UNSPECIFIED, pesTimeUs);
-    }
+    seiReader.flush();
+    // Simulate end of current NAL unit and start an unspecified one to trigger output of current
+    // sample
+    endNalUnit(totalBytesWritten, 0, 0, pesTimeUs);
+    startNalUnit(totalBytesWritten, 0, NalUnitUtil.H265_NAL_UNIT_TYPE_UNSPECIFIED, pesTimeUs);
   }
 
   @RequiresNonNull("sampleReader")

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

@@ -124,7 +124,7 @@ public void consume(ParsableByteArray data) {
   }
 
   @Override
-  public void packetFinished(boolean isEndOfInput) {
+  public void packetFinished() {
     // Asserts that createTracks has been called.
     checkNotNull(output);
     if (!writingSample || sampleSize == 0 || sampleBytesRead != sampleSize) {

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

@@ -157,11 +157,6 @@ public void consume(ParsableByteArray data) throws ParserException {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Parses an AudioMuxElement as defined in 14496-3:2009, Section 1.7.3.1, Table 1.41.
    *

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

@@ -122,11 +122,6 @@ public void consume(ParsableByteArray data) {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Attempts to locate the start of the next frame header.
    *

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

@@ -220,11 +220,6 @@ public void consume(ParsableByteArray data) throws ParserException {
     }
   }
 
-  @Override
-  public void packetFinished(boolean isEndOfInput) {
-    // Do nothing.
-  }
-
   /**
    * Copies data from the provided {@code source} into a given {@code target}, attempting to fill
    * the target buffer up to its limit.

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

@@ -108,12 +108,14 @@ public void consume(ParsableByteArray data, @Flags int flags) throws ParserExcep
             Log.w(TAG, "Unexpected start indicator: expected " + payloadSize + " more bytes");
           }
           // Either way, notify the reader that it has now finished.
-          boolean isEndOfInput = (data.limit() == 0);
-          reader.packetFinished(isEndOfInput);
+          reader.packetFinished();
           break;
         default:
           throw new IllegalStateException();
       }
+      if (data.limit() == 0) {
+        reader.endOfInputReached();
+      }
       setState(STATE_READING_HEADER);
     }
 
@@ -149,8 +151,7 @@ && continueRead(data, /* target= */ null, extendedHeaderLength)) {
           if (payloadSize != C.LENGTH_UNSET) {
             payloadSize -= readLength;
             if (payloadSize == 0) {
-              // There are bytes left in data, see above, so this is not the end of input
-              reader.packetFinished(/* isEndOfInput= */ false);
+              reader.packetFinished();
               setState(STATE_READING_HEADER);
             }
           }
@@ -168,16 +169,18 @@ && continueRead(data, /* target= */ null, extendedHeaderLength)) {
    *     otherwise.
    */
   public boolean canConsumeSynthesizedEmptyPusi(boolean isModeHls) {
-    // Pusi only payload to trigger end of sample data is only applicable if
-    // pes does not have a length field and body is being read, another exclusion
-    // is due to H262 streams possibly having, in HLS mode, a pes across more than one segment
-    // which would trigger committing an unfinished sample in the middle of the access unit
-
     // Only call parseHeader if isModeHls is true and can parse header as some HLS streams may
     // contain packages.
     boolean headerParsed = !isModeHls || parseHeader();
-    return state == STATE_READING_BODY
-        && payloadSize == C.LENGTH_UNSET
+    // Either the PES packet does not have a length field and the body is being read (where the end
+    // of the stream also signals the end of the body), or we are waiting for the next PES packet to
+    // start and the end of the stream means there won't be any more data.
+    return ((state == STATE_READING_BODY && payloadSize == C.LENGTH_UNSET)
+            || (state == STATE_READING_HEADER))
+        // An empty PES packet with only the PUSI (Payload Unit Start Indicator) flag set, used to
+        // trigger the end of the sample data, is not applicable for H262 streams. These streams
+        // may, in HLS mode, have a PES packet spanning across more than one segment, which would
+        // incorrectly trigger committing an unfinished sample in the middle of the access unit.
         && !(isModeHls && reader instanceof H262Reader)
         && headerParsed;
   }

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

@@ -191,16 +191,19 @@ public int read(ExtractorInput input, PositionHolder seekPosition) throws IOExce
     long peekBytesLeft =
         inputLength != C.LENGTH_UNSET ? inputLength - input.getPeekPosition() : C.LENGTH_UNSET;
     if (peekBytesLeft != C.LENGTH_UNSET && peekBytesLeft < 4) {
+      onEndOfInput();
       return RESULT_END_OF_INPUT;
     }
     // First peek and check what type of start code is next.
     if (!input.peekFully(psPacketBuffer.getData(), 0, 4, true)) {
+      onEndOfInput();
       return RESULT_END_OF_INPUT;
     }
 
     psPacketBuffer.setPosition(0);
     int nextStartCode = psPacketBuffer.readInt();
     if (nextStartCode == MPEG_PROGRAM_END_CODE) {
+      onEndOfInput();
       return RESULT_END_OF_INPUT;
     } else if (nextStartCode == PACK_START_CODE) {
       // Now peek the rest of the pack_header.
@@ -294,6 +297,12 @@ public int read(ExtractorInput input, PositionHolder seekPosition) throws IOExce
 
   // Internals.
 
+  private void onEndOfInput() {
+    for (int i = 0; i < psPayloadReaders.size(); i++) {
+      psPayloadReaders.valueAt(i).consumeEndOfInput();
+    }
+  }
+
   @RequiresNonNull("output")
   private void maybeOutputSeekMap(long inputLength) {
     if (!hasOutputSeekMap) {
@@ -360,7 +369,11 @@ public void consume(ParsableByteArray data) throws ParserException {
       pesPayloadReader.packetStarted(timeUs, TsPayloadReader.FLAG_DATA_ALIGNMENT_INDICATOR);
       pesPayloadReader.consume(data);
       // We always have complete PES packets with program stream.
-      pesPayloadReader.packetFinished(/* isEndOfInput= */ false);
+      pesPayloadReader.packetFinished();
+    }
+
+    private void consumeEndOfInput() {
+      pesPayloadReader.endOfInputReached();
     }
 
     private void parseHeader() {

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

@@ -559,7 +559,7 @@ public void packetStarted(long pesTimeUs, @TsPayloadReader.Flags int flags) {}
     public void consume(ParsableByteArray data) {}
 
     @Override
-    public void packetFinished(boolean isEndOfInput) {
+    public void packetFinished() {
       packetsRead++;
     }
 

libraries/test_data/src/test/assets/extractordumps/ts/sample_h262_mpeg_audio.ps.0.dump

@@ -34,7 +34,7 @@ track 192:
     data = length 418, hash 79CF71F8
 track 224:
   total output bytes = 44056
-  sample count = 2
+  sample count = 3
   format 0:
     id = 224
     containerMimeType = video/mp2p
@@ -51,4 +51,8 @@ track 224:
     time = 80000
     flags = 0
     data = length 17831, hash 5C5A57F5
+  sample 2:
+    time = 120000
+    flags = 0
+    data = length 5579, hash 72A6C31F
 tracksEnded = true

libraries/test_data/src/test/assets/extractordumps/ts/sample_h262_mpeg_audio.ps.1.dump

@@ -18,7 +18,7 @@ track 192:
     sampleRate = 44100
 track 224:
   total output bytes = 33949
-  sample count = 1
+  sample count = 2
   format 0:
     id = 224
     containerMimeType = video/mp2p
@@ -31,4 +31,8 @@ track 224:
     time = 80000
     flags = 0
     data = length 17831, hash 5C5A57F5
+  sample 1:
+    time = 120000
+    flags = 0
+    data = length 5579, hash 72A6C31F
 tracksEnded = true

libraries/test_data/src/test/assets/extractordumps/ts/sample_h262_mpeg_audio.ps.2.dump

@@ -18,7 +18,7 @@ track 192:
     sampleRate = 44100
 track 224:
   total output bytes = 19791
-  sample count = 0
+  sample count = 1
   format 0:
     id = 224
     containerMimeType = video/mp2p
@@ -27,4 +27,8 @@ track 224:
     height = 426
     initializationData:
       data = length 22, hash 743CC6F8
+  sample 0:
+    time = 120000
+    flags = 0
+    data = length 5579, hash 72A6C31F
 tracksEnded = true

libraries/test_data/src/test/assets/extractordumps/ts/sample_h262_mpeg_audio.ps.unknown_length.dump

@@ -31,7 +31,7 @@ track 192:
     data = length 418, hash 79CF71F8
 track 224:
   total output bytes = 44056
-  sample count = 2
+  sample count = 3
   format 0:
     id = 224
     containerMimeType = video/mp2p
@@ -48,4 +48,8 @@ track 224:
     time = 80000
     flags = 0
     data = length 17831, hash 5C5A57F5
+  sample 2:
+    time = 120000
+    flags = 0
+    data = length 5579, hash 72A6C31F
 tracksEnded = true