From e999937d75d8e8c40b06add376ebac423b0c2079 Mon Sep 17 00:00:00 2001
From: Keuin <keuinx@gmail.com>
Date: Sat, 10 Sep 2022 16:05:10 +0800
Subject: Fix task is not properly restarted when the live is closed and
 started again. Use more friendly log format to replace golang's default
 `log.Logger`. (not completed) Cleaner task status management.

---
 common/bytesize.go | 19 +++++++++++++++++++
 common/copy.go     | 32 +++++++++++++++++---------------
 2 files changed, 36 insertions(+), 15 deletions(-)
 create mode 100644 common/bytesize.go

(limited to 'common')

diff --git a/common/bytesize.go b/common/bytesize.go
new file mode 100644
index 0000000..9b3aa97
--- /dev/null
+++ b/common/bytesize.go
@@ -0,0 +1,19 @@
+package common
+
+import "fmt"
+
+func PrettyBytes(b uint64) string {
+	if b < 1000 {
+		return fmt.Sprintf("%d Byte", b)
+	}
+	if b < 1000_000 {
+		return fmt.Sprintf("%.2f KiB", float64(b)/1024)
+	}
+	if b < 1000_000_000 {
+		return fmt.Sprintf("%.2f MiB", float64(b)/1024/1024)
+	}
+	if b < 1000_000_000_000 {
+		return fmt.Sprintf("%.2f GiB", float64(b)/1024/1024/1024)
+	}
+	return fmt.Sprintf("%.2f TiB", float64(b)/1024/1024/1024/1024)
+}
diff --git a/common/copy.go b/common/copy.go
index bee6515..1d273dc 100644
--- a/common/copy.go
+++ b/common/copy.go
@@ -36,29 +36,31 @@ func CopyToFileWithBuffer(
 	}()
 
 	for {
-		if err = ctx.Err(); err != nil {
+		select {
+		case <-ctx.Done():
 			return
-		}
-		nRead, err = in.Read(buffer[off:Min[int](off+chunkSize, bufSize)])
-		if err != nil {
-			return
-		}
-		off += nRead
-		if off == bufSize {
-			// buffer is full
-			var nWritten int
-			nWritten, err = out.Write(buffer)
+		default:
+			nRead, err = in.Read(buffer[off:Min[int](off+chunkSize, bufSize)])
 			if err != nil {
 				return
 			}
-			if syncFile {
-				err = out.Sync()
+			off += nRead
+			if off == bufSize {
+				// buffer is full
+				var nWritten int
+				nWritten, err = out.Write(buffer)
 				if err != nil {
 					return
 				}
+				if syncFile {
+					err = out.Sync()
+					if err != nil {
+						return
+					}
+				}
+				written += int64(nWritten)
+				off = 0
 			}
-			written += int64(nWritten)
-			off = 0
 		}
 	}
 }
-- 
cgit v1.2.3