Skip to content

Commit

Permalink
Write to buffer before writing to file for TS segments generation. (#790
Browse files Browse the repository at this point in the history 
)

The refactoring is needed to address #554.
  • Loading branch information
@sr1990
sr1990 authored Jul 4, 2020
1 parent 540c0aa commit db5413e
Show file tree
Hide file tree
Showing 6 changed files with 176 additions and 245 deletions.
94 changes: 52 additions & 42 deletions packager/media/formats/mp2t/ts_segmenter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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) {
Expand All @@ -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.");
}
}
Expand All @@ -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);
}

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);
}
current_segment_path_.clear();
segment_started_ = false;

return Status::OK;
}

Expand Down
26 changes: 12 additions & 14 deletions packager/media/formats/mp2t/ts_segmenter.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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_;
Expand All @@ -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;
std::unique_ptr<PesPacketGenerator> pes_packet_generator_;

BufferWriter segment_buffer_;

// 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_;
// 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_;

int64_t segment_start_timestamp_ = -1;
DISALLOW_COPY_AND_ASSIGN(TsSegmenter);
};

Expand Down
50 changes: 24 additions & 26 deletions packager/media/formats/mp2t/ts_segmenter_unittest.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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);
}
};

Expand Down Expand Up @@ -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));

Expand All @@ -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.
Expand All @@ -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);
Expand All @@ -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));

Expand All @@ -236,21 +237,19 @@ 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));

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));

Expand Down Expand Up @@ -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) {
Expand All @@ -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));
Expand Down Expand Up @@ -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));

Expand All @@ -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(
Expand All @@ -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());
Expand Down
Loading

0 comments on commit db5413e

Please sign in to comment.