Support key frame events in MuxerListener and HlsNotifier

Issue: #287

Change-Id: I33b91b73988fc40e113bd2e627d08f549a7e3dc5

packager/hls/base/hls_notifier.h

@@ -57,6 +57,17 @@ class HlsNotifier {
                                 uint64_t start_byte_offset,
                                 uint64_t size) = 0;
 
+  /// Called on every key frame. For Video only.
+  /// @param stream_id is the value set by NotifyNewStream().
+  /// @param timestamp is the timesamp of the key frame in timescale units
+  ///        passed in @a media_info.
+  /// @param start_byte_offset is the offset of where the keyframe starts.
+  /// @param size is the size in bytes.
+  virtual bool NotifyKeyFrame(uint32_t stream_id,
+                              uint64_t timestamp,
+                              uint64_t start_byte_offset,
+                              uint64_t size) = 0;
+
   /// @param stream_id is the value set by NotifyNewStream().
   /// @param timestamp is the timestamp of the CueEvent.
   /// @return true on success, false otherwise.

packager/hls/base/mock_media_playlist.h

@@ -31,6 +31,10 @@ class MockMediaPlaylist : public MediaPlaylist {
                     uint64_t duration,
                     uint64_t start_byte_offset,
                     uint64_t size));
+  MOCK_METHOD3(AddKeyFrame,
+               void(uint64_t timestamp,
+                    uint64_t start_byte_offset,
+                    uint64_t size));
   MOCK_METHOD6(AddEncryptionInfo,
                void(EncryptionMethod method,
                     const std::string& url,

packager/hls/base/simple_hls_notifier.cc

@@ -367,6 +367,21 @@ bool SimpleHlsNotifier::NotifyNewSegment(uint32_t stream_id,
   return true;
 }
 
+bool SimpleHlsNotifier::NotifyKeyFrame(uint32_t stream_id,
+                                       uint64_t timestamp,
+                                       uint64_t start_byte_offset,
+                                       uint64_t size) {
+  base::AutoLock auto_lock(lock_);
+  auto stream_iterator = stream_map_.find(stream_id);
+  if (stream_iterator == stream_map_.end()) {
+    LOG(ERROR) << "Cannot find stream with ID: " << stream_id;
+    return false;
+  }
+  auto& media_playlist = stream_iterator->second->media_playlist;
+  media_playlist->AddKeyFrame(timestamp, start_byte_offset, size);
+  return true;
+}
+
 bool SimpleHlsNotifier::NotifyCueEvent(uint32_t stream_id, uint64_t timestamp) {
   base::AutoLock auto_lock(lock_);
   auto stream_iterator = stream_map_.find(stream_id);

packager/hls/base/simple_hls_notifier.h

@@ -56,6 +56,10 @@ class SimpleHlsNotifier : public HlsNotifier {
                         uint64_t duration,
                         uint64_t start_byte_offset,
                         uint64_t size) override;
+  bool NotifyKeyFrame(uint32_t stream_id,
+                      uint64_t timestamp,
+                      uint64_t start_byte_offset,
+                      uint64_t size) override;
   bool NotifyCueEvent(uint32_t container_id, uint64_t timestamp) override;
   bool NotifyEncryptionUpdate(
       uint32_t stream_id,

packager/hls/base/simple_hls_notifier_unittest.cc

@@ -459,6 +459,23 @@ TEST_F(SimpleHlsNotifierTest, NotifyNewSegment) {
   EXPECT_TRUE(notifier.Flush());
 }
 
+TEST_F(SimpleHlsNotifierTest, NotifyKeyFrame) {
+  // Pointer released by SimpleHlsNotifier.
+  MockMediaPlaylist* mock_media_playlist =
+      new MockMediaPlaylist(kVodPlaylist, "playlist.m3u8", "", "");
+  SimpleHlsNotifier notifier(hls_params_);
+  const uint32_t stream_id =
+      SetupStream(kCencProtectionScheme, mock_media_playlist, &notifier);
+
+  const uint64_t kTimestamp = 12345;
+  const uint64_t kStartByteOffset = 888;
+  const uint64_t kSize = 555;
+  EXPECT_CALL(*mock_media_playlist,
+              AddKeyFrame(kTimestamp, kStartByteOffset, kSize));
+  EXPECT_TRUE(
+      notifier.NotifyKeyFrame(stream_id, kTimestamp, kStartByteOffset, kSize));
+}
+
 TEST_F(SimpleHlsNotifierTest, NotifyNewSegmentWithoutStreamsRegistered) {
   SimpleHlsNotifier notifier(hls_params_);
   EXPECT_TRUE(notifier.Init());

packager/media/event/combined_muxer_listener.cc

@@ -65,6 +65,14 @@ void CombinedMuxerListener::OnNewSegment(const std::string& file_name,
   }
 }
 
+void CombinedMuxerListener::OnKeyFrame(uint64_t timestamp,
+                                       uint64_t start_byte_offset,
+                                       uint64_t size) {
+  for (auto& listener : muxer_listeners_) {
+    listener->OnKeyFrame(timestamp, start_byte_offset, size);
+  }
+}
+
 void CombinedMuxerListener::OnCueEvent(uint64_t timestamp,
                                        const std::string& cue_data) {
   for (auto& listener : muxer_listeners_) {

packager/media/event/combined_muxer_listener.h

@@ -39,6 +39,9 @@ class CombinedMuxerListener : public MuxerListener {
                     uint64_t start_time,
                     uint64_t duration,
                     uint64_t segment_file_size) override;
+  void OnKeyFrame(uint64_t timestamp,
+                  uint64_t start_byte_offset,
+                  uint64_t size);
   void OnCueEvent(uint64_t timestamp, const std::string& cue_data) override;
 
  private:

packager/media/event/event_info.h

@@ -22,16 +22,29 @@ struct SegmentEventInfo {
   uint64_t segment_file_size;
 };
 
+struct KeyFrameEvent {
+  uint64_t timestamp;
+  uint64_t start_byte_offset;
+  uint64_t size;
+};
+
 // This stores data passed into OnCueEvent() for VOD.
 struct CueEventInfo {
   uint64_t timestamp;
 };
 
+enum class EventInfoType {
+  kSegment,
+  kKeyFrame,
+  kCue,
+};
+
 // This stores data for lazy event callback for VOD.
 struct EventInfo {
-  bool is_cue_event;
+  EventInfoType type;
   union {
     SegmentEventInfo segment_info;
+    KeyFrameEvent key_frame;
     CueEventInfo cue_event_info;
   };
 };

packager/media/event/hls_notify_muxer_listener.cc

@@ -18,10 +18,12 @@ namespace media {
 
 HlsNotifyMuxerListener::HlsNotifyMuxerListener(
     const std::string& playlist_name,
+    bool iframes_only,
     const std::string& ext_x_media_name,
     const std::string& ext_x_media_group_id,
     hls::HlsNotifier* hls_notifier)
     : playlist_name_(playlist_name),
+      iframes_only_(iframes_only),
       ext_x_media_name_(ext_x_media_name),
       ext_x_media_group_id_(ext_x_media_group_id),
       hls_notifier_(hls_notifier) {
@@ -157,10 +159,8 @@ void HlsNotifyMuxerListener::OnMediaEnd(const MediaRanges& media_ranges,
     const size_t num_subsegments = subsegment_ranges.size();
     size_t subsegment_index = 0;
     for (const auto& event_info : event_info_) {
-      if (event_info.is_cue_event) {
+      switch (event_info.type) {
-        hls_notifier_->NotifyCueEvent(stream_id_,
+        case EventInfoType::kSegment:
-                                      event_info.cue_event_info.timestamp);
-      } else {
           if (subsegment_index < num_subsegments) {
             const Range& range = subsegment_ranges[subsegment_index];
             hls_notifier_->NotifyNewSegment(
@@ -170,6 +170,17 @@ void HlsNotifyMuxerListener::OnMediaEnd(const MediaRanges& media_ranges,
                 range.end + 1 - range.start);
           }
           ++subsegment_index;
+          break;
+        case EventInfoType::kKeyFrame:
+          hls_notifier_->NotifyKeyFrame(stream_id_,
+                                        event_info.key_frame.timestamp,
+                                        event_info.key_frame.start_byte_offset,
+                                        event_info.key_frame.size);
+          break;
+        case EventInfoType::kCue:
+          hls_notifier_->NotifyCueEvent(stream_id_,
+                                        event_info.cue_event_info.timestamp);
+          break;
       }
     }
     if (subsegment_index != num_subsegments) {
@@ -187,7 +198,7 @@ void HlsNotifyMuxerListener::OnNewSegment(const std::string& file_name,
                                           uint64_t segment_file_size) {
   if (!media_info_.has_segment_template()) {
     EventInfo event_info;
-    event_info.is_cue_event = false;
+    event_info.type = EventInfoType::kSegment;
     event_info.segment_info = {start_time, duration, segment_file_size};
     event_info_.push_back(event_info);
   } else {
@@ -200,12 +211,29 @@ void HlsNotifyMuxerListener::OnNewSegment(const std::string& file_name,
   }
 }
 
+void HlsNotifyMuxerListener::OnKeyFrame(uint64_t timestamp,
+                                        uint64_t start_byte_offset,
+                                        uint64_t size) {
+  if (!iframes_only_)
+    return;
+  if (!media_info_.has_segment_template()) {
+    EventInfo event_info;
+    event_info.type = EventInfoType::kKeyFrame;
+    event_info.key_frame = {timestamp, start_byte_offset, size};
+    event_info_.push_back(event_info);
+  } else {
+    const bool result = hls_notifier_->NotifyKeyFrame(stream_id_, timestamp,
+                                                      start_byte_offset, size);
+    LOG_IF(WARNING, !result) << "Failed to add new segment.";
+  }
+}
+
 void HlsNotifyMuxerListener::OnCueEvent(uint64_t timestamp,
                                         const std::string& cue_data) {
   // Not using |cue_data| at this moment.
   if (!media_info_.has_segment_template()) {
     EventInfo event_info;
-    event_info.is_cue_event = true;
+    event_info.type = EventInfoType::kCue;
     event_info.cue_event_info = {timestamp};
     event_info_.push_back(event_info);
   } else {

packager/media/event/hls_notify_muxer_listener.h

@@ -25,6 +25,8 @@ namespace media {
 class HlsNotifyMuxerListener : public MuxerListener {
  public:
   /// @param playlist_name is the name of the playlist for the muxer's stream.
+  /// @param iframes_only if true, indicates that it is for iframes-only
+  ///        playlist.
   /// @param ext_x_media_name is the name of this playlist. This is the
   ///        value of the NAME attribute for EXT-X-MEDIA, it is not the same as
   ///        @a playlist_name. This may be empty for video.
@@ -33,6 +35,7 @@ class HlsNotifyMuxerListener : public MuxerListener {
   ///        video.
   /// @param hls_notifier used by this listener. Ownership does not transfer.
   HlsNotifyMuxerListener(const std::string& playlist_name,
+                         bool iframes_only,
                          const std::string& ext_x_media_name,
                          const std::string& ext_x_media_group_id,
                          hls::HlsNotifier* hls_notifier);
@@ -58,6 +61,9 @@ class HlsNotifyMuxerListener : public MuxerListener {
                     uint64_t start_time,
                     uint64_t duration,
                     uint64_t segment_file_size) override;
+  void OnKeyFrame(uint64_t timestamp,
+                  uint64_t start_byte_offset,
+                  uint64_t size);
   void OnCueEvent(uint64_t timestamp, const std::string& cue_data) override;
   /// @}
 
@@ -66,6 +72,7 @@ class HlsNotifyMuxerListener : public MuxerListener {
   HlsNotifyMuxerListener& operator=(const HlsNotifyMuxerListener&) = delete;
 
   const std::string playlist_name_;
+  const bool iframes_only_;
   const std::string ext_x_media_name_;
   const std::string ext_x_media_group_id_;
   hls::HlsNotifier* const hls_notifier_;

packager/media/event/hls_notify_muxer_listener_unittest.cc

@@ -16,9 +16,11 @@
 namespace shaka {
 namespace media {
 
+using ::testing::_;
+using ::testing::Bool;
 using ::testing::Return;
 using ::testing::StrEq;
-using ::testing::_;
+using ::testing::TestWithParam;
 
 namespace {
 
@@ -40,6 +42,11 @@ class MockHlsNotifier : public hls::HlsNotifier {
                     uint64_t duration,
                     uint64_t start_byte_offset,
                     uint64_t size));
+  MOCK_METHOD4(NotifyKeyFrame,
+               bool(uint32_t stream_id,
+                    uint64_t timestamp,
+                    uint64_t start_byte_offset,
+                    uint64_t size));
   MOCK_METHOD2(NotifyCueEvent, bool(uint32_t stream_id, uint64_t timestamp));
   MOCK_METHOD5(
       NotifyEncryptionUpdate,
@@ -65,6 +72,7 @@ const uint8_t kAnyData[] = {
 // This value doesn't really affect the test, it's not used by the
 // implementation.
 const bool kInitialEncryptionInfo = true;
+const bool kIFramesOnlyPlaylist = true;
 
 const char kDefaultPlaylistName[] = "default_playlist.m3u8";
 const char kDefaultName[] = "DEFAULTNAME";
@@ -87,6 +95,7 @@ class HlsNotifyMuxerListenerTest : public ::testing::Test {
  protected:
   HlsNotifyMuxerListenerTest()
       : listener_(kDefaultPlaylistName,
+                  !kIFramesOnlyPlaylist,
                   kDefaultName,
                   kDefaultGroupId,
                   &mock_notifier_) {}
@@ -365,8 +374,7 @@ TEST_F(HlsNotifyMuxerListenerTest, NoSegmentTemplateOnMediaEnd) {
   const uint64_t kFileSize = 756739;
 
   listener_.OnCueEvent(kStartTime, "dummy cue data");
-  listener_.OnNewSegment("filename.mp4", kStartTime, kDuration,
+  listener_.OnNewSegment("filename.mp4", kStartTime, kDuration, kFileSize);
-                         kFileSize);
   MuxerListener::MediaRanges ranges;
   Range init_range;
   init_range.start = 0;
@@ -386,8 +394,8 @@ TEST_F(HlsNotifyMuxerListenerTest, NoSegmentTemplateOnMediaEnd) {
 
   EXPECT_CALL(mock_notifier_, NotifyCueEvent(_, kStartTime));
   EXPECT_CALL(mock_notifier_,
-              NotifyNewSegment(_, StrEq("filename.mp4"), kStartTime,
+              NotifyNewSegment(_, StrEq("filename.mp4"), kStartTime, kDuration,
-                               kDuration, kSegmentStartOffset, kFileSize));
+                               kSegmentStartOffset, kFileSize));
   listener_.OnMediaEnd(ranges, 200000);
 }
 
@@ -443,5 +451,73 @@ TEST_F(HlsNotifyMuxerListenerTest,
   listener_.OnMediaEnd(ranges, 200000);
 }
 
+class HlsNotifyMuxerListenerKeyFrameTest : public TestWithParam<bool> {
+ public:
+  HlsNotifyMuxerListenerKeyFrameTest()
+      : listener_(kDefaultPlaylistName,
+                  GetParam(),
+                  kDefaultName,
+                  kDefaultGroupId,
+                  &mock_notifier_) {}
+
+  MockHlsNotifier mock_notifier_;
+  HlsNotifyMuxerListener listener_;
+};
+
+TEST_P(HlsNotifyMuxerListenerKeyFrameTest, WithSegmentTemplate) {
+  ON_CALL(mock_notifier_, NotifyNewStream(_, _, _, _, _))
+      .WillByDefault(Return(true));
+  VideoStreamInfoParameters video_params = GetDefaultVideoStreamInfoParams();
+  std::shared_ptr<StreamInfo> video_stream_info =
+      CreateVideoStreamInfo(video_params);
+  MuxerOptions muxer_options;
+  muxer_options.segment_template = "$Number$.mp4";
+  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))
+      .Times(GetParam() ? 1 : 0);
+  listener_.OnKeyFrame(kKeyFrameTimestamp, kKeyFrameStartByteOffset,
+                       kKeyFrameSize);
+}
+
+// Verify that the notifier is called for every key frame in OnMediaEnd if
+// segment_template is not set.
+TEST_P(HlsNotifyMuxerListenerKeyFrameTest, NoSegmentTemplate) {
+  ON_CALL(mock_notifier_, NotifyNewStream(_, _, _, _, _))
+      .WillByDefault(Return(true));
+  VideoStreamInfoParameters video_params = GetDefaultVideoStreamInfoParams();
+  std::shared_ptr<StreamInfo> video_stream_info =
+      CreateVideoStreamInfo(video_params);
+  MuxerOptions muxer_options;
+  muxer_options.output_file_name = "filename.mp4";
+  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);
+
+  EXPECT_CALL(mock_notifier_,
+              NotifyKeyFrame(_, kKeyFrameTimestamp, kKeyFrameStartByteOffset,
+                             kKeyFrameSize))
+      .Times(GetParam() ? 1 : 0);
+  MuxerListener::MediaRanges ranges;
+  // The value does not matter for this test.
+  ranges.subsegment_ranges.resize(1);
+  listener_.OnMediaEnd(ranges, 200000);
+}
+
+INSTANTIATE_TEST_CASE_P(InstantiationName,
+                        HlsNotifyMuxerListenerKeyFrameTest,
+                        Bool());
+
 }  // namespace media
 }  // namespace shaka

packager/media/event/mock_muxer_listener.h

@@ -62,6 +62,11 @@ class MockMuxerListener : public MuxerListener {
                     uint64_t duration,
                     uint64_t segment_file_size));
 
+  MOCK_METHOD3(OnKeyFrame,
+               void(uint64_t timestamp,
+                    uint64_t start_byte_offset,
+                    uint64_t size));
+
   MOCK_METHOD2(OnCueEvent,
                void(uint64_t timestamp, const std::string& cue_data));
 };

packager/media/event/mpd_notify_muxer_listener.cc

@@ -133,13 +133,19 @@ void MpdNotifyMuxerListener::OnMediaEnd(const MediaRanges& media_ranges,
   // TODO(rkuroiwa): Use media_ranges.subsegment_ranges instead of caching the
   // subsegments.
   for (const auto& event_info : event_info_) {
-    if (event_info.is_cue_event) {
+    switch (event_info.type) {
-      mpd_notifier_->NotifyCueEvent(id, event_info.cue_event_info.timestamp);
+      case EventInfoType::kSegment:
-    } else {
         mpd_notifier_->NotifyNewSegment(
             id, event_info.segment_info.start_time,
             event_info.segment_info.duration,
             event_info.segment_info.segment_file_size);
+        break;
+      case EventInfoType::kKeyFrame:
+        // NO-OP for DASH.
+        break;
+      case EventInfoType::kCue:
+        mpd_notifier_->NotifyCueEvent(id, event_info.cue_event_info.timestamp);
+        break;
     }
   }
   event_info_.clear();
@@ -158,12 +164,18 @@ void MpdNotifyMuxerListener::OnNewSegment(const std::string& file_name,
       mpd_notifier_->Flush();
   } else {
     EventInfo event_info;
-    event_info.is_cue_event = false;
+    event_info.type = EventInfoType::kSegment;
     event_info.segment_info = {start_time, duration, segment_file_size};
     event_info_.push_back(event_info);
   }
 }
 
+void MpdNotifyMuxerListener::OnKeyFrame(uint64_t timestamp,
+                                        uint64_t start_byte_offset,
+                                        uint64_t size) {
+  // NO-OP for DASH.
+}
+
 void MpdNotifyMuxerListener::OnCueEvent(uint64_t timestamp,
                                         const std::string& cue_data) {
   // Not using |cue_data| at this moment.
@@ -171,7 +183,7 @@ void MpdNotifyMuxerListener::OnCueEvent(uint64_t timestamp,
     mpd_notifier_->NotifyCueEvent(notification_id_, timestamp);
   } else {
     EventInfo event_info;
-    event_info.is_cue_event = true;
+    event_info.type = EventInfoType::kCue;
     event_info.cue_event_info = {timestamp};
     event_info_.push_back(event_info);
   }

packager/media/event/mpd_notify_muxer_listener.h

@@ -50,6 +50,9 @@ class MpdNotifyMuxerListener : public MuxerListener {
                     uint64_t start_time,
                     uint64_t duration,
                     uint64_t segment_file_size) override;
+  void OnKeyFrame(uint64_t timestamp,
+                  uint64_t start_byte_offset,
+                  uint64_t size);
   void OnCueEvent(uint64_t timestamp, const std::string& cue_data) override;
   /// @}
 

packager/media/event/muxer_listener.h

@@ -128,6 +128,15 @@ class MuxerListener {
                             uint64_t duration,
                             uint64_t segment_file_size) = 0;
 
+  /// Called when there is a new key frame. For Video only. Note that it should
+  /// be called before OnNewSegment is called on the containing segment.
+  /// @param timestamp is in terms of the timescale of the media.
+  /// @param start_byte_offset is the offset of where the key frame starts.
+  /// @param size is size in bytes.
+  virtual void OnKeyFrame(uint64_t timestamp,
+                          uint64_t start_byte_offset,
+                          uint64_t size) = 0;
+
   /// Called when there is a new Ad Cue, which should align with (sub)segments.
   /// @param timestamp indicate the cue timestamp.
   /// @param cue_data is the data of the cue.

packager/media/event/muxer_listener_factory.cc

@@ -56,8 +56,9 @@ std::unique_ptr<MuxerListener> CreateHlsListenerInternal(
     hls_playlist_name = base::StringPrintf("stream_%d.m3u8", stream_index);
   }
 
-  std::unique_ptr<MuxerListener> listener(
+  const bool kIFramesOnly = true;
-      new HlsNotifyMuxerListener(hls_playlist_name, name, group_id, notifier));
+  std::unique_ptr<MuxerListener> listener(new HlsNotifyMuxerListener(
+      hls_playlist_name, !kIFramesOnly, name, group_id, notifier));
   return listener;
 }
 }  // namespace

packager/media/event/vod_media_info_dump_muxer_listener.cc

@@ -88,6 +88,10 @@ void VodMediaInfoDumpMuxerListener::OnNewSegment(const std::string& file_name,
                                                  uint64_t duration,
                                                  uint64_t segment_file_size) {}
 
+void VodMediaInfoDumpMuxerListener::OnKeyFrame(uint64_t timestamp,
+                                               uint64_t start_byte_offset,
+                                               uint64_t size) {}
+
 void VodMediaInfoDumpMuxerListener::OnCueEvent(uint64_t timestamp,
                                                const std::string& cue_data) {
   NOTIMPLEMENTED();

packager/media/event/vod_media_info_dump_muxer_listener.h

@@ -50,6 +50,9 @@ class VodMediaInfoDumpMuxerListener : public MuxerListener {
                     uint64_t start_time,
                     uint64_t duration,
                     uint64_t segment_file_size) override;
+  void OnKeyFrame(uint64_t timestamp,
+                  uint64_t start_byte_offset,
+                  uint64_t size);
   void OnCueEvent(uint64_t timestamp, const std::string& cue_data) override;
   /// @}