Only segment before key frames

- Segmenter should not finish a segment unless the next sample is a key
  frame.
- Renamed PesPacketGenerator::Finalize() to Flush().
- Use duration from the sample instead of the one copied to PesPacket.
- Remove duration field from PesPacket.

Issue #84

Change-Id: Icd90e65fd63fdeb955e7abac3473b0b54db6ac4a

packager/media/formats/mp2t/pes_packet.h

@@ -46,14 +46,6 @@ class PesPacket {
     pts_ = pts;
   }
 
-  /// Duration is not really part of PES but its here to calculate stream's
-  /// duration.
-  /// @return duration of this PES in timescale.
-  int64_t duration() const { return duration_; }
-  /// @param duration of this PES.
-  void set_duration(int64_t duration) { duration_ = duration; }
-
-  /// @return data carried by this PES, the payload.
   const std::vector<uint8_t>& data() const { return data_; }
   /// @return mutable data for this PES.
   std::vector<uint8_t>* mutable_data() { return &data_; }
@@ -65,8 +57,6 @@ class PesPacket {
   int64_t dts_ = -1;
   int64_t pts_ = -1;
 
-  int64_t duration_ = 0;
-
   std::vector<uint8_t> data_;
 
   DISALLOW_COPY_AND_ASSIGN(PesPacket);

packager/media/formats/mp2t/pes_packet_generator.cc

@@ -64,7 +64,6 @@ bool PesPacketGenerator::PushSample(scoped_refptr<MediaSample> sample) {
   if (!current_processing_pes_)
     current_processing_pes_.reset(new PesPacket());
 
-  current_processing_pes_->set_duration(sample->duration());
   current_processing_pes_->set_pts(timescale_scale_ * sample->pts());
   current_processing_pes_->set_dts(timescale_scale_ * sample->dts());
   if (stream_type_ == kStreamVideo) {

packager/media/formats/mp2t/pes_packet_generator.h

@@ -56,10 +56,8 @@ class PesPacketGenerator {
   /// @return Next PES packet that is ready.
   virtual scoped_ptr<PesPacket> GetNextPesPacket();
 
-  /// Flush the object. This may create more PesPackets with the stored
-  /// samples.
-  /// It is safe to call NumberOfReadyPesPackets() and GetNextPesPacket() after
-  /// this.
+  /// Flush the object.
+  /// This may increase NumberOfReadyPesPackets().
   /// @return true on success, false otherwise.
   virtual bool Flush();
 

packager/media/formats/mp2t/ts_segmenter.cc

@@ -39,31 +39,32 @@ Status TsSegmenter::Initialize(const StreamInfo& stream_info) {
 }
 
 Status TsSegmenter::Finalize() {
-  if (!pes_packet_generator_->Flush()) {
-    return Status(error::MUXER_FAILURE,
-                  "Failed to finalize PesPacketGenerator.");
-  }
-
-  Status status = WritePesPacketsToFiles();
-  if (!status.ok())
-    return status;
-
-  if (!ts_writer_file_opened_)
-    return Status::OK;
-
-  if (!ts_writer_->FinalizeSegment())
-    return Status(error::MUXER_FAILURE, "Failed to finalize TsPacketWriter.");
-  ts_writer_file_opened_ = false;
-  return Status::OK;
+  return Flush();
 }
 
+// First checks whether the sample is a key frame. If so and the segment has
+// passed the segment duration, then flush the generator and write all the data
+// to file.
 Status TsSegmenter::AddSample(scoped_refptr<MediaSample> sample) {
+  const bool passed_segment_duration =
+      current_segment_total_sample_duration_ > muxer_options_.segment_duration;
+  if (sample->is_key_frame() && passed_segment_duration) {
+    Status status = Flush();
+    if (!status.ok())
+      return status;
+  }
+
+  if (!ts_writer_file_opened_ && !sample->is_key_frame())
+    LOG(WARNING) << "A segment will start with a non key frame.";
+
   if (!pes_packet_generator_->PushSample(sample)) {
     return Status(error::MUXER_FAILURE,
                   "Failed to add sample to PesPacketGenerator.");
   }
-  // TODO(rkuriowa): Only segment files before a key frame.
-  return WritePesPacketsToFiles();
+
+  current_segment_total_sample_duration_ += sample->duration() / kTsTimescale;
+
+  return WritePesPacketsToFile();
 }
 
 void TsSegmenter::InjectTsWriterForTesting(scoped_ptr<TsWriter> writer) {
@@ -79,16 +80,6 @@ void TsSegmenter::SetTsWriterFileOpenedForTesting(bool value) {
   ts_writer_file_opened_ = value;
 }
 
-Status TsSegmenter::FinalizeSegmentIfPastSegmentDuration() {
-  if (current_segment_duration_ > muxer_options_.segment_duration) {
-    if (!ts_writer_->FinalizeSegment())
-      return Status(error::FILE_FAILURE, "Failed to finalize segment.");
-    ts_writer_file_opened_ = false;
-    current_segment_duration_ = 0.0;
-  }
-  return Status::OK;
-}
-
 Status TsSegmenter::OpenNewSegmentIfClosed(uint32_t next_pts) {
   if (ts_writer_file_opened_)
     return Status::OK;
@@ -101,7 +92,7 @@ Status TsSegmenter::OpenNewSegmentIfClosed(uint32_t next_pts) {
   return Status::OK;
 }
 
-Status TsSegmenter::WritePesPacketsToFiles() {
+Status TsSegmenter::WritePesPacketsToFile() {
   while (pes_packet_generator_->NumberOfReadyPesPackets() > 0u) {
     scoped_ptr<PesPacket> pes_packet =
         pes_packet_generator_->GetNextPesPacket();
@@ -110,20 +101,33 @@ Status TsSegmenter::WritePesPacketsToFiles() {
     if (!status.ok())
       return status;
 
-    const double pes_packet_duration = pes_packet->duration();
-
     if (!ts_writer_->AddPesPacket(pes_packet.Pass()))
       return Status(error::MUXER_FAILURE, "Failed to add PES packet.");
-
-    current_segment_duration_ += pes_packet_duration / kTsTimescale;
-
-    status = FinalizeSegmentIfPastSegmentDuration();
-    if (!status.ok())
-      return status;
   }
   return Status::OK;
 }
 
+Status TsSegmenter::Flush() {
+  if (!pes_packet_generator_->Flush()) {
+    return Status(error::MUXER_FAILURE,
+                  "Failed to flush PesPacketGenerator.");
+  }
+  Status status = WritePesPacketsToFile();
+  if (!status.ok())
+    return status;
+
+  // This method may be called from Finalize() so ts_writer_file_opened_ could
+  // be false.
+  if (ts_writer_file_opened_) {
+    if (!ts_writer_->FinalizeSegment()) {
+      return Status(error::MUXER_FAILURE, "Failed to finalize TsWriter.");
+    }
+    ts_writer_file_opened_ = false;
+  }
+  current_segment_total_sample_duration_ = 0.0;
+  return Status::OK;
+}
+
 }  // namespace mp2t
 }  // namespace media
 }  // namespace edash_packager

packager/media/formats/mp2t/ts_segmenter.h

@@ -53,18 +53,21 @@ class TsSegmenter {
   void SetTsWriterFileOpenedForTesting(bool value);
 
  private:
-  Status FinalizeSegmentIfPastSegmentDuration();
   Status OpenNewSegmentIfClosed(uint32_t next_pts);
 
-  // Writes PES packets (carried in TsPackets) to file(s). This will
-  // segment appropriately. The state of file may be open or closed after
-  // calling this.
-  Status WritePesPacketsToFiles();
+  // Writes PES packets (carried in TsPackets) to a file. If a file is not open,
+  // it will open one. This will not close the file.
+  Status WritePesPacketsToFile();
+
+  // Flush all the samples that are (possibly) buffered and write them to the
+  // current segment, this will close the file. If a file is not already opened
+  // before calling this, this will open one and write them to file.
+  Status Flush();
 
   const MuxerOptions& muxer_options_;
 
   // in seconds.
-  double current_segment_duration_ = 0.0;
+  double current_segment_total_sample_duration_ = 0.0;
 
   // Used for segment template.
   uint64_t segment_number_ = 0;

packager/media/formats/mp2t/ts_segmenter_unittest.cc

@@ -153,10 +153,8 @@ TEST_F(TsSegmenterTest, AddSample) {
       .WillOnce(Return(true));
 
   // The pointer is released inside the segmenter.
-  PesPacket* pes = new PesPacket();
-  pes->set_duration(kTimeScale);
   EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
-      .WillOnce(Return(pes));
+      .WillOnce(Return(new PesPacket()));
 
   segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
   segmenter.InjectPesPacketGeneratorForTesting(
@@ -168,7 +166,9 @@ TEST_F(TsSegmenterTest, AddSample) {
 
 // Verify the case where the segment is long enough and the current segment
 // should be closed.
-TEST_F(TsSegmenterTest, PastSegmentDuration) {
+// This will add 2 samples and verify that the first segment is closed when the
+// second sample is added.
+TEST_F(TsSegmenterTest, PassedSegmentDuration) {
   scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
       kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
       kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
@@ -178,8 +178,7 @@ TEST_F(TsSegmenterTest, PastSegmentDuration) {
   options.segment_template = "file$Number$.ts";
   TsSegmenter segmenter(options);
 
-  ON_CALL(*mock_ts_writer_, TimeScale())
-      .WillByDefault(Return(kTimeScale));
+  ON_CALL(*mock_ts_writer_, TimeScale()).WillByDefault(Return(kTimeScale));
 
   EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
   EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
@@ -188,17 +187,40 @@ TEST_F(TsSegmenterTest, PastSegmentDuration) {
   const uint8_t kAnyData[] = {
       0x01, 0x0F, 0x3C,
   };
-  scoped_refptr<MediaSample> sample =
+  scoped_refptr<MediaSample> sample1 =
       MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+  scoped_refptr<MediaSample> sample2 =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+
+  // 11 seconds > 10 seconds (segment duration).
+  // Expect the segment to be finalized.
+  sample1->set_duration(kTimeScale * 11);
+
+  // Doesn't really matter how long this is.
+  sample2->set_duration(kTimeScale * 7);
 
   Sequence writer_sequence;
   EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts")))
       .InSequence(writer_sequence)
       .WillOnce(Return(true));
 
-  EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_)).WillOnce(Return(true));
+  EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_))
+      .Times(2)
+      .WillRepeatedly(Return(true));
 
   Sequence ready_pes_sequence;
+  // First AddSample().
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+  // When Flush() is called, inside second AddSample().
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+  // Still inside AddSample() but after Flush().
   EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
       .InSequence(ready_pes_sequence)
       .WillOnce(Return(1u));
@@ -206,26 +228,35 @@ TEST_F(TsSegmenterTest, PastSegmentDuration) {
       .InSequence(ready_pes_sequence)
       .WillOnce(Return(0u));
 
+  EXPECT_CALL(*mock_pes_packet_generator_, Flush())
+      .WillOnce(Return(true));
+
   EXPECT_CALL(*mock_ts_writer_, FinalizeSegment())
       .InSequence(writer_sequence)
       .WillOnce(Return(true));
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file2.ts")))
+      .InSequence(writer_sequence)
+      .WillOnce(Return(true));
 
   EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
-      .WillOnce(Return(true));
+      .Times(2)
+      .WillRepeatedly(Return(true));
 
-  // The pointer is released inside the segmenter.
-  PesPacket* pes = new PesPacket();
-  // 11 seconds > 10 seconds (segment duration).
-  // Expect the segment to be finalized.
-  pes->set_duration(11 * kTimeScale);
+  // The pointers are released inside the segmenter.
+  Sequence pes_packet_sequence;
   EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
-      .WillOnce(Return(pes));
+      .InSequence(pes_packet_sequence)
+      .WillOnce(Return(new PesPacket()));
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(pes_packet_sequence)
+      .WillOnce(Return(new PesPacket()));
 
   segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
   segmenter.InjectPesPacketGeneratorForTesting(
       mock_pes_packet_generator_.Pass());
   EXPECT_OK(segmenter.Initialize(*stream_info));
-  EXPECT_OK(segmenter.AddSample(sample));
+  EXPECT_OK(segmenter.AddSample(sample1));
+  EXPECT_OK(segmenter.AddSample(sample2));
 }
 
 // Finalize right after Initialize(). The writer will not be initialized.
@@ -244,6 +275,8 @@ TEST_F(TsSegmenterTest, InitializeThenFinalize) {
       .WillOnce(Return(true));
 
   EXPECT_CALL(*mock_pes_packet_generator_, Flush()).WillOnce(Return(true));
+  ON_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .WillByDefault(Return(0));
 
   segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
   segmenter.InjectPesPacketGeneratorForTesting(
@@ -284,8 +317,8 @@ TEST_F(TsSegmenterTest, Finalize) {
   EXPECT_OK(segmenter.Finalize());
 }
 
-// Verify that it can generate multiple segments.
-TEST_F(TsSegmenterTest, MultipleSegments) {
+// Verify that it won't finish a segment if the sample is not a key frame.
+TEST_F(TsSegmenterTest, SegmentOnlyBeforeKeyFrame) {
   scoped_refptr<VideoStreamInfo> stream_info(new VideoStreamInfo(
       kTrackId, kTimeScale, kDuration, kH264VideoCodec, kCodecString, kLanguage,
       kWidth, kHeight, kPixelWidth, kPixelHeight, kTrickPlayRate,
@@ -295,8 +328,7 @@ TEST_F(TsSegmenterTest, MultipleSegments) {
   options.segment_template = "file$Number$.ts";
   TsSegmenter segmenter(options);
 
-  ON_CALL(*mock_ts_writer_, TimeScale())
-      .WillByDefault(Return(kTimeScale));
+  ON_CALL(*mock_ts_writer_, TimeScale()).WillByDefault(Return(kTimeScale));
 
   EXPECT_CALL(*mock_ts_writer_, Initialize(_)).WillOnce(Return(true));
   EXPECT_CALL(*mock_pes_packet_generator_, Initialize(_))
@@ -305,11 +337,24 @@ TEST_F(TsSegmenterTest, MultipleSegments) {
   const uint8_t kAnyData[] = {
       0x01, 0x0F, 0x3C,
   };
-  scoped_refptr<MediaSample> sample =
+  scoped_refptr<MediaSample> key_frame_sample1 =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
+  scoped_refptr<MediaSample> non_key_frame_sample =
+      MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), !kIsKeyFrame);
+  scoped_refptr<MediaSample> key_frame_sample2 =
       MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame);
 
+  // 11 seconds > 10 seconds (segment duration).
+  key_frame_sample1->set_duration(kTimeScale * 11);
+
+  // But since the second sample is not a key frame, it shouldn't be segmented.
+  non_key_frame_sample->set_duration(kTimeScale * 7);
+
+  // Since this is a key frame, it should be segmented when this is added.
+  key_frame_sample2->set_duration(kTimeScale * 3);
+
   EXPECT_CALL(*mock_pes_packet_generator_, PushSample(_))
-      .Times(2)
+      .Times(3)
       .WillRepeatedly(Return(true));
 
   Sequence writer_sequence;
@@ -318,25 +363,33 @@ TEST_F(TsSegmenterTest, MultipleSegments) {
       .WillOnce(Return(true));
 
   Sequence ready_pes_sequence;
+  // First AddSample().
   EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
       .InSequence(ready_pes_sequence)
       .WillOnce(Return(1u));
-
-  // The pointer is released inside the segmenter.
-  PesPacket* pes = new PesPacket();
-  // 11 seconds > 10 seconds (segment duration).
-  // Expect the segment to be finalized.
-  pes->set_duration(11 * kTimeScale);
-  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
       .InSequence(ready_pes_sequence)
-      .WillOnce(Return(pes));
-
+      .WillOnce(Return(0u));
+  // Second AddSample().
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+  // Third AddSample(), in Flush().
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(0u));
+  // Third AddSample() after Flush().
+  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
+      .InSequence(ready_pes_sequence)
+      .WillOnce(Return(1u));
   EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
       .InSequence(ready_pes_sequence)
       .WillOnce(Return(0u));
 
-  EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
-      .InSequence(writer_sequence)
+  EXPECT_CALL(*mock_pes_packet_generator_, Flush())
       .WillOnce(Return(true));
 
   EXPECT_CALL(*mock_ts_writer_, FinalizeSegment())
@@ -348,33 +401,29 @@ TEST_F(TsSegmenterTest, MultipleSegments) {
       .InSequence(writer_sequence)
       .WillOnce(Return(true));
 
-  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
-      .InSequence(ready_pes_sequence)
-      .WillOnce(Return(1u));
-
-  // The pointer is released inside the segmenter.
-  pes = new PesPacket();
-  // 7 < 10 seconds, If FinalizeSegment() is called AddSample will fail (due to
-  // mock returning false by default).
-  pes->set_duration(7 * kTimeScale);
-  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
-      .InSequence(ready_pes_sequence)
-      .WillOnce(Return(pes));
-
-  EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets())
-      .InSequence(ready_pes_sequence)
-      .WillOnce(Return(0u));
-
   EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_))
-      .InSequence(writer_sequence)
-      .WillOnce(Return(true));
+      .Times(3)
+      .WillRepeatedly(Return(true));
+
+  // The pointers are released inside the segmenter.
+  Sequence pes_packet_sequence;
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(pes_packet_sequence)
+      .WillOnce(Return(new PesPacket()));
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(pes_packet_sequence)
+      .WillOnce(Return(new PesPacket()));
+  EXPECT_CALL(*mock_pes_packet_generator_, GetNextPesPacketMock())
+      .InSequence(pes_packet_sequence)
+      .WillOnce(Return(new PesPacket()));
 
   segmenter.InjectTsWriterForTesting(mock_ts_writer_.Pass());
   segmenter.InjectPesPacketGeneratorForTesting(
       mock_pes_packet_generator_.Pass());
   EXPECT_OK(segmenter.Initialize(*stream_info));
-  EXPECT_OK(segmenter.AddSample(sample));
-  EXPECT_OK(segmenter.AddSample(sample));
+  EXPECT_OK(segmenter.AddSample(key_frame_sample1));
+  EXPECT_OK(segmenter.AddSample(non_key_frame_sample));
+  EXPECT_OK(segmenter.AddSample(key_frame_sample2));
 }
 
 }  // namespace mp2t

packager/media/formats/mp2t/ts_writer_unittest.cc

@@ -220,7 +220,6 @@ TEST_F(TsWriterTest, AddPesPacket) {
 
   scoped_ptr<PesPacket> pes(new PesPacket());
   pes->set_stream_id(0xE0);
-  pes->set_duration(99000);
   pes->set_pts(0x900);
   pes->set_dts(0x900);
   const uint8_t kAnyData[] = {
@@ -285,7 +284,6 @@ TEST_F(TsWriterTest, BigPesPacket) {
   EXPECT_TRUE(ts_writer_.NewSegment(test_file_name_));
 
   scoped_ptr<PesPacket> pes(new PesPacket());
-  pes->set_duration(99000);
   pes->set_pts(0);
   pes->set_dts(0);
   // A little over 2 TS Packets (3 TS Packets).
@@ -323,7 +321,6 @@ TEST_F(TsWriterTest, PesPtsZeroNoDts) {
 
   scoped_ptr<PesPacket> pes(new PesPacket());
   pes->set_stream_id(0xE0);
-  pes->set_duration(99000);
   pes->set_pts(0x0);
   const uint8_t kAnyData[] = {
       0x12, 0x88, 0x4F, 0x4A,