From 32fbadfff7205f94a5089ec8ff2fc1cef30d325d Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Mon, 12 Sep 2022 02:59:36 +0800
Subject: Feature: use alternative file extension name when the downloading is
 not completed.

---
 common/filename.go  |  7 +++++++
 recording/config.go |  5 +++--
 recording/runner.go | 38 ++++++++++++++++++++++++++++++++------
 3 files changed, 42 insertions(+), 8 deletions(-)
 create mode 100644 common/filename.go

diff --git a/common/filename.go b/common/filename.go
new file mode 100644
index 0000000..aedd377
--- /dev/null
+++ b/common/filename.go
@@ -0,0 +1,7 @@
+package common
+
+import "fmt"
+
+func CombineFileName(base string, ext string) string {
+	return fmt.Sprintf("%s.%s", base, ext)
+}
diff --git a/recording/config.go b/recording/config.go
index 76a6f3a..ff3ae2a 100644
--- a/recording/config.go
+++ b/recording/config.go
@@ -21,8 +21,9 @@ type TransportConfig struct {
 }
 
 type DownloadConfig struct {
-	SaveDirectory        string `mapstructure:"save_directory"`
-	DiskWriteBufferBytes int    `mapstructure:"disk_write_buffer_bytes"`
+	SaveDirectory                    string `mapstructure:"save_directory"`
+	DiskWriteBufferBytes             int    `mapstructure:"disk_write_buffer_bytes"`
+	UseSpecialExtNameBeforeFinishing bool   `mapstructure:"use_special_ext_name_when_downloading"`
 }
 
 type WatchConfig struct {
diff --git a/recording/runner.go b/recording/runner.go
index 039b760..fddb562 100644
--- a/recording/runner.go
+++ b/recording/runner.go
@@ -25,6 +25,7 @@ type TaskResult struct {
 }
 
 const kReadChunkSize = 128 * 1024
+const kSpecialExtName = "partial"
 
 // runTaskWithAutoRestart
 // start a monitor&download task.
@@ -193,12 +194,37 @@ func record(
 	}
 	streamSource := urlInfo.Data.URLs[0]
 
-	fileName := fmt.Sprintf(
-		"%s.%s",
-		GenerateFileName(profile.Data.Title, time.Now()),
-		common.Errorable[string](common.GetFileExtensionFromUrl(streamSource.URL)).OrElse("flv"),
-	)
-	filePath := path.Join(task.Download.SaveDirectory, fileName)
+	var extName string
+
+	// the real extension name (without renaming)
+	originalExtName := common.Errorable[string](common.GetFileExtensionFromUrl(streamSource.URL)).OrElse("flv")
+
+	if task.TaskConfig.Download.UseSpecialExtNameBeforeFinishing {
+		extName = kSpecialExtName
+	} else {
+		extName = originalExtName
+	}
+
+	baseName := GenerateFileName(profile.Data.Title, time.Now())
+	fileName := common.CombineFileName(baseName, extName)
+	saveDir := task.Download.SaveDirectory
+	filePath := path.Join(saveDir, fileName)
+
+	// rename the extension name to originalExtName when finish writing
+	defer func() {
+		if extName == originalExtName {
+			return
+		}
+		from := filePath
+		to := path.Join(saveDir, common.CombineFileName(baseName, originalExtName))
+		err := os.Rename(from, to)
+		if err != nil {
+			task.logger.Error("Cannot rename %v to %v: %v", from, to, err)
+			return
+		}
+		task.logger.Info("Rename file \"%s\" to \"%s\".", from, to)
+	}()
+
 	file, err := os.OpenFile(filePath, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0644)
 	if err != nil {
 		task.logger.Error("Cannot open file for writing: %v", err)
-- 
cgit v1.2.3