From 11d69896c063ccbc775466d761f901108912a458 Mon Sep 17 00:00:00 2001
From: sr90 <sr1990003@gmail.com>
Date: Mon, 27 Jul 2020 21:32:56 -0700
Subject: [PATCH] Write to memory file before writing to file for WebM segments
 generation (#803)

The change is needed to address #554.
---
 .../formats/webm/multi_segment_segmenter.cc   | 28 +++++++++++++------
 .../formats/webm/multi_segment_segmenter.h    |  2 ++
 2 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/packager/media/formats/webm/multi_segment_segmenter.cc b/packager/media/formats/webm/multi_segment_segmenter.cc
index 4a2f24bae2..5cd46dddae 100644
--- a/packager/media/formats/webm/multi_segment_segmenter.cc
+++ b/packager/media/formats/webm/multi_segment_segmenter.cc
@@ -32,17 +32,28 @@ Status MultiSegmentSegmenter::FinalizeSegment(uint64_t start_timestamp,
     return Status(error::FILE_FAILURE, "Error finalizing segment.");
 
   if (!is_subsegment) {
-    const std::string segment_name = writer_->file()->file_name();
+    std::string segment_name =
+        GetSegmentName(options().segment_template, start_timestamp,
+                       num_segment_, options().bandwidth);
+
     // Close the file, which also does flushing, to make sure the file is
     // written before manifest is updated.
     RETURN_IF_ERROR(writer_->Close());
 
+    if (!File::Copy(temp_file_name_.c_str(), segment_name.c_str()))
+      return Status(error::FILE_FAILURE, "Failure to copy memory file.");
+
+    if (!File::Delete(temp_file_name_.c_str()))
+      return Status(error::FILE_FAILURE, "Failure to delete memory file.");
+
+    num_segment_++;
+
     if (muxer_listener()) {
       const uint64_t size = cluster()->Size();
       muxer_listener()->OnNewSegment(segment_name, start_timestamp,
                                      duration_timestamp, size);
     }
-    VLOG(1) << "WEBM file '" << writer_->file()->file_name() << "' finalized.";
+    VLOG(1) << "WEBM file '" << segment_name << "' finalized.";
   }
   return Status::OK;
 }
@@ -77,15 +88,16 @@ Status MultiSegmentSegmenter::DoFinalize() {
 Status MultiSegmentSegmenter::NewSegment(uint64_t start_timestamp,
                                          bool is_subsegment) {
   if (!is_subsegment) {
-    // Create a new file for the new segment.
-    std::string segment_name =
-        GetSegmentName(options().segment_template, start_timestamp,
-                       num_segment_, options().bandwidth);
+    temp_file_name_ =
+        "memory://" + GetSegmentName(options().segment_template,
+                                     start_timestamp, num_segment_,
+                                     options().bandwidth);
+
     writer_.reset(new MkvWriter);
-    Status status = writer_->Open(segment_name);
+    Status status = writer_->Open(temp_file_name_);
+
     if (!status.ok())
       return status;
-    num_segment_++;
   }
 
   const uint64_t start_timecode = FromBmffTimestamp(start_timestamp);
diff --git a/packager/media/formats/webm/multi_segment_segmenter.h b/packager/media/formats/webm/multi_segment_segmenter.h
index 4d788ca793..a740994260 100644
--- a/packager/media/formats/webm/multi_segment_segmenter.h
+++ b/packager/media/formats/webm/multi_segment_segmenter.h
@@ -8,6 +8,7 @@
 #define PACKAGER_MEDIA_FORMATS_WEBM_MULTI_SEGMENT_SEGMENTER_H_
 
 #include <memory>
+
 #include "packager/media/formats/webm/mkv_writer.h"
 #include "packager/media/formats/webm/segmenter.h"
 #include "packager/status.h"
@@ -47,6 +48,7 @@ class MultiSegmentSegmenter : public Segmenter {
 
   std::unique_ptr<MkvWriter> writer_;
   uint32_t num_segment_;
+  std::string temp_file_name_;
 
   DISALLOW_COPY_AND_ASSIGN(MultiSegmentSegmenter);
 };