From a1becee58561451daea4bd69989b0806cfa9ceab Mon Sep 17 00:00:00 2001
From: DCCooper <1866858@gmail.com>
Date: Tue, 16 Nov 2021 17:31:58 +0800
Subject: [PATCH 22/29] perf:use bufio reader instead ioutil.ReadFile

reason: read file with fixed chunk size instead of
the whole file into memory cause memory pressure

Signed-off-by: DCCooper <1866858@gmail.com>
---
 util/cipher.go      | 38 +++++++++++++++++++-----
 util/cipher_test.go | 84 +++++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 114 insertions(+), 8 deletions(-)

diff --git a/util/cipher.go b/util/cipher.go
index ecbbc47..a5e3125 100644
--- a/util/cipher.go
+++ b/util/cipher.go
@@ -14,6 +14,7 @@
 package util
 
 import (
+	"bufio"
 	"crypto"
 	"crypto/aes"
 	"crypto/cipher"
@@ -234,6 +235,33 @@ func ReadPublicKey(path string) (rsa.PublicKey, error) {
 	return *key, nil
 }
 
+func checkSumReader(path string) (string, error) {
+	const bufferSize = 32 * 1024 // 32KB
+
+	file, err := os.Open(filepath.Clean(path))
+	if err != nil {
+		return "", errors.Wrapf(err, "hash file failed")
+	}
+	defer func() {
+		if cErr := file.Close(); cErr != nil && err == nil {
+			err = cErr
+		}
+	}()
+	buf := make([]byte, bufferSize)
+	reader := bufio.NewReader(file)
+	hasher := sha256.New()
+	for {
+		switch n, err := reader.Read(buf); err {
+		case nil:
+			hasher.Write(buf[:n])
+		case io.EOF:
+			return fmt.Sprintf("%x", hasher.Sum(nil)), nil
+		default:
+			return "", err
+		}
+	}
+}
+
 func hashFile(path string) (string, error) {
 	cleanPath := filepath.Clean(path)
 	if f, err := os.Stat(cleanPath); err != nil {
@@ -242,12 +270,7 @@ func hashFile(path string) (string, error) {
 		return "", errors.New("failed to hash directory")
 	}
 
-	file, err := ioutil.ReadFile(cleanPath) // nolint:gosec
-	if err != nil {
-		return "", errors.Wrapf(err, "hash file failed")
-	}
-
-	return fmt.Sprintf("%x", sha256.Sum256(file)), nil
+	return checkSumReader(path)
 }
 
 func hashDir(path string) (string, error) {
@@ -261,11 +284,10 @@ func hashDir(path string) (string, error) {
 			return nil
 		}
 		if !info.IsDir() {
-			f, err := ioutil.ReadFile(cleanPath) // nolint:gosec
+			fileHash, err := hashFile(cleanPath)
 			if err != nil {
 				return err
 			}
-			fileHash := fmt.Sprintf("%x", sha256.Sum256(f))
 			checkSum = fmt.Sprintf("%s%s", checkSum, fileHash)
 		}
 		return nil
diff --git a/util/cipher_test.go b/util/cipher_test.go
index bab6dfe..4bbe894 100644
--- a/util/cipher_test.go
+++ b/util/cipher_test.go
@@ -15,10 +15,13 @@ package util
 
 import (
 	"crypto"
+	"crypto/rand"
 	"crypto/sha1"
 	"crypto/sha256"
 	"crypto/sha512"
+	"fmt"
 	"hash"
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
@@ -31,6 +34,9 @@ import (
 )
 
 const (
+	sizeKB        = 1024
+	sizeMB        = 1024 * sizeKB
+	sizeGB        = 1024 * sizeMB
 	maxRepeatTime = 1000000
 )
 
@@ -453,3 +459,81 @@ func TestCheckSum(t *testing.T) {
 		})
 	}
 }
+
+func createFileWithSize(path string, size int) error {
+	file, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	_, err = io.CopyN(file, rand.Reader, int64(size))
+	return err
+}
+
+func benchmarkSHA256SumWithFileSize(b *testing.B, fileSize int) {
+	b.ReportAllocs()
+	filepath := fs.NewFile(b, b.Name())
+	defer filepath.Remove()
+	_ = createFileWithSize(filepath.Path(), fileSize)
+	b.ResetTimer()
+	for n := 0; n < b.N; n++ {
+		_, _ = SHA256Sum(filepath.Path())
+	}
+}
+
+func BenchmarkSHA256Sum(b *testing.B) {
+	tests := []struct {
+		fileSuffix string
+		fileSize   int
+	}{
+		{fileSuffix: "100MB", fileSize: 100 * sizeMB},
+		{fileSuffix: "200MB", fileSize: 200 * sizeMB},
+		{fileSuffix: "500MB", fileSize: 500 * sizeMB},
+		{fileSuffix: "1GB", fileSize: 1 * sizeGB},
+		{fileSuffix: "2GB", fileSize: 2 * sizeGB},
+		{fileSuffix: "4GB", fileSize: 4 * sizeGB},
+		{fileSuffix: "8GB", fileSize: 8 * sizeGB},
+	}
+
+	for _, t := range tests {
+		name := fmt.Sprintf("BenchmarkSHA256SumWithFileSize_%s", t.fileSuffix)
+		b.Run(name, func(b *testing.B) {
+			benchmarkSHA256SumWithFileSize(b, t.fileSize)
+		})
+	}
+}
+
+func TestCreateFileWithSize(t *testing.T) {
+	newFile := fs.NewFile(t, t.Name())
+	defer newFile.Remove()
+	type args struct {
+		path string
+		size int
+	}
+	tests := []struct {
+		name    string
+		args    args
+		wantErr bool
+	}{
+		{
+			name: "TC-generate 500MB file",
+			args: args{
+				path: newFile.Path(),
+				size: 500 * sizeMB,
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			err := createFileWithSize(tt.args.path, tt.args.size)
+			if (err != nil) != tt.wantErr {
+				t.Errorf("createFileWithSize() error = %v, wantErr %v", err, tt.wantErr)
+			}
+			if err == nil {
+				file, _ := os.Stat(tt.args.path)
+				if file.Size() != int64(tt.args.size) {
+					t.Errorf("createFileWithSize() size = %v, actually %v", tt.args.size, file.Size())
+				}
+			}
+		})
+	}
+}
-- 
1.8.3.1