From db5413ed7a116c8bb0dc7bc695bb3aad72b66fbe Mon Sep 17 00:00:00 2001 From: sr90 Date: Sat, 4 Jul 2020 15:18:30 -0700 Subject: [PATCH] Write to buffer before writing to file for TS segments generation. (#790) The refactoring is needed to address #554. --- packager/media/formats/mp2t/ts_segmenter.cc | 94 +++++---- packager/media/formats/mp2t/ts_segmenter.h | 26 ++- .../formats/mp2t/ts_segmenter_unittest.cc | 50 +++-- packager/media/formats/mp2t/ts_writer.cc | 48 ++--- packager/media/formats/mp2t/ts_writer.h | 21 +- .../media/formats/mp2t/ts_writer_unittest.cc | 182 +++++++----------- 6 files changed, 176 insertions(+), 245 deletions(-) diff --git a/packager/media/formats/mp2t/ts_segmenter.cc b/packager/media/formats/mp2t/ts_segmenter.cc index aaea5ced17..4bedb1e32b 100644 --- a/packager/media/formats/mp2t/ts_segmenter.cc +++ b/packager/media/formats/mp2t/ts_segmenter.cc @@ -15,6 +15,7 @@ #include "packager/media/formats/mp2t/pes_packet.h" #include "packager/media/formats/mp2t/program_map_table_writer.h" #include "packager/status.h" +#include "packager/status_macros.h" namespace shaka { namespace media { @@ -102,14 +103,14 @@ Status TsSegmenter::AddSample(const MediaSample& sample) { if (sample.is_encrypted()) ts_writer_->SignalEncrypted(); - if (!ts_writer_file_opened_ && !sample.is_key_frame()) + if (!segment_started_ && !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."); } - return WritePesPacketsToFile(); + return WritePesPackets(); } void TsSegmenter::InjectTsWriterForTesting(std::unique_ptr writer) { @@ -121,47 +122,41 @@ void TsSegmenter::InjectPesPacketGeneratorForTesting( pes_packet_generator_ = std::move(generator); } -void TsSegmenter::SetTsWriterFileOpenedForTesting(bool value) { - ts_writer_file_opened_ = value; +void TsSegmenter::SetSegmentStartedForTesting(bool value) { + segment_started_ = value; } -Status TsSegmenter::OpenNewSegmentIfClosed(int64_t next_pts) { - if (ts_writer_file_opened_) +Status TsSegmenter::StartSegmentIfNeeded(int64_t next_pts) { + if (segment_started_) return Status::OK; - const std::string segment_name = - GetSegmentName(muxer_options_.segment_template, next_pts, - segment_number_++, muxer_options_.bandwidth); - if (!ts_writer_->NewSegment(segment_name)) - return Status(error::MUXER_FAILURE, "Failed to initilize TsPacketWriter."); - current_segment_path_ = segment_name; - ts_writer_file_opened_ = true; + segment_start_timestamp_ = next_pts; + if (!ts_writer_->NewSegment(&segment_buffer_)) + return Status(error::MUXER_FAILURE, "Failed to initialize new segment."); + segment_started_ = true; return Status::OK; } -Status TsSegmenter::WritePesPacketsToFile() { +Status TsSegmenter::WritePesPackets() { while (pes_packet_generator_->NumberOfReadyPesPackets() > 0u) { std::unique_ptr pes_packet = pes_packet_generator_->GetNextPesPacket(); - Status status = OpenNewSegmentIfClosed(pes_packet->pts()); + Status status = StartSegmentIfNeeded(pes_packet->pts()); if (!status.ok()) return status; if (listener_ && IsVideoCodec(codec_) && pes_packet->is_key_frame()) { - base::Optional start_pos = ts_writer_->GetFilePosition(); + uint64_t start_pos = segment_buffer_.Size(); const int64_t timestamp = pes_packet->pts(); - if (!ts_writer_->AddPesPacket(std::move(pes_packet))) + if (!ts_writer_->AddPesPacket(std::move(pes_packet), &segment_buffer_)) return Status(error::MUXER_FAILURE, "Failed to add PES packet."); - base::Optional 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); + uint64_t end_pos = segment_buffer_.Size(); + + listener_->OnKeyFrame(timestamp, start_pos, end_pos - start_pos); } else { - if (!ts_writer_->AddPesPacket(std::move(pes_packet))) + if (!ts_writer_->AddPesPacket(std::move(pes_packet), &segment_buffer_)) return Status(error::MUXER_FAILURE, "Failed to add PES packet."); } } @@ -171,30 +166,45 @@ Status TsSegmenter::WritePesPacketsToFile() { Status TsSegmenter::FinalizeSegment(uint64_t start_timestamp, uint64_t duration) { if (!pes_packet_generator_->Flush()) { - return Status(error::MUXER_FAILURE, - "Failed to flush PesPacketGenerator."); + return Status(error::MUXER_FAILURE, "Failed to flush PesPacketGenerator."); } - Status status = WritePesPacketsToFile(); + Status status = WritePesPackets(); if (!status.ok()) return status; - // This method may be called from Finalize() so ts_writer_file_opened_ could + // This method may be called from Finalize() so segment_started_ could // be false. - if (ts_writer_file_opened_) { - if (!ts_writer_->FinalizeSegment()) { - return Status(error::MUXER_FAILURE, "Failed to finalize TsWriter."); - } - if (listener_) { - const int64_t file_size = - File::GetFileSize(current_segment_path_.c_str()); - listener_->OnNewSegment(current_segment_path_, - start_timestamp * timescale_scale_ + - transport_stream_timestamp_offset_, - duration * timescale_scale_, file_size); - } - ts_writer_file_opened_ = false; + if (!segment_started_) + return Status::OK; + std::string segment_path = + GetSegmentName(muxer_options_.segment_template, segment_start_timestamp_, + segment_number_++, muxer_options_.bandwidth); + + const int64_t file_size = segment_buffer_.Size(); + std::unique_ptr segment_file; + segment_file.reset(File::Open(segment_path.c_str(), "w")); + if (!segment_file) { + return Status(error::FILE_FAILURE, + "Cannot open file for write " + segment_path); } - current_segment_path_.clear(); + + RETURN_IF_ERROR(segment_buffer_.WriteToFile(segment_file.get())); + + if (!segment_file.release()->Close()) { + return Status( + error::FILE_FAILURE, + "Cannot close file " + segment_path + + ", possibly file permission issue or running out of disk space."); + } + + if (listener_) { + listener_->OnNewSegment(segment_path, + start_timestamp * timescale_scale_ + + transport_stream_timestamp_offset_, + duration * timescale_scale_, file_size); + } + segment_started_ = false; + return Status::OK; } diff --git a/packager/media/formats/mp2t/ts_segmenter.h b/packager/media/formats/mp2t/ts_segmenter.h index b140040eb7..595e36f9e2 100644 --- a/packager/media/formats/mp2t/ts_segmenter.h +++ b/packager/media/formats/mp2t/ts_segmenter.h @@ -68,14 +68,13 @@ class TsSegmenter { std::unique_ptr generator); /// Only for testing. - void SetTsWriterFileOpenedForTesting(bool value); - + void SetSegmentStartedForTesting(bool value); + private: - Status OpenNewSegmentIfClosed(int64_t next_pts); + Status StartSegmentIfNeeded(int64_t next_pts); - // 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(); + // Writes PES packets (carried in TsPackets) to a buffer. + Status WritePesPackets(); const MuxerOptions& muxer_options_; MuxerListener* const listener_; @@ -93,16 +92,15 @@ class TsSegmenter { uint64_t segment_number_ = 0; std::unique_ptr ts_writer_; - // Set to true if TsWriter::NewFile() succeeds, set to false after - // TsWriter::FinalizeFile() succeeds. - bool ts_writer_file_opened_ = false; + + BufferWriter segment_buffer_; + + // Set to true if segment_buffer_ is initialized, set to false after + // FinalizeSegment() succeeds. + bool segment_started_ = false; std::unique_ptr pes_packet_generator_; - // For OnNewSegment(). - // Path of the current segment so that File::GetFileSize() can be used after - // the segment has been finalized. - std::string current_segment_path_; - + int64_t segment_start_timestamp_ = -1; DISALLOW_COPY_AND_ASSIGN(TsSegmenter); }; diff --git a/packager/media/formats/mp2t/ts_segmenter_unittest.cc b/packager/media/formats/mp2t/ts_segmenter_unittest.cc index aa0fd31a9f..61f2289d77 100644 --- a/packager/media/formats/mp2t/ts_segmenter_unittest.cc +++ b/packager/media/formats/mp2t/ts_segmenter_unittest.cc @@ -14,6 +14,7 @@ #include "packager/media/formats/mp2t/program_map_table_writer.h" #include "packager/media/formats/mp2t/ts_segmenter.h" #include "packager/status_test_util.h" +#include "packager/media/base/macros.h" namespace shaka { namespace media { @@ -81,15 +82,17 @@ class MockTsWriter : public TsWriter { // Create a bogus pmt writer, which we don't really care. new VideoProgramMapTableWriter(kUnknownCodec))) {} - MOCK_METHOD1(NewSegment, bool(const std::string& file_name)); + MOCK_METHOD1(NewSegment, bool(BufferWriter* buffer_writer)); MOCK_METHOD0(SignalEncrypted, void()); - MOCK_METHOD0(FinalizeSegment, bool()); // Similar to the hack above but takes a std::unique_ptr. - MOCK_METHOD1(AddPesPacketMock, bool(PesPacket* pes_packet)); - bool AddPesPacket(std::unique_ptr pes_packet) override { + MOCK_METHOD2(AddPesPacketMock, bool(PesPacket* pes_packet, + BufferWriter* buffer_writer)); + bool AddPesPacket(std::unique_ptr pes_packet, + BufferWriter* buffer_writer) override { + buffer_writer->AppendArray(kAnyData, arraysize(kAnyData)); // No need to keep the pes packet around for the current tests. - return AddPesPacketMock(pes_packet.get()); + return AddPesPacketMock(pes_packet.get(), buffer_writer); } }; @@ -144,7 +147,7 @@ TEST_F(TsSegmenterTest, AddSample) { MediaSample::CopyFrom(kAnyData, arraysize(kAnyData), kIsKeyFrame); Sequence writer_sequence; - EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts"))) + EXPECT_CALL(*mock_ts_writer_, NewSegment(_)) .InSequence(writer_sequence) .WillOnce(Return(true)); @@ -159,7 +162,7 @@ TEST_F(TsSegmenterTest, AddSample) { .InSequence(ready_pes_sequence) .WillOnce(Return(0u)); - EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_)) + EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_, _)) .WillOnce(Return(true)); // The pointer is released inside the segmenter. @@ -186,7 +189,7 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) { kTransferCharacteristics, kTrickPlayFactor, kNaluLengthSize, kLanguage, kIsEncrypted)); MuxerOptions options; - options.segment_template = "file$Number$.ts"; + options.segment_template = "memory://file$Number$.ts"; MockMuxerListener mock_listener; TsSegmenter segmenter(options, &mock_listener); @@ -204,15 +207,13 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) { // Doesn't really matter how long this is. sample2->set_duration(kInputTimescale * 7); - // (Finalize is not called at the end of this test so) Expect one segment - // event. The length should be the same as the above sample that exceeds the - // duration. EXPECT_CALL(mock_listener, - OnNewSegment("file1.ts", kFirstPts * kTimeScale / kInputTimescale, + OnNewSegment("memory://file1.ts", + kFirstPts * kTimeScale / kInputTimescale, kTimeScale * 11, _)); Sequence writer_sequence; - EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file1.ts"))) + EXPECT_CALL(*mock_ts_writer_, NewSegment(_)) .InSequence(writer_sequence) .WillOnce(Return(true)); @@ -236,6 +237,7 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) { 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)); @@ -243,14 +245,11 @@ TEST_F(TsSegmenterTest, PassedSegmentDuration) { 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"))) + EXPECT_CALL(*mock_ts_writer_, NewSegment(_)) .InSequence(writer_sequence) .WillOnce(Return(true)); - EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_)) + EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_, _)) .Times(2) .WillRepeatedly(Return(true)); @@ -319,14 +318,13 @@ TEST_F(TsSegmenterTest, FinalizeSegment) { EXPECT_CALL(*mock_pes_packet_generator_, Flush()).WillOnce(Return(true)); EXPECT_CALL(*mock_pes_packet_generator_, NumberOfReadyPesPackets()) .WillOnce(Return(0u)); - EXPECT_CALL(*mock_ts_writer_, FinalizeSegment()).WillOnce(Return(true)); segmenter.InjectPesPacketGeneratorForTesting( std::move(mock_pes_packet_generator_)); EXPECT_OK(segmenter.Initialize(*stream_info)); segmenter.InjectTsWriterForTesting(std::move(mock_ts_writer_)); - segmenter.SetTsWriterFileOpenedForTesting(true); - EXPECT_OK(segmenter.FinalizeSegment(0, 100 /* arbitrary duration */)); + + EXPECT_OK(segmenter.FinalizeSegment(0, 100 /* arbitrary duration*/)); } TEST_F(TsSegmenterTest, EncryptedSample) { @@ -338,14 +336,13 @@ TEST_F(TsSegmenterTest, EncryptedSample) { kIsEncrypted)); MuxerOptions options; - options.segment_template = "file$Number$.ts"; + options.segment_template = "memory://file$Number$.ts"; MockMuxerListener mock_listener; TsSegmenter segmenter(options, &mock_listener); ON_CALL(*mock_ts_writer_, NewSegment(_)).WillByDefault(Return(true)); - ON_CALL(*mock_ts_writer_, FinalizeSegment()).WillByDefault(Return(true)); - ON_CALL(*mock_ts_writer_, AddPesPacketMock(_)).WillByDefault(Return(true)); + ON_CALL(*mock_ts_writer_, AddPesPacketMock(_,_)).WillByDefault(Return(true)); ON_CALL(*mock_pes_packet_generator_, Initialize(_)) .WillByDefault(Return(true)); ON_CALL(*mock_pes_packet_generator_, Flush()).WillByDefault(Return(true)); @@ -384,7 +381,7 @@ TEST_F(TsSegmenterTest, EncryptedSample) { .InSequence(ready_pes_sequence) .WillOnce(Return(0u)); - EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_)) + EXPECT_CALL(*mock_ts_writer_, AddPesPacketMock(_, _)) .Times(2) .WillRepeatedly(Return(true)); @@ -397,6 +394,8 @@ TEST_F(TsSegmenterTest, EncryptedSample) { .InSequence(pes_packet_sequence) .WillOnce(Return(new PesPacket())); + EXPECT_CALL(mock_listener, OnNewSegment("memory://file1.ts", _, _, _)); + MockTsWriter* mock_ts_writer_raw = mock_ts_writer_.get(); segmenter.InjectPesPacketGeneratorForTesting( @@ -405,7 +404,6 @@ TEST_F(TsSegmenterTest, EncryptedSample) { EXPECT_OK(segmenter.Initialize(*stream_info)); segmenter.InjectTsWriterForTesting(std::move(mock_ts_writer_)); EXPECT_OK(segmenter.AddSample(*sample1)); - EXPECT_OK(segmenter.FinalizeSegment(1, sample1->duration())); // Signal encrypted if sample is encrypted. EXPECT_CALL(*mock_ts_writer_raw, SignalEncrypted()); diff --git a/packager/media/formats/mp2t/ts_writer.cc b/packager/media/formats/mp2t/ts_writer.cc index d8dc540178..36bf3a73b0 100644 --- a/packager/media/formats/mp2t/ts_writer.cc +++ b/packager/media/formats/mp2t/ts_writer.cc @@ -87,9 +87,9 @@ void WritePtsOrDts(uint8_t leading_bits, writer->AppendInt(fifth_byte); } -bool WritePesToFile(const PesPacket& pes, - ContinuityCounter* continuity_counter, - File* file) { +bool WritePesToBuffer(const PesPacket& pes, + ContinuityCounter* continuity_counter, + BufferWriter* current_buffer) { // The size of the length field. const int kAdaptationFieldLengthSize = 1; // The size of the flags field. @@ -153,7 +153,9 @@ bool WritePesToFile(const PesPacket& pes, !kPayloadUnitStartIndicator, pid, !kHasPcr, 0, continuity_counter, &output_writer); } - return output_writer.WriteToFile(file).ok(); + + current_buffer->AppendBuffer(output_writer); + return true; } } // namespace @@ -163,16 +165,7 @@ TsWriter::TsWriter(std::unique_ptr pmt_writer) TsWriter::~TsWriter() {} -bool TsWriter::NewSegment(const std::string& file_name) { - if (current_file_) { - LOG(ERROR) << "File " << current_file_->file_name() << " still open."; - return false; - } - current_file_.reset(File::Open(file_name.c_str(), "w")); - if (!current_file_) { - LOG(ERROR) << "Failed to open file " << file_name; - return false; - } +bool TsWriter::NewSegment(BufferWriter* buffer) { BufferWriter psi; WritePatToBuffer(kPat, arraysize(kPat), &pat_continuity_counter_, &psi); @@ -185,11 +178,7 @@ bool TsWriter::NewSegment(const std::string& file_name) { return false; } } - - if (!psi.WriteToFile(current_file_.get()).ok()) { - LOG(ERROR) << "Failed to write PSI to file."; - return false; - } + buffer->AppendBuffer(psi); return true; } @@ -198,15 +187,12 @@ void TsWriter::SignalEncrypted() { encrypted_ = true; } -bool TsWriter::FinalizeSegment() { - return current_file_.release()->Close(); -} +bool TsWriter::AddPesPacket(std::unique_ptr pes_packet, + BufferWriter* buffer) { -bool TsWriter::AddPesPacket(std::unique_ptr pes_packet) { - DCHECK(current_file_); - if (!WritePesToFile(*pes_packet, &elementary_stream_continuity_counter_, - current_file_.get())) { - LOG(ERROR) << "Failed to write pes to file."; + if (!WritePesToBuffer(*pes_packet, &elementary_stream_continuity_counter_, + buffer)) { + LOG(ERROR) << "Failed to write pes to buffer."; return false; } @@ -214,14 +200,6 @@ bool TsWriter::AddPesPacket(std::unique_ptr pes_packet) { return true; } -base::Optional 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 diff --git a/packager/media/formats/mp2t/ts_writer.h b/packager/media/formats/mp2t/ts_writer.h index 43e4c0d09a..dc5057444d 100644 --- a/packager/media/formats/mp2t/ts_writer.h +++ b/packager/media/formats/mp2t/ts_writer.h @@ -16,6 +16,7 @@ #include "packager/file/file.h" #include "packager/file/file_closer.h" #include "packager/media/formats/mp2t/continuity_counter.h" +#include "packager/media/base/buffer_writer.h" namespace shaka { namespace media { @@ -32,27 +33,19 @@ class TsWriter { virtual ~TsWriter(); /// This will fail if the current segment is not finalized. - /// @param file_name is the output file name. - /// @param encrypted must be true if the new segment is encrypted. + /// @param buffer to write segment data. /// @return true on success, false otherwise. - virtual bool NewSegment(const std::string& file_name); + virtual bool NewSegment(BufferWriter* buffer); /// Signals the writer that the rest of the segments are encrypted. virtual void SignalEncrypted(); - /// Flush all the pending PesPackets that have not been written to file and - /// close the file. - /// @return true on success, false otherwise. - virtual bool FinalizeSegment(); - - /// Add PesPacket to the instance. PesPacket might not get written to file + /// Add PesPacket to the instance. PesPacket might not be added to the buffer /// immediately. /// @param pes_packet gets added to the writer. + /// @param buffer to write pes packet. /// @return true on success, false otherwise. - virtual bool AddPesPacket(std::unique_ptr pes_packet); - - /// @return current file position on success, nullopt otherwise. - base::Optional GetFilePosition(); + virtual bool AddPesPacket(std::unique_ptr pes_packet, BufferWriter* buffer); private: TsWriter(const TsWriter&) = delete; @@ -65,8 +58,6 @@ class TsWriter { ContinuityCounter elementary_stream_continuity_counter_; std::unique_ptr pmt_writer_; - - std::unique_ptr current_file_; }; } // namespace mp2t diff --git a/packager/media/formats/mp2t/ts_writer_unittest.cc b/packager/media/formats/mp2t/ts_writer_unittest.cc index 83ceedd0d4..b2d5582c22 100644 --- a/packager/media/formats/mp2t/ts_writer_unittest.cc +++ b/packager/media/formats/mp2t/ts_writer_unittest.cc @@ -83,28 +83,7 @@ ACTION(WriteTwoPmts) { class TsWriterTest : public ::testing::Test { protected: - // Using different file names for each test so that the tests can be run in - // parallel. - void SetUp() override { - base::CreateTemporaryFile(&test_file_path_); - // TODO(rkuroiwa): Use memory file prefix once its exposed. - test_file_name_ = - std::string(kLocalFilePrefix) + test_file_path_.AsUTF8Unsafe(); - } - - void TearDown() override { - const bool kRecursive = true; - base::DeleteFile(test_file_path_, !kRecursive); - } - - bool ReadFileToVector(const base::FilePath& path, std::vector* out) { - std::string content; - if (!base::ReadFileToString(path, &content)) - return false; - out->assign(content.begin(), content.end()); - return true; - } - + // Checks whether |actual|'s prefix matches with |prefix| and the suffix // matches with |suffix|. If there is padding, then padding_length specifies // how long the padding is between prefix and suffix. @@ -130,9 +109,6 @@ class TsWriterTest : public ::testing::Test { actual_suffix); } - std::string test_file_name_; - - base::FilePath test_file_path_; }; // Verify that PAT and PMT are correct for clear segment. @@ -142,15 +118,12 @@ TEST_F(TsWriterTest, ClearH264Psi) { std::unique_ptr mock_pmt_writer( new MockProgramMapTableWriter()); EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(WriteOnePmt()); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - ASSERT_TRUE(ts_writer.FinalizeSegment()); - - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); // 2 TS Packets one for PAT and the fake PMT data. - ASSERT_EQ(376u, content.size()); + ASSERT_EQ(376u, buffer_writer.Size()); const uint8_t kExpectedPatPrefix[] = { 0x47, // Sync byte. @@ -180,9 +153,9 @@ TEST_F(TsWriterTest, ClearH264Psi) { EXPECT_NO_FATAL_FAILURE(ExpectTsPacketEqual( kExpectedPatPrefix, kExpectedPatPrefixSize, 165, kExpectedPatPayload, - arraysize(kExpectedPatPayload), content.data())); + arraysize(kExpectedPatPayload), buffer_writer.Buffer())); - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); } @@ -191,16 +164,15 @@ TEST_F(TsWriterTest, ClearAacPmt) { new MockProgramMapTableWriter()); EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(WriteOnePmt()); + BufferWriter buffer_writer; + TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - ASSERT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); // 2 TS Packets. PAT, PMT. - ASSERT_EQ(376u, content.size()); + ASSERT_EQ(376u, buffer_writer.Size()); - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); } @@ -212,18 +184,15 @@ TEST_F(TsWriterTest, ClearLeadH264Pmt) { EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)) .WillOnce(WriteTwoPmts()); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); + + ASSERT_EQ(564u, buffer_writer.Size()); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); - - ASSERT_EQ(564u, content.size()); - - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 2 * kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + 2 * kTsPacketSize, kTsPacketSize)); } @@ -232,8 +201,9 @@ TEST_F(TsWriterTest, ClearSegmentPmtFailure) { new MockProgramMapTableWriter()); EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(Return(false)); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_FALSE(ts_writer.NewSegment(test_file_name_)); + EXPECT_FALSE(ts_writer.NewSegment(&buffer_writer)); } // Check the encrypted segments' PMT (after clear lead). @@ -244,21 +214,18 @@ TEST_F(TsWriterTest, EncryptedSegmentsH264Pmt) { EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(Return(true)); EXPECT_CALL(*mock_pmt_writer, EncryptedSegmentPmt(_)).WillOnce(WriteOnePmt()); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); + buffer_writer.Clear(); // Overwrite the file but as encrypted segment. ts_writer.SignalEncrypted(); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); + ASSERT_EQ(376u, buffer_writer.Size()); - ASSERT_EQ(376u, content.size()); - - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); } @@ -269,12 +236,12 @@ TEST_F(TsWriterTest, EncryptedSegmentPmtFailure) { EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(Return(true)); EXPECT_CALL(*mock_pmt_writer, EncryptedSegmentPmt(_)).WillOnce(Return(false)); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); ts_writer.SignalEncrypted(); - EXPECT_FALSE(ts_writer.NewSegment(test_file_name_)); + EXPECT_FALSE(ts_writer.NewSegment(&buffer_writer)); } // Same as ClearLeadH264Pmt but for AAC. @@ -284,18 +251,15 @@ TEST_F(TsWriterTest, ClearLeadAacPmt) { EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)) .WillOnce(WriteTwoPmts()); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - ASSERT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); + + ASSERT_EQ(564u, buffer_writer.Size()); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); - - ASSERT_EQ(564u, content.size()); - - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 2 * kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + 2 * kTsPacketSize, kTsPacketSize)); } @@ -307,21 +271,18 @@ TEST_F(TsWriterTest, EncryptedSegmentsAacPmt) { EXPECT_CALL(*mock_pmt_writer, ClearSegmentPmt(_)).WillOnce(Return(true)); EXPECT_CALL(*mock_pmt_writer, EncryptedSegmentPmt(_)).WillOnce(WriteOnePmt()); + BufferWriter buffer_writer; TsWriter ts_writer(std::move(mock_pmt_writer)); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); - + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); + + buffer_writer.Clear(); // Overwrite the file but as encrypted segment. ts_writer.SignalEncrypted(); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); - EXPECT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); + + ASSERT_EQ(376u, buffer_writer.Size()); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); - - ASSERT_EQ(376u, content.size()); - - EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + kTsPacketSize, + EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize, kTsPacketSize)); } @@ -329,7 +290,8 @@ TEST_F(TsWriterTest, EncryptedSegmentsAacPmt) { TEST_F(TsWriterTest, AddPesPacket) { TsWriter ts_writer(std::unique_ptr( new VideoProgramMapTableWriter(kCodecForTesting))); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); + BufferWriter buffer_writer; + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); std::unique_ptr pes(new PesPacket()); pes->set_stream_id(0xE0); @@ -340,13 +302,11 @@ TEST_F(TsWriterTest, AddPesPacket) { }; pes->mutable_data()->assign(kAnyData, kAnyData + arraysize(kAnyData)); - EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes))); - ASSERT_TRUE(ts_writer.FinalizeSegment()); - - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); + EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer)); + // 3 TS Packets. PAT, PMT, and PES. - ASSERT_EQ(564u, content.size()); + + ASSERT_EQ(564u, buffer_writer.Size()); const int kPesStartPosition = 376; @@ -384,14 +344,16 @@ TEST_F(TsWriterTest, AddPesPacket) { EXPECT_NO_FATAL_FAILURE(ExpectTsPacketEqual( kExpectedOutputPrefix, arraysize(kExpectedOutputPrefix), 153, kExpectedPayload, arraysize(kExpectedPayload), - content.data() + kPesStartPosition)); + buffer_writer.Buffer() + kPesStartPosition)); } // Verify that PES packet > 64KiB can be handled. TEST_F(TsWriterTest, BigPesPacket) { TsWriter ts_writer(std::unique_ptr( new VideoProgramMapTableWriter(kCodecForTesting))); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); + + BufferWriter buffer_writer; + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); std::unique_ptr pes(new PesPacket()); pes->set_pts(0); @@ -400,23 +362,20 @@ TEST_F(TsWriterTest, BigPesPacket) { const std::vector big_data(400, 0x23); *pes->mutable_data() = big_data; - EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes))); - ASSERT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer)); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); // The first TsPacket can only carry // 177 (TS packet size - header - adaptation_field) - 19 (PES header data) = // 158 bytes of the PES packet payload. // So this should create // 2 + 1 + ceil((400 - 158) / 184) = 5 TsPackets. // Where 184 is the maxium payload of a TS packet. - EXPECT_EQ(5u * 188, content.size()); + EXPECT_EQ(5u * 188, buffer_writer.Size()); // Check continuity counter. - EXPECT_EQ(0, (content[2 * 188 + 3] & 0xF)); - EXPECT_EQ(1, (content[3 * 188 + 3] & 0xF)); - EXPECT_EQ(2, (content[4 * 188 + 3] & 0xF)); + EXPECT_EQ(0, (buffer_writer.Buffer()[2 * 188 + 3] & 0xF)); + EXPECT_EQ(1, (buffer_writer.Buffer()[3 * 188 + 3] & 0xF)); + EXPECT_EQ(2, (buffer_writer.Buffer()[4 * 188 + 3] & 0xF)); } // Bug found in code review. It should check whether PTS is present not whether @@ -424,7 +383,9 @@ TEST_F(TsWriterTest, BigPesPacket) { TEST_F(TsWriterTest, PesPtsZeroNoDts) { TsWriter ts_writer(std::unique_ptr( new VideoProgramMapTableWriter(kCodecForTesting))); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); + + BufferWriter buffer_writer; + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); std::unique_ptr pes(new PesPacket()); pes->set_stream_id(0xE0); @@ -434,13 +395,10 @@ TEST_F(TsWriterTest, PesPtsZeroNoDts) { }; pes->mutable_data()->assign(kAnyData, kAnyData + arraysize(kAnyData)); - EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes))); - ASSERT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer)); - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); // 3 TS Packets. PAT, PMT, and PES. - ASSERT_EQ(564u, content.size()); + ASSERT_EQ(564u, buffer_writer.Size()); const int kPesStartPosition = 376; @@ -473,7 +431,7 @@ TEST_F(TsWriterTest, PesPtsZeroNoDts) { EXPECT_NO_FATAL_FAILURE(ExpectTsPacketEqual( kExpectedOutputPrefix, arraysize(kExpectedOutputPrefix), 158, kExpectedPayload, arraysize(kExpectedPayload), - content.data() + kPesStartPosition)); + buffer_writer.Buffer() + kPesStartPosition)); } // Verify that TS packet with payload 183 is handled correctly, e.g. @@ -481,7 +439,8 @@ TEST_F(TsWriterTest, PesPtsZeroNoDts) { TEST_F(TsWriterTest, TsPacketPayload183Bytes) { TsWriter ts_writer(std::unique_ptr( new VideoProgramMapTableWriter(kCodecForTesting))); - EXPECT_TRUE(ts_writer.NewSegment(test_file_name_)); + BufferWriter buffer_writer; + EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer)); std::unique_ptr pes(new PesPacket()); pes->set_stream_id(0xE0); @@ -495,8 +454,7 @@ TEST_F(TsWriterTest, TsPacketPayload183Bytes) { std::vector pes_payload(157 + 183, 0xAF); *pes->mutable_data() = pes_payload; - EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes))); - ASSERT_TRUE(ts_writer.FinalizeSegment()); + EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer)); const uint8_t kExpectedOutputPrefix[] = { 0x47, // Sync byte. @@ -506,15 +464,13 @@ TEST_F(TsWriterTest, TsPacketPayload183Bytes) { 0x00, // Adaptation Field length, 1 byte padding. }; - std::vector content; - ASSERT_TRUE(ReadFileToVector(test_file_path_, &content)); // 4 TsPackets. PAT, PMT, TsPacket with PES header, TsPacket rest of PES // payload. - ASSERT_EQ(752u, content.size()); + ASSERT_EQ(752u, buffer_writer.Size()); const int kPesStartPosition = 564; - std::vector actual_prefix(content.data() + kPesStartPosition, - content.data() + kPesStartPosition + 5); + std::vector actual_prefix(buffer_writer.Buffer() + kPesStartPosition, + buffer_writer.Buffer() + kPesStartPosition + 5); EXPECT_EQ( std::vector(kExpectedOutputPrefix, kExpectedOutputPrefix + 5), actual_prefix);