From 40fa6d4be17ac4e0c87feecd0eb5e296d4bcc165 Mon Sep 17 00:00:00 2001
From: Roy Xiang <git@royxiang.me>
Date: Mon, 29 Jul 2024 10:57:34 +0800
Subject: [PATCH] feat(rclone): respect RCLONE_TRANSFERS in environment
 variables

---
 rclone/main.go   | 49 +++++++++++++++++++++++++++++-------------------
 rclone/worker.go | 14 ++++----------
 2 files changed, 34 insertions(+), 29 deletions(-)

diff --git a/rclone/main.go b/rclone/main.go
index 4d3d5ac..2ea8bfe 100644
--- a/rclone/main.go
+++ b/rclone/main.go
@@ -26,15 +26,11 @@ var (
 	delayBeforeTransfer time.Duration
 	excludeFileTypes    []string
 
-	argMultiThreadCutoff  int64
-	argLargeFileTransfers int64
-	argSmallFileTransfers int64
-	argMaxTransfers       int64
-
 	cmdEnv        []string
 	moveArgs      []string
 	largeFileArgs []string
 	smallFileArgs []string
+	folderWeight  int64
 
 	taskChan    chan *putio.FileInfo
 	transferSem *semaphore.Weighted
@@ -45,25 +41,17 @@ var (
 
 func init() {
 	rcGlobalConfig := fs.GetConfig(nil)
-	argMultiThreadCutoff = int64(rcGlobalConfig.MultiThreadCutoff)
-	argLargeFileTransfers = int64(rcGlobalConfig.Transfers)
-	argSmallFileTransfers = argLargeFileTransfers * 2
-	argMaxTransfers = argSmallFileTransfers + 2
+	argMultiThreadCutoff := int64(rcGlobalConfig.MultiThreadCutoff)
+	argLargeFileTransfers := int64(rcGlobalConfig.Transfers)
+	argSmallFileTransfers := argLargeFileTransfers * 2
 
 	moveArgs = []string{
 		"--check-first",
 		"--no-traverse",
 		"--use-mmap",
 	}
-	largeFileArgs = []string{
-		fmt.Sprintf("--transfers=%d", argLargeFileTransfers),
-		fmt.Sprintf("--checkers=%d", rcGlobalConfig.Checkers),
-		fmt.Sprintf("--min-size=%db", argMultiThreadCutoff),
-	}
-	smallFileArgs = []string{
-		fmt.Sprintf("--transfers=%d", argSmallFileTransfers),
-		fmt.Sprintf("--checkers=%d", rcGlobalConfig.Checkers*2),
-	}
+	largeFileArgs = make([]string, 0, 4)
+	smallFileArgs = make([]string, 0, 3)
 
 	osEnv := os.Environ()
 	maxTransfers := 0
@@ -98,6 +86,15 @@ func init() {
 				largeFileArgs = append(largeFileArgs, filterArgs)
 				smallFileArgs = append(smallFileArgs, filterArgs)
 			}
+		case "RCLONE_TRANSFERS":
+			transfers, err := strconv.ParseInt(pair[1], 10, 64)
+			if err != nil {
+				break
+			}
+			if transfers > argLargeFileTransfers {
+				argSmallFileTransfers = transfers * 2
+			}
+			argLargeFileTransfers = transfers
 		default:
 			if pair[0] == "HOME" || strings.HasPrefix(pair[0], "RCLONE_") {
 				cmdEnv = append(cmdEnv, env)
@@ -108,8 +105,22 @@ func init() {
 	accessToken := parseRCloneConfig()
 	Put = putio.New(accessToken, maxTransfers)
 
+	maxWeight := argLargeFileTransfers + 1
+	folderWeight = maxWeight - 1
+	largeFileArgs = append(
+		largeFileArgs,
+		fmt.Sprintf("--min-size=%db", argMultiThreadCutoff),
+		fmt.Sprintf("--transfers=%d", argLargeFileTransfers),
+		fmt.Sprintf("--checkers=%d", argLargeFileTransfers*2),
+	)
+	smallFileArgs = append(
+		largeFileArgs,
+		fmt.Sprintf("--transfers=%d", argSmallFileTransfers),
+		fmt.Sprintf("--checkers=%d", argSmallFileTransfers*2),
+	)
+
 	taskChan = make(chan *putio.FileInfo, 1)
-	transferSem = semaphore.NewWeighted(int64(argMaxTransfers))
+	transferSem = semaphore.NewWeighted(maxWeight)
 }
 
 func Start() {
diff --git a/rclone/worker.go b/rclone/worker.go
index 1eca822..b6f769f 100644
--- a/rclone/worker.go
+++ b/rclone/worker.go
@@ -93,11 +93,11 @@ func moveFolder(folder *putio.FileInfo) {
 		return
 	}
 
-	if err := transferSem.Acquire(context.Background(), argSmallFileTransfers); err != nil {
+	if err := transferSem.Acquire(context.Background(), folderWeight); err != nil {
 		log.Printf("Failed acquiring semaphore while moving folder %s", folder.Name)
 		return
 	}
-	defer transferSem.Release(argSmallFileTransfers)
+	defer transferSem.Release(folderWeight)
 
 	if folder.Size > 0 {
 		log.Printf("Moving folder %s...", folder.Name)
@@ -130,17 +130,11 @@ func moveFile(file *putio.FileInfo) {
 		return
 	}
 
-	var weight int64
-	if file.Size < argMultiThreadCutoff {
-		weight = 1
-	} else {
-		weight = 2
-	}
-	if err := transferSem.Acquire(context.Background(), weight); err != nil {
+	if err := transferSem.Acquire(context.Background(), 1); err != nil {
 		log.Printf("Failed acquiring semaphore while moving file %s", file.Name)
 		return
 	}
-	defer transferSem.Release(weight)
+	defer transferSem.Release(1)
 
 	newFilename := file.Name
 	if strings.HasPrefix(file.ContentType, putio.ContentTypeVideo) {