Mirror-Tools
/
shaka-packager
mirror of https://github.com/shaka-project/shaka-packager.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Write to buffer before writing to file for TS segments generation. (#790) 

The refactoring is needed to address #554.
This commit is contained in:
sr90 2020-07-04 15:18:30 -07:00 committed by GitHub
parent 540c0aaffb
commit db5413ed7a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 176 additions and 245 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

94
packager/media/formats/mp2t/ts_segmenter.cc
View File

@ -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<TsWriter> 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<PesPacket> 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<uint64_t> 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<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);
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<File, FileCloser> 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;
}

26
packager/media/formats/mp2t/ts_segmenter.h
View File

@ -68,14 +68,13 @@ class TsSegmenter {
std::unique_ptr<PesPacketGenerator> 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<TsWriter> 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<PesPacketGenerator> 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);
};

50
packager/media/formats/mp2t/ts_segmenter_unittest.cc
View File

@ -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<PesPacket> pes_packet) override {
MOCK_METHOD2(AddPesPacketMock, bool(PesPacket* pes_packet,
BufferWriter* buffer_writer));
bool AddPesPacket(std::unique_ptr<PesPacket> 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());

48
packager/media/formats/mp2t/ts_writer.cc
View File

@ -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<ProgramMapTableWriter> 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<PesPacket> pes_packet,
BufferWriter* buffer) {
bool TsWriter::AddPesPacket(std::unique_ptr<PesPacket> 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<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

21
packager/media/formats/mp2t/ts_writer.h
View File

@ -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<PesPacket> pes_packet);
/// @return current file position on success, nullopt otherwise.
base::Optional<uint64_t> GetFilePosition();
virtual bool AddPesPacket(std::unique_ptr<PesPacket> pes_packet, BufferWriter* buffer);
private:
TsWriter(const TsWriter&) = delete;
@ -65,8 +58,6 @@ class TsWriter {
ContinuityCounter elementary_stream_continuity_counter_;
std::unique_ptr<ProgramMapTableWriter> pmt_writer_;
std::unique_ptr<File, FileCloser> current_file_;
};
} // namespace mp2t

182
packager/media/formats/mp2t/ts_writer_unittest.cc
View File

@ -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<uint8_t>* 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<MockProgramMapTableWriter> 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<uint8_t> 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<uint8_t> 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<uint8_t> 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<uint8_t> 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<uint8_t> 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<uint8_t> 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<ProgramMapTableWriter>(
new VideoProgramMapTableWriter(kCodecForTesting)));
EXPECT_TRUE(ts_writer.NewSegment(test_file_name_));
BufferWriter buffer_writer;
EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
std::unique_ptr<PesPacket> 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<uint8_t> 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<ProgramMapTableWriter>(
new VideoProgramMapTableWriter(kCodecForTesting)));
EXPECT_TRUE(ts_writer.NewSegment(test_file_name_));
BufferWriter buffer_writer;
EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
std::unique_ptr<PesPacket> pes(new PesPacket());
pes->set_pts(0);
@ -400,23 +362,20 @@ TEST_F(TsWriterTest, BigPesPacket) {
const std::vector<uint8_t> 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<uint8_t> 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<ProgramMapTableWriter>(
new VideoProgramMapTableWriter(kCodecForTesting)));
EXPECT_TRUE(ts_writer.NewSegment(test_file_name_));
BufferWriter buffer_writer;
EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
std::unique_ptr<PesPacket> 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<uint8_t> 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<ProgramMapTableWriter>(
new VideoProgramMapTableWriter(kCodecForTesting)));
EXPECT_TRUE(ts_writer.NewSegment(test_file_name_));
BufferWriter buffer_writer;
EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
std::unique_ptr<PesPacket> pes(new PesPacket());
pes->set_stream_id(0xE0);
@ -495,8 +454,7 @@ TEST_F(TsWriterTest, TsPacketPayload183Bytes) {
std::vector<uint8_t> 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<uint8_t> 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<uint8_t> actual_prefix(content.data() + kPesStartPosition,
content.data() + kPesStartPosition + 5);
std::vector<uint8_t> actual_prefix(buffer_writer.Buffer() + kPesStartPosition,
buffer_writer.Buffer() + kPesStartPosition + 5);
EXPECT_EQ(
std::vector<uint8_t>(kExpectedOutputPrefix, kExpectedOutputPrefix + 5),
actual_prefix);