diff --git a/packager/media/formats/webvtt/webvtt_pipeline_unittest.cc b/packager/media/formats/webvtt/webvtt_pipeline_unittest.cc index 565858891c..3423e7ae43 100644 --- a/packager/media/formats/webvtt/webvtt_pipeline_unittest.cc +++ b/packager/media/formats/webvtt/webvtt_pipeline_unittest.cc @@ -26,7 +26,6 @@ TEST(WebVttTextPipelineTest, SegmentedOutput) { const char* kOutput1 = "memory://output/template-1.vtt"; const char* kOutput2 = "memory://output/template-2.vtt"; const char* kOutput3 = "memory://output/template-3.vtt"; - const char* kOutput4 = "memory://output/template-4.vtt"; const char* kInputFile = "memory://input.vtt"; const char* kInput = @@ -54,14 +53,8 @@ TEST(WebVttTextPipelineTest, SegmentedOutput) { "Thank you.\n"; // Segment One - // 00:00:00.000 to 00:00:10.000 - const char* kExpectedOutput1 = - "WEBVTT\n" - "\n"; - - // Segment Two // 00:00:10.000 to 00:00:20.000 - const char* kExpectedOutput2 = + const char* kExpectedOutput1 = "WEBVTT\n" "\n" "1\n" @@ -69,9 +62,9 @@ TEST(WebVttTextPipelineTest, SegmentedOutput) { "This blade has a dark past.\n" "\n"; - // Segment Three + // Segment Two // 00:00:20.000 to 00:00:30.000 - const char* kExpectedOutput3 = + const char* kExpectedOutput2 = "WEBVTT\n" "\n" "1\n" @@ -87,9 +80,9 @@ TEST(WebVttTextPipelineTest, SegmentedOutput) { "You're a fool for traveling alone,\nso completely unprepared.\n" "\n"; - // Segment Four + // Segment Three // 00:00:30.000 to 00:00:40.000 - const char* kExpectedOutput4 = + const char* kExpectedOutput3 = "WEBVTT\n" "\n" "3\n" @@ -135,7 +128,6 @@ TEST(WebVttTextPipelineTest, SegmentedOutput) { ASSERT_FILE_STREQ(kOutput1, kExpectedOutput1); ASSERT_FILE_STREQ(kOutput2, kExpectedOutput2); ASSERT_FILE_STREQ(kOutput3, kExpectedOutput3); - ASSERT_FILE_STREQ(kOutput4, kExpectedOutput4); } } // namespace media } // namespace shaka diff --git a/packager/media/formats/webvtt/webvtt_segmenter.cc b/packager/media/formats/webvtt/webvtt_segmenter.cc index d122599fa3..1eb5e5458a 100644 --- a/packager/media/formats/webvtt/webvtt_segmenter.cc +++ b/packager/media/formats/webvtt/webvtt_segmenter.cc @@ -33,14 +33,17 @@ Status WebVttSegmenter::Process(std::unique_ptr stream_data) { } Status WebVttSegmenter::OnFlushRequest(size_t input_stream_index) { - Status status; - while (status.ok() && samples_.size()) { - // It is not possible for there to be any gaps, or else we would have - // already ended the segments before them. So just close the last remaining - // open segments. - OnSegmentEnd(samples_.top().segment); + while (segment_map_.size()) { + uint64_t segment = segment_map_.begin()->first; + + // OnSegmentEnd will remove the segment from the map. + Status status = OnSegmentEnd(segment); + if (!status.ok()) { + return status; + } } - return status.ok() ? FlushAllDownstreams() : status; + + return FlushAllDownstreams(); } Status WebVttSegmenter::OnTextSample(std::shared_ptr sample) { @@ -56,25 +59,34 @@ Status WebVttSegmenter::OnTextSample(std::shared_ptr sample) { // Samples must always be advancing. If a sample comes in out of order, // skip the sample. - if (samples_.size() && samples_.top().segment > start_segment) { + if (head_segment_ > start_segment) { LOG(WARNING) << "New sample has arrived out of order. Skipping sample " << "as segment start is " << start_segment << " and segment " - << "head is " << samples_.top().segment << "."; + << "head is " << head_segment_ << "."; return Status::OK; } - for (uint64_t segment = start_segment; segment <= ending_segment; segment++) { - WebVttSegmentedTextSample seg_sample; - seg_sample.segment = segment; - seg_sample.sample = sample; + // Now that we are accepting this new segment, its starting segment is our + // new head segment. + head_segment_ = start_segment; - samples_.push(seg_sample); + // Add the sample to each segment it spans. + for (uint64_t segment = start_segment; segment <= ending_segment; segment++) { + segment_map_[segment].push_back(sample); } // Output all the segments that come before the start of this cue's first // segment. - for (; current_segment_ < start_segment; current_segment_++) { - Status status = OnSegmentEnd(current_segment_); + while (segment_map_.size()) { + // Since the segments are in accending order, we can break out of the loop + // once we catch-up to the new samples starting segment. + const uint64_t segment = segment_map_.begin()->first; + if (segment >= start_segment) { + break; + } + + // Output the segment. If there is an error, there is no reason to continue. + Status status = OnSegmentEnd(segment); if (!status.ok()) { return status; } @@ -85,10 +97,15 @@ Status WebVttSegmenter::OnTextSample(std::shared_ptr sample) { Status WebVttSegmenter::OnSegmentEnd(uint64_t segment) { Status status; - while (status.ok() && samples_.size() && samples_.top().segment == segment) { - status.Update( - DispatchTextSample(kStreamIndex, std::move(samples_.top().sample))); - samples_.pop(); + + const auto found = segment_map_.find(segment); + if (found != segment_map_.end()) { + for (const auto& sample : found->second) { + status.Update(DispatchTextSample(kStreamIndex, sample)); + } + + // Stop storing the segment. + segment_map_.erase(found); } // Only send the segment info if all the samples were accepted. diff --git a/packager/media/formats/webvtt/webvtt_segmenter.h b/packager/media/formats/webvtt/webvtt_segmenter.h index d3d19a78cc..97bd2578c9 100644 --- a/packager/media/formats/webvtt/webvtt_segmenter.h +++ b/packager/media/formats/webvtt/webvtt_segmenter.h @@ -9,8 +9,8 @@ #include -#include -#include +#include +#include #include "packager/media/base/media_handler.h" #include "packager/status.h" @@ -18,29 +18,6 @@ namespace shaka { namespace media { -// Because a text sample can be in multiple segments, this struct -// allows us to associate a segment with a sample. This allows us -// to easily sort samples base on segment then time. -struct WebVttSegmentedTextSample { - uint64_t segment = 0; - std::shared_ptr sample; -}; - -class WebVttSegmentedTextSampleCompare { - public: - bool operator()(const WebVttSegmentedTextSample& left, - const WebVttSegmentedTextSample& right) const { - // If the samples are in the same segment, then the start time is the - // only way to order the two segments. - if (left.segment == right.segment) { - return left.sample->start_time() > right.sample->start_time(); - } - - // Time will not matter as the samples are not in the same segment. - return left.segment > right.segment; - } -}; - class WebVttSegmenter : public MediaHandler { public: explicit WebVttSegmenter(uint64_t segment_duration_ms); @@ -59,12 +36,14 @@ class WebVttSegmenter : public MediaHandler { Status OnSegmentEnd(uint64_t segment); - uint64_t current_segment_ = 0; uint64_t segment_duration_ms_; - std::priority_queue, - WebVttSegmentedTextSampleCompare> - samples_; + + using WebVttSample = std::shared_ptr; + using WebVttSegment = std::vector; + + // Mapping of segment number to segment. + std::map segment_map_; + uint64_t head_segment_ = 0; }; } // namespace media diff --git a/packager/media/formats/webvtt/webvtt_segmenter_unittest.cc b/packager/media/formats/webvtt/webvtt_segmenter_unittest.cc index d8347c1728..731dcb4ef1 100644 --- a/packager/media/formats/webvtt/webvtt_segmenter_unittest.cc +++ b/packager/media/formats/webvtt/webvtt_segmenter_unittest.cc @@ -142,7 +142,7 @@ TEST_F(WebVttSegmenterTest, CreatesSegmentsForCues) { // | | // | | [---B---] // | | -TEST_F(WebVttSegmenterTest, DoesntSkipsEmptySegments) { +TEST_F(WebVttSegmenterTest, SkipsEmptySegments) { const uint64_t kSampleDuration = kSegmentDuration / 2; { @@ -160,13 +160,7 @@ TEST_F(WebVttSegmenterTest, DoesntSkipsEmptySegments) { OnProcess(IsSegmentInfo(kStreamIndex, kStartTimeSigned, kSegmentDuration, !kSubSegment, !kEncrypted))); - // Segment two (empty) - EXPECT_CALL(*Output(kOutputIndex), - OnProcess(IsSegmentInfo( - kStreamIndex, kStartTimeSigned + kSegmentDuration, - kSegmentDuration, !kSubSegment, !kEncrypted))); - - // Segment Three + // Segment Two EXPECT_CALL(*Output(kOutputIndex), OnProcess(IsTextSample( kId[1], kStartTime + 2 * kSegmentDuration,