From 154aea7c63545abc4b160b4f49be61f71400c46c Mon Sep 17 00:00:00 2001
From: CorentinB <corentin@archive.org>
Date: Wed, 18 Nov 2020 01:33:32 +0100
Subject: [PATCH] Refactoring

---
 utils.go |  7 -------
 write.go | 13 +++++++++----
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/utils.go b/utils.go
index 2c2d925..c22380d 100644
--- a/utils.go
+++ b/utils.go
@@ -15,13 +15,6 @@ import (
 	"github.com/klauspost/compress/zstd"
 )
 
-func tempFileDone(filePath string) {
-	err := os.Rename(filePath, filePath+".done")
-	if err != nil {
-		panic(err)
-	}
-}
-
 // GetSHA1 return the SHA1 of a []byte,
 // can be used to fill the WARC-Payload-Digest header
 func GetSHA1(content []byte) string {
diff --git a/write.go b/write.go
index 74f341f..a49308a 100644
--- a/write.go
+++ b/write.go
@@ -50,6 +50,15 @@ type Record struct {
 // 	CLRF
 // 	CLRF
 func (w *Writer) WriteRecord(r *Record) (recordID string, err error) {
+	defer func(r *Record) {
+		if r.PayloadPath != "" {
+			err := os.Rename(r.PayloadPath, r.PayloadPath+".done")
+			if err != nil {
+				panic(err)
+			}
+		}
+	}(r)
+
 	// Generate record ID
 	recordID = uuid.NewV4().String()
 
@@ -92,7 +101,6 @@ func (w *Writer) WriteRecord(r *Record) (recordID string, err error) {
 
 		_, err = io.WriteString(w.fileWriter, "\r\n")
 		if err != nil {
-			tempFileDone(r.PayloadPath)
 			return recordID, err
 		}
 
@@ -102,13 +110,11 @@ func (w *Writer) WriteRecord(r *Record) (recordID string, err error) {
 		for {
 			count, err := bufferedReader.Read(buffer)
 			if err != nil && err != io.EOF {
-				tempFileDone(r.PayloadPath)
 				return recordID, err
 			}
 
 			_, err = io.WriteString(w.fileWriter, string(buffer))
 			if err != nil {
-				tempFileDone(r.PayloadPath)
 				return recordID, err
 			}
 
@@ -117,7 +123,6 @@ func (w *Writer) WriteRecord(r *Record) (recordID string, err error) {
 				break
 			}
 		}
-		tempFileDone(r.PayloadPath)
 
 		_, err = io.WriteString(w.fileWriter, "\r\n\r\n")
 		if err != nil {