支持相对路径

Author: wilac-pv

service/s3/s3manager/file_counter.go

@@ -0,0 +1,21 @@
+package s3manager
+
+import "sync/atomic"
+
+type FileCounter struct {
+	TotalNum   int64
+	SuccessNum int64
+	FailNum    int64
+}
+
+func (fc *FileCounter) addTotalNum(num int64) {
+	atomic.AddInt64(&fc.TotalNum, num)
+}
+
+func (fc *FileCounter) addSuccessNum(num int64) {
+	atomic.AddInt64(&fc.SuccessNum, num)
+}
+
+func (fc *FileCounter) addFailNum(num int64) {
+	atomic.AddInt64(&fc.FailNum, num)
+}

service/s3/s3manager/upload.go

@@ -10,6 +10,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"os/user"
 	"path/filepath"
 	"sort"
 	"strings"
@@ -559,17 +560,21 @@ func (uploader *Uploader) UploadDir(rootDir string, bucket, prefix string) error
 	if !strings.HasSuffix(prefix, "/") && len(prefix) > 0 {
 		prefix = prefix + "/"
 	}
-	if !strings.HasSuffix(rootDir, "/") && len(rootDir) > 0 {
-		rootDir = rootDir + "/"
+	rootDir, err := uploader.toAbs(rootDir)
+	if err != nil {
+		return err
 	}
+
 	chFiles := make(chan fileInfoType)
 	var consumerWgc sync.WaitGroup
+	var fileCounter FileCounter
 	for i := 0; i < uploader.opts.Jobs; i++ {
 		consumerWgc.Add(1)
 		go func() {
 			defer consumerWgc.Done()
 			for file := range chFiles {
-				uploader.upload(file, consumerWgc)
+				fileCounter.addTotalNum(1)
+				uploader.upload(file, &fileCounter)
 			}
 		}()
 	}
@@ -590,9 +595,31 @@ func (uploader *Uploader) UploadDir(rootDir string, bucket, prefix string) error
 	})
 	close(chFiles)
 	consumerWgc.Wait()
+	fmt.Printf("Done. Total num: %d, success num: %d, fail num: %d \n", fileCounter.TotalNum, fileCounter.SuccessNum, fileCounter.FailNum)
 	return nil
 }
 
+func (uploader *Uploader) toAbs(rootDir string) (string, error) {
+	if rootDir == "~" || strings.HasPrefix(rootDir, "~/") {
+		currentUser, err := user.Current()
+		if err != nil {
+			log.Fatalf(err.Error())
+			return rootDir, err
+		}
+
+		homeDir := currentUser.HomeDir
+		rootDir = strings.Replace(rootDir, "~", homeDir, 1)
+	}
+	rootDir, err := filepath.Abs(rootDir)
+	if err != nil {
+		return rootDir, err
+	}
+	if !strings.HasSuffix(rootDir, "/") && len(rootDir) > 0 {
+		rootDir = rootDir + "/"
+	}
+	return rootDir, nil
+}
+
 func (uploader *Uploader) uploadFile(fileIfo fileInfoType, call func(success bool)) {
 
 	file, err := os.Open(fileIfo.filePath)
@@ -625,12 +652,15 @@ func makeObjectName(RootDir, Prefix, filePath string) string {
 	objectName := Prefix + resDir
 	return objectName
 }
-func (uploader *Uploader) upload(file fileInfoType, consumerWgc sync.WaitGroup) {
+
+func (uploader *Uploader) upload(file fileInfoType, fileCounter *FileCounter) {
 	uploader.uploadFile(file, func(success bool) {
-		consumerWgc.Done()
-		//completed.Inc()
 		if success {
+			fileCounter.addSuccessNum(1)
 			fmt.Println(fmt.Sprintf("%s successfully uploaded ", file.objectKey))
+		} else {
+			fileCounter.addFailNum(1)
+			fmt.Println(fmt.Sprintf("%s upload failed", file.objectKey))
 		}
 	})
 }