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

@@ -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) {
+void TsSegmenter::SetSegmentStartedForTesting(bool value) {
-  ts_writer_file_opened_ = value;
+  segment_started_ = value;
 }
 
-Status TsSegmenter::OpenNewSegmentIfClosed(int64_t next_pts) {
+Status TsSegmenter::StartSegmentIfNeeded(int64_t next_pts) {
-  if (ts_writer_file_opened_)
+  if (segment_started_)
     return Status::OK;
-  const std::string segment_name =
+  segment_start_timestamp_ = next_pts;
-      GetSegmentName(muxer_options_.segment_template, next_pts,
+  if (!ts_writer_->NewSegment(&segment_buffer_))
-                     segment_number_++, muxer_options_.bandwidth);
+    return Status(error::MUXER_FAILURE, "Failed to initialize new segment.");
-  if (!ts_writer_->NewSegment(segment_name))
+  segment_started_ = true;
-    return Status(error::MUXER_FAILURE, "Failed to initilize TsPacketWriter.");
-  current_segment_path_ = segment_name;
-  ts_writer_file_opened_ = 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();
+      uint64_t end_pos = segment_buffer_.Size();
-      if (!start_pos || !end_pos) {
+      
-        return Status(error::MUXER_FAILURE,
+      listener_->OnKeyFrame(timestamp, start_pos, end_pos - start_pos);
-                      "Failed to get file position in WritePesPacketsToFile.");
-      }
-      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,
+    return Status(error::MUXER_FAILURE, "Failed to flush PesPacketGenerator.");
-                  "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 (!segment_started_)
-    if (!ts_writer_->FinalizeSegment()) {
+    return Status::OK;
-      return Status(error::MUXER_FAILURE, "Failed to finalize TsWriter.");
+  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_) {
-      const int64_t file_size =
+    listener_->OnNewSegment(segment_path,
-          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;
+  segment_started_ = false;
-  }
+  
-  current_segment_path_.clear();
   return Status::OK;
 }
 

packager/media/formats/mp2t/ts_segmenter.h

@@ -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,
+  // Writes PES packets (carried in TsPackets) to a buffer.
-  // it will open one. This will not close the file.
+  Status WritePesPackets();
-  Status WritePesPacketsToFile();
 
   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.
+  BufferWriter segment_buffer_;
-  bool ts_writer_file_opened_ = false;
+
+  // 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().
+  int64_t segment_start_timestamp_ = -1;
-  // Path of the current segment so that File::GetFileSize() can be used after
-  // the segment has been finalized.
-  std::string current_segment_path_;
-
   DISALLOW_COPY_AND_ASSIGN(TsSegmenter);
 };
 

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));
+  MOCK_METHOD2(AddPesPacketMock, bool(PesPacket* pes_packet,
-  bool AddPesPacket(std::unique_ptr<PesPacket> pes_packet) override {
+			  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())
+  EXPECT_CALL(*mock_ts_writer_, NewSegment(_))
-      .InSequence(writer_sequence)
-      .WillOnce(Return(true));
-  EXPECT_CALL(*mock_ts_writer_, NewSegment(StrEq("file2.ts")))
       .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());

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,
+bool WritePesToBuffer(const PesPacket& pes,
                       ContinuityCounter* continuity_counter,
-                    File* file) {
+                      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) {
+bool TsWriter::NewSegment(BufferWriter* buffer) {
-  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;
-  }
 
   BufferWriter psi;
   WritePatToBuffer(kPat, arraysize(kPat), &pat_continuity_counter_, &psi);
@@ -185,11 +178,7 @@ bool TsWriter::NewSegment(const std::string& file_name) {
       return false;
     }
   }
-
+  buffer->AppendBuffer(psi);
-  if (!psi.WriteToFile(current_file_.get()).ok()) {
-    LOG(ERROR) << "Failed to write PSI to file.";
-    return false;
-  }
 
   return true;
 }
@@ -198,15 +187,12 @@ void TsWriter::SignalEncrypted() {
   encrypted_ = true;
 }
 
-bool TsWriter::FinalizeSegment() {
+bool TsWriter::AddPesPacket(std::unique_ptr<PesPacket> pes_packet,
-  return current_file_.release()->Close();
+		            BufferWriter* buffer) {
-}
 
-bool TsWriter::AddPesPacket(std::unique_ptr<PesPacket> pes_packet) {
+  if (!WritePesToBuffer(*pes_packet, &elementary_stream_continuity_counter_,
-  DCHECK(current_file_);
+                        buffer)) {
-  if (!WritePesToFile(*pes_packet, &elementary_stream_continuity_counter_,
+    LOG(ERROR) << "Failed to write pes to buffer.";
-                      current_file_.get())) {
-    LOG(ERROR) << "Failed to write pes to file.";
     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

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 buffer to write segment data.
-  /// @param encrypted must be true if the new segment is encrypted.
   /// @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
+  /// Add PesPacket to the instance. PesPacket might not be added to the buffer
-  /// close the file.
-  /// @return true on success, false otherwise.
-  virtual bool FinalizeSegment();
-
-  /// Add PesPacket to the instance. PesPacket might not get written to file
   /// 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);
+  virtual bool AddPesPacket(std::unique_ptr<PesPacket> pes_packet, BufferWriter* buffer);
-
-  /// @return current file position on success, nullopt otherwise.
-  base::Optional<uint64_t> GetFilePosition();
 
  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

packager/media/formats/mp2t/ts_writer_unittest.cc

@@ -83,27 +83,6 @@ 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
@@ -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_));
+  EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
-
-  std::vector<uint8_t> content;
-  ASSERT_TRUE(ReadFileToVector(test_file_path_, &content));
   // 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_));
+  EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
 
-  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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
   
-  std::vector<uint8_t> content;
+  ASSERT_EQ(564u, buffer_writer.Size());
-  ASSERT_TRUE(ReadFileToVector(test_file_path_, &content));
 
-  ASSERT_EQ(564u, content.size());
+  EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize,
-
-  EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
 
+  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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
 
-  std::vector<uint8_t> content;
+  ASSERT_EQ(376u, buffer_writer.Size());
-  ASSERT_TRUE(ReadFileToVector(test_file_path_, &content));
 
-  ASSERT_EQ(376u, content.size());
+  EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize,
-
-  EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
 
   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_));
+  EXPECT_TRUE(ts_writer.NewSegment(&buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
   
-  std::vector<uint8_t> content;
+  ASSERT_EQ(564u, buffer_writer.Size());
-  ASSERT_TRUE(ReadFileToVector(test_file_path_, &content));
 
-  ASSERT_EQ(564u, content.size());
+  EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize,
-
-  EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
   
+  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.NewSegment(&buffer_writer));
-  EXPECT_TRUE(ts_writer.FinalizeSegment());
   
-  std::vector<uint8_t> content;
+  ASSERT_EQ(376u, buffer_writer.Size());
-  ASSERT_TRUE(ReadFileToVector(test_file_path_, &content));
 
-  ASSERT_EQ(376u, content.size());
+  EXPECT_EQ(0, memcmp(kMockPmtWriterData, buffer_writer.Buffer() + kTsPacketSize,
-
-  EXPECT_EQ(0, memcmp(kMockPmtWriterData, content.data() + 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)));
+  EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
   
-  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;
 
@@ -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)));
+  EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
 
-  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(0, (buffer_writer.Buffer()[2 * 188 + 3] & 0xF));
-  EXPECT_EQ(1, (content[3 * 188 + 3] & 0xF));
+  EXPECT_EQ(1, (buffer_writer.Buffer()[3 * 188 + 3] & 0xF));
-  EXPECT_EQ(2, (content[4 * 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)));
+  EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
 
-  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)));
+  EXPECT_TRUE(ts_writer.AddPesPacket(std::move(pes), &buffer_writer));
-  ASSERT_TRUE(ts_writer.FinalizeSegment());
 
   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,
+  std::vector<uint8_t> actual_prefix(buffer_writer.Buffer() + kPesStartPosition,
-                                     content.data() + kPesStartPosition + 5);
+                                buffer_writer.Buffer() + kPesStartPosition + 5);
   EXPECT_EQ(
       std::vector<uint8_t>(kExpectedOutputPrefix, kExpectedOutputPrefix + 5),
       actual_prefix);