Support KeyFrame events in TS and MP4

Issue: #287

Change-Id: I7c50853c7cc61c5fadbb620f44f7574a27dc2b68

packager/media/event/event_info.h

@@ -24,7 +24,8 @@ struct SegmentEventInfo {
 
 struct KeyFrameEvent {
   uint64_t timestamp;
-  uint64_t start_byte_offset;
+  // In segment for multi-segment, in subsegment for single-segment.
+  uint64_t start_offset_in_segment;
   uint64_t size;
 };
 

packager/media/event/hls_notify_muxer_listener.cc

@@ -172,10 +172,15 @@ void HlsNotifyMuxerListener::OnMediaEnd(const MediaRanges& media_ranges,
           ++subsegment_index;
           break;
         case EventInfoType::kKeyFrame:
-          hls_notifier_->NotifyKeyFrame(stream_id_,
+          if (subsegment_index < num_subsegments) {
-                                        event_info.key_frame.timestamp,
+            const uint64_t segment_start_offset =
-                                        event_info.key_frame.start_byte_offset,
+                subsegment_ranges[subsegment_index].start;
+            hls_notifier_->NotifyKeyFrame(
+                stream_id_, event_info.key_frame.timestamp,
+                segment_start_offset +
+                    event_info.key_frame.start_offset_in_segment,
                 event_info.key_frame.size);
+          }
           break;
         case EventInfoType::kCue:
           hls_notifier_->NotifyCueEvent(stream_id_,

packager/media/event/hls_notify_muxer_listener_unittest.cc

@@ -69,6 +69,20 @@ const uint8_t kAnyData[] = {
   0xFF, 0x78, 0xAA, 0x6B,
 };
 
+const uint64_t kSegmentStartOffset = 10000;
+const uint64_t kSegmentStartTime = 19283;
+const uint64_t kSegmentDuration = 98028;
+const uint64_t kSegmentSize = 756739;
+
+const uint64_t kCueStartTime = kSegmentStartTime;
+
+const uint64_t kKeyFrameTimestamp = 20123;
+const uint64_t kKeyFrameStartByteOffset = 3456;
+const uint64_t kKeyFrameSize = 543234;
+
+static_assert(kKeyFrameStartByteOffset < kSegmentSize, "");
+static_assert(kKeyFrameStartByteOffset + kKeyFrameSize <= kSegmentSize, "");
+
 // This value doesn't really affect the test, it's not used by the
 // implementation.
 const bool kInitialEncryptionInfo = true;
@@ -343,16 +357,14 @@ TEST_F(HlsNotifyMuxerListenerTest, OnNewSegmentAndCueEvent) {
   listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
                          MuxerListener::kContainerMpeg2ts);
 
-  const uint64_t kStartTime = 19283;
+  EXPECT_CALL(mock_notifier_, NotifyCueEvent(_, kCueStartTime));
-  const uint64_t kDuration = 98028;
+  EXPECT_CALL(
-  const uint64_t kFileSize = 756739;
+      mock_notifier_,
-  EXPECT_CALL(mock_notifier_, NotifyCueEvent(_, kStartTime));
+      NotifyNewSegment(_, StrEq("new_segment_name10.ts"), kSegmentStartTime,
-  EXPECT_CALL(mock_notifier_,
+                       kSegmentDuration, _, kSegmentSize));
-              NotifyNewSegment(_, StrEq("new_segment_name10.ts"), kStartTime,
+  listener_.OnCueEvent(kCueStartTime, "dummy cue data");
-                               kDuration, _, kFileSize));
+  listener_.OnNewSegment("new_segment_name10.ts", kSegmentStartTime,
-  listener_.OnCueEvent(kStartTime, "dummy cue data");
+                         kSegmentDuration, kSegmentSize);
-  listener_.OnNewSegment("new_segment_name10.ts", kStartTime, kDuration,
-                         kFileSize);
 }
 
 // Verify that the notifier is called for every segment in OnMediaEnd if
@@ -368,13 +380,9 @@ TEST_F(HlsNotifyMuxerListenerTest, NoSegmentTemplateOnMediaEnd) {
   listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
                          MuxerListener::kContainerMpeg2ts);
 
-  const uint64_t kSegmentStartOffset = 10000;
+  listener_.OnCueEvent(kCueStartTime, "dummy cue data");
-  const uint64_t kStartTime = 19283;
+  listener_.OnNewSegment("filename.mp4", kSegmentStartTime, kSegmentDuration,
-  const uint64_t kDuration = 98028;
+                         kSegmentSize);
-  const uint64_t kFileSize = 756739;
-
-  listener_.OnCueEvent(kStartTime, "dummy cue data");
-  listener_.OnNewSegment("filename.mp4", kStartTime, kDuration, kFileSize);
   MuxerListener::MediaRanges ranges;
   Range init_range;
   init_range.start = 0;
@@ -386,16 +394,17 @@ TEST_F(HlsNotifyMuxerListenerTest, NoSegmentTemplateOnMediaEnd) {
   std::vector<Range> segment_ranges;
   Range segment_range;
   segment_range.start = kSegmentStartOffset;
-  segment_range.end = kSegmentStartOffset + kFileSize - 1;
+  segment_range.end = kSegmentStartOffset + kSegmentSize - 1;
   segment_ranges.push_back(segment_range);
   ranges.init_range = init_range;
   ranges.index_range = index_range;
   ranges.subsegment_ranges = segment_ranges;
 
-  EXPECT_CALL(mock_notifier_, NotifyCueEvent(_, kStartTime));
+  EXPECT_CALL(mock_notifier_, NotifyCueEvent(_, kCueStartTime));
-  EXPECT_CALL(mock_notifier_,
+  EXPECT_CALL(
-              NotifyNewSegment(_, StrEq("filename.mp4"), kStartTime, kDuration,
+      mock_notifier_,
-                               kSegmentStartOffset, kFileSize));
+      NotifyNewSegment(_, StrEq("filename.mp4"), kSegmentStartTime,
+                       kSegmentDuration, kSegmentStartOffset, kSegmentSize));
   listener_.OnMediaEnd(ranges, 200000);
 }
 
@@ -414,13 +423,8 @@ TEST_F(HlsNotifyMuxerListenerTest,
   listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
                          MuxerListener::kContainerMpeg2ts);
 
-  const uint64_t kSegmentStartOffset = 10000;
+  listener_.OnNewSegment("filename.mp4", kSegmentStartTime, kSegmentDuration,
-  const uint64_t kStartTime = 19283;
+                         kSegmentSize);
-  const uint64_t kDuration = 98028;
-  const uint64_t kFileSize = 756739;
-
-  listener_.OnNewSegment("filename.mp4", kStartTime, kDuration,
-                         kFileSize);
   MuxerListener::MediaRanges ranges;
   Range init_range;
   init_range.start = 0;
@@ -433,7 +437,7 @@ TEST_F(HlsNotifyMuxerListenerTest,
 
   Range segment_range1;
   segment_range1.start = kSegmentStartOffset;
-  segment_range1.end = kSegmentStartOffset + kFileSize - 1;
+  segment_range1.end = kSegmentStartOffset + kSegmentSize - 1;
   segment_ranges.push_back(segment_range1);
 
   Range segment_range2;
@@ -445,9 +449,10 @@ TEST_F(HlsNotifyMuxerListenerTest,
   ranges.index_range = index_range;
   ranges.subsegment_ranges = segment_ranges;
 
-  EXPECT_CALL(mock_notifier_,
+  EXPECT_CALL(
-              NotifyNewSegment(_, StrEq("filename.mp4"), kStartTime,
+      mock_notifier_,
-                               kDuration, kSegmentStartOffset, kFileSize));
+      NotifyNewSegment(_, StrEq("filename.mp4"), kSegmentStartTime,
+                       kSegmentDuration, kSegmentStartOffset, kSegmentSize));
   listener_.OnMediaEnd(ranges, 200000);
 }
 
@@ -475,9 +480,6 @@ TEST_P(HlsNotifyMuxerListenerKeyFrameTest, WithSegmentTemplate) {
   listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
                          MuxerListener::kContainerMpeg2ts);
 
-  const uint64_t kKeyFrameTimestamp = 20123;
-  const uint64_t kKeyFrameStartByteOffset = 3456;
-  const uint64_t kKeyFrameSize = 543234;
   EXPECT_CALL(mock_notifier_,
               NotifyKeyFrame(_, kKeyFrameTimestamp, kKeyFrameStartByteOffset,
                              kKeyFrameSize))
@@ -499,19 +501,24 @@ TEST_P(HlsNotifyMuxerListenerKeyFrameTest, NoSegmentTemplate) {
   listener_.OnMediaStart(muxer_options, *video_stream_info, 90000,
                          MuxerListener::kContainerMpeg2ts);
 
-  const uint64_t kKeyFrameTimestamp = 20123;
-  const uint64_t kKeyFrameStartByteOffset = 3456;
-  const uint64_t kKeyFrameSize = 543234;
   listener_.OnKeyFrame(kKeyFrameTimestamp, kKeyFrameStartByteOffset,
                        kKeyFrameSize);
+  listener_.OnNewSegment("filename.mp4", kSegmentStartTime, kSegmentDuration,
+                         kSegmentSize);
 
   EXPECT_CALL(mock_notifier_,
-              NotifyKeyFrame(_, kKeyFrameTimestamp, kKeyFrameStartByteOffset,
+              NotifyKeyFrame(_, kKeyFrameTimestamp,
+                             kSegmentStartOffset + kKeyFrameStartByteOffset,
                              kKeyFrameSize))
       .Times(GetParam() ? 1 : 0);
+  EXPECT_CALL(
+      mock_notifier_,
+      NotifyNewSegment(_, StrEq("filename.mp4"), kSegmentStartTime,
+                       kSegmentDuration, kSegmentStartOffset, kSegmentSize));
+
   MuxerListener::MediaRanges ranges;
-  // The value does not matter for this test.
+  ranges.subsegment_ranges.push_back(
-  ranges.subsegment_ranges.resize(1);
+      {kSegmentStartOffset, kSegmentStartOffset + kSegmentSize - 1});
   listener_.OnMediaEnd(ranges, 200000);
 }
 

packager/media/formats/mp2t/pes_packet.h

@@ -46,6 +46,11 @@ class PesPacket {
     pts_ = pts;
   }
 
+  /// @return whether it is a key frame.
+  bool is_key_frame() const { return is_key_frame_; }
+  /// @param is_key_frame indicates whether it is a key frame.
+  void set_is_key_frame(bool is_key_frame) { is_key_frame_ = is_key_frame; }
+
   const std::vector<uint8_t>& data() const { return data_; }
   /// @return mutable data for this PES.
   std::vector<uint8_t>* mutable_data() { return &data_; }
@@ -56,6 +61,7 @@ class PesPacket {
   // These values mean "not set" when the value is less than 0.
   int64_t dts_ = -1;
   int64_t pts_ = -1;
+  bool is_key_frame_ = false;
 
   std::vector<uint8_t> data_;
 

packager/media/formats/mp2t/pes_packet_generator.cc

@@ -76,6 +76,7 @@ bool PesPacketGenerator::PushSample(const MediaSample& sample) {
   if (!current_processing_pes_)
     current_processing_pes_.reset(new PesPacket());
 
+  current_processing_pes_->set_is_key_frame(sample.is_key_frame());
   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/ts_segmenter.cc

@@ -143,8 +143,23 @@ Status TsSegmenter::WritePesPacketsToFile() {
     if (!status.ok())
       return status;
 
+    if (listener_ && IsVideoCodec(codec_) && pes_packet->is_key_frame()) {
+      base::Optional<uint64_t> start_pos = ts_writer_->GetFilePosition();
+
+      const int64_t timestamp = pes_packet->pts();
       if (!ts_writer_->AddPesPacket(std::move(pes_packet)))
         return Status(error::MUXER_FAILURE, "Failed to add PES packet.");
+
+      base::Optional<uint64_t> end_pos = ts_writer_->GetFilePosition();
+      if (!start_pos || !end_pos) {
+        return Status(error::MUXER_FAILURE,
+                      "Failed to get file position in WritePesPacketsToFile.");
+      }
+      listener_->OnKeyFrame(timestamp, *start_pos, *end_pos - *start_pos);
+    } else {
+      if (!ts_writer_->AddPesPacket(std::move(pes_packet)))
+        return Status(error::MUXER_FAILURE, "Failed to add PES packet.");
+    }
   }
   return Status::OK;
 }

packager/media/formats/mp2t/ts_writer.cc

@@ -214,6 +214,14 @@ bool TsWriter::AddPesPacket(std::unique_ptr<PesPacket> pes_packet) {
   return true;
 }
 
+base::Optional<uint64_t> TsWriter::GetFilePosition() {
+  if (!current_file_)
+    return base::nullopt;
+  uint64_t position;
+  return current_file_->Tell(&position) ? base::make_optional(position)
+                                        : base::nullopt;
+}
+
 }  // namespace mp2t
 }  // namespace media
 }  // namespace shaka

packager/media/formats/mp2t/ts_writer.h

@@ -12,6 +12,7 @@
 #include <memory>
 #include <vector>
 
+#include "packager/base/optional.h"
 #include "packager/file/file.h"
 #include "packager/file/file_closer.h"
 #include "packager/media/formats/mp2t/continuity_counter.h"
@@ -50,6 +51,9 @@ class TsWriter {
   /// @return true on success, false otherwise.
   virtual bool AddPesPacket(std::unique_ptr<PesPacket> pes_packet);
 
+  /// @return current file position on success, nullopt otherwise.
+  base::Optional<uint64_t> GetFilePosition();
+
  private:
   TsWriter(const TsWriter&) = delete;
   TsWriter& operator=(const TsWriter&) = delete;

packager/media/formats/mp4/fragmenter.cc

@@ -8,10 +8,11 @@
 
 #include <limits>
 
-#include "packager/media/base/buffer_writer.h"
 #include "packager/media/base/audio_stream_info.h"
+#include "packager/media/base/buffer_writer.h"
 #include "packager/media/base/media_sample.h"
 #include "packager/media/formats/mp4/box_definitions.h"
+#include "packager/media/formats/mp4/key_frame_info.h"
 
 namespace shaka {
 namespace media {
@@ -89,6 +90,12 @@ Status Fragmenter::AddSample(const MediaSample& sample) {
         !stream_info_->encryption_config().constant_iv.empty(), traf_);
   }
 
+  if (stream_info_->stream_type() == StreamType::kStreamVideo &&
+      sample.is_key_frame()) {
+    key_frame_infos_.push_back({static_cast<uint64_t>(sample.pts()),
+                                data_->Size(), sample.data_size()});
+  }
+
   data_->AppendArray(sample.data(), sample.data_size());
   fragment_duration_ += sample.duration();
 
@@ -132,6 +139,7 @@ Status Fragmenter::InitializeFragment(int64_t first_sample_dts) {
   earliest_presentation_time_ = kInvalidTime;
   first_sap_time_ = kInvalidTime;
   data_.reset(new BufferWriter());
+  key_frame_infos_.clear();
   return Status::OK;
 }
 

packager/media/formats/mp4/fragmenter.h

@@ -22,6 +22,7 @@ class StreamInfo;
 
 namespace mp4 {
 
+struct KeyFrameInfo;
 struct SegmentReference;
 struct TrackFragment;
 
@@ -62,6 +63,9 @@ class Fragmenter {
   bool fragment_initialized() const { return fragment_initialized_; }
   bool fragment_finalized() const { return fragment_finalized_; }
   BufferWriter* data() { return data_.get(); }
+  const std::vector<KeyFrameInfo>& key_frame_infos() const {
+    return key_frame_infos_;
+  }
 
   /// Set the flag use_decoding_timestamp_in_timeline, which if set to true, use
   /// decoding timestamp instead of presentation timestamp in media timeline,
@@ -96,6 +100,8 @@ class Fragmenter {
   int64_t earliest_presentation_time_;
   int64_t first_sap_time_;
   std::unique_ptr<BufferWriter> data_;
+  // Saves key frames information, for Video.
+  std::vector<KeyFrameInfo> key_frame_infos_;
 
   DISALLOW_COPY_AND_ASSIGN(Fragmenter);
 };

packager/media/formats/mp4/key_frame_info.h

@@ -0,0 +1,27 @@
+// Copyright 2018 Google Inc. All rights reserved.
+//
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file or at
+// https://developers.google.com/open-source/licenses/bsd
+
+#ifndef PACKAGER_MEDIA_FORMATS_MP4_KEY_FRAME_INFO_H_
+#define PACKAGER_MEDIA_FORMATS_MP4_KEY_FRAME_INFO_H_
+
+#include <cstdint>
+
+namespace shaka {
+namespace media {
+namespace mp4 {
+
+/// Tracks key frame information.
+struct KeyFrameInfo {
+  uint64_t timestamp;
+  uint64_t start_byte_offset;
+  uint64_t size;
+};
+
+}  // namespace mp4
+}  // namespace media
+}  // namespace shaka
+
+#endif  // PACKAGER_MEDIA_FORMATS_MP4_KEY_FRAME_INFO_H_

packager/media/formats/mp4/mp4.gyp

@@ -28,6 +28,7 @@
         'decoding_time_iterator.h',
         'fragmenter.cc',
         'fragmenter.h',
+        'key_frame_info.h',
         'mp4_media_parser.cc',
         'mp4_media_parser.h',
         'mp4_muxer.cc',

packager/media/formats/mp4/multi_segment_segmenter.cc

@@ -16,6 +16,7 @@
 #include "packager/media/base/muxer_util.h"
 #include "packager/media/event/muxer_listener.h"
 #include "packager/media/formats/mp4/box_definitions.h"
+#include "packager/media/formats/mp4/key_frame_info.h"
 
 namespace shaka {
 namespace media {
@@ -172,12 +173,22 @@ Status MultiSegmentSegmenter::WriteSegment() {
   if (options().mp4_params.num_subsegments_per_sidx >= 0)
     sidx()->Write(buffer.get());
 
-  const size_t segment_size = buffer->Size() + fragment_buffer()->Size();
+  const size_t segment_header_size = buffer->Size();
+  const size_t segment_size = segment_header_size + fragment_buffer()->Size();
   DCHECK_NE(segment_size, 0u);
 
   Status status = buffer->WriteToFile(file);
-  if (status.ok())
+  if (status.ok()) {
+    if (muxer_listener()) {
+      for (const KeyFrameInfo& key_frame_info : key_frame_infos()) {
+        muxer_listener()->OnKeyFrame(
+            key_frame_info.timestamp,
+            segment_header_size + key_frame_info.start_byte_offset,
+            key_frame_info.size);
+      }
+    }
     status = fragment_buffer()->WriteToFile(file);
+  }
 
   if (!file->Close())
     LOG(WARNING) << "Failed to close the file properly: " << file_name;

packager/media/formats/mp4/segmenter.cc

@@ -18,6 +18,7 @@
 #include "packager/media/event/progress_listener.h"
 #include "packager/media/formats/mp4/box_definitions.h"
 #include "packager/media/formats/mp4/fragmenter.h"
+#include "packager/media/formats/mp4/key_frame_info.h"
 #include "packager/version/version.h"
 
 namespace shaka {
@@ -191,8 +192,13 @@ Status Segmenter::FinalizeSegment(size_t stream_id,
   // Write the fragment to buffer.
   moof_->Write(fragment_buffer_.get());
   mdat.WriteHeader(fragment_buffer_.get());
-  for (const std::unique_ptr<Fragmenter>& fragmenter : fragmenters_)
+  for (const std::unique_ptr<Fragmenter>& fragmenter : fragmenters_) {
+    for (const KeyFrameInfo& key_frame_info : fragmenter->key_frame_infos()) {
+      key_frame_infos_.push_back(key_frame_info);
+      key_frame_infos_.back().start_byte_offset += fragment_buffer_->Size();
+    }
     fragment_buffer_->AppendBuffer(*fragmenter->data());
+  }
 
   // Increase sequence_number for next fragment.
   ++moof_->header.sequence_number;
@@ -203,6 +209,7 @@ Status Segmenter::FinalizeSegment(size_t stream_id,
     Status status = DoFinalizeSegment();
     // Reset segment information to initial state.
     sidx_->references.clear();
+    key_frame_infos_.clear();
     return status;
   }
   return Status::OK;

packager/media/formats/mp4/segmenter.h

@@ -33,6 +33,7 @@ class StreamInfo;
 namespace mp4 {
 
 class Fragmenter;
+struct KeyFrameInfo;
 
 /// This class defines the Segmenter which is responsible for organizing
 /// fragments into segments/subsegments and package them into a MP4 file.
@@ -112,6 +113,9 @@ class Segmenter {
   SegmentIndex* sidx() { return sidx_.get(); }
   MuxerListener* muxer_listener() { return muxer_listener_; }
   uint64_t progress_target() { return progress_target_; }
+  const std::vector<KeyFrameInfo>& key_frame_infos() const {
+    return key_frame_infos_;
+  }
 
   void set_progress_target(uint64_t progress_target) {
     progress_target_ = progress_target;
@@ -142,6 +146,7 @@ class Segmenter {
   uint64_t accumulated_progress_ = 0u;
   uint32_t sample_duration_ = 0u;
   std::vector<uint64_t> stream_durations_;
+  std::vector<KeyFrameInfo> key_frame_infos_;
 
   DISALLOW_COPY_AND_ASSIGN(Segmenter);
 };

packager/media/formats/mp4/single_segment_segmenter.cc

@@ -14,6 +14,7 @@
 #include "packager/media/base/muxer_options.h"
 #include "packager/media/event/progress_listener.h"
 #include "packager/media/formats/mp4/box_definitions.h"
+#include "packager/media/formats/mp4/key_frame_info.h"
 
 namespace shaka {
 namespace media {
@@ -217,6 +218,15 @@ Status SingleSegmentSegmenter::DoFinalizeSegment() {
   }
   vod_sidx_->references.push_back(vod_ref);
 
+  if (muxer_listener()) {
+    for (const KeyFrameInfo& key_frame_info : key_frame_infos()) {
+      // Unlike multisegment-segmenter, there is no (sub)segment header (styp,
+      // sidx), so this is already the offset within the (sub)segment.
+      muxer_listener()->OnKeyFrame(key_frame_info.timestamp,
+                                   key_frame_info.start_byte_offset,
+                                   key_frame_info.size);
+    }
+  }
   // Append fragment buffer to temp file.
   size_t segment_size = fragment_buffer()->Size();
   Status status = fragment_buffer()->WriteToFile(temp_file_.get());