feat: add configurable thread count for file processing

krau · Feb 21, 2025 · d7cd2ed · d7cd2ed
1 parent ed21b65
commit d7cd2ed
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/config/viper.go b/config/viper.go
@@ -9,10 +9,12 @@ import (
 )
 
 type Config struct {
-	Workers      int          `toml:"workers" mapstructure:"workers"`
-	Retry        int          `toml:"retry" mapstructure:"retry"`
-	NoCleanCache bool         `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"`
-	Users        []userConfig `toml:"users" mapstructure:"users" json:"users"`
+	Workers      int  `toml:"workers" mapstructure:"workers"`
+	Retry        int  `toml:"retry" mapstructure:"retry"`
+	NoCleanCache bool `toml:"no_clean_cache" mapstructure:"no_clean_cache" json:"no_clean_cache"`
+	Threads      int  `toml:"threads" mapstructure:"threads" json:"threads"`
+
+	Users []userConfig `toml:"users" mapstructure:"users" json:"users"`
 
 	Temp     tempConfig      `toml:"temp" mapstructure:"temp"`
 	Log      logConfig       `toml:"log" mapstructure:"log"`
@@ -67,6 +69,7 @@ func Init() error {
 
 	viper.SetDefault("workers", 3)
 	viper.SetDefault("retry", 3)
+	viper.SetDefault("threads", 4)
 
 	viper.SetDefault("telegram.app_id", 1025907)
 	viper.SetDefault("telegram.app_hash", "452b0359b988148995f22ff0f4229750")

diff --git a/core/utils.go b/core/utils.go
@@ -162,9 +162,9 @@ func fixTaskFileExt(task *types.Task, localFilePath string) {
 func getTaskThreads(fileSize int64) int {
 	threads := 1
 	if fileSize > 1024*1024*100 {
-		threads = 4
+		threads = config.Cfg.Threads
 	} else if fileSize > 1024*1024*50 {
-		threads = 2
+		threads = config.Cfg.Threads / 2
 	}
 	return threads
 }