diff --git a/config/viper.go b/config/viper.go index b66c7e8..6c0754d 100644 --- a/config/viper.go +++ b/config/viper.go @@ -8,6 +8,9 @@ import ( ) type Config struct { + Threads int `toml:"threads" mapstructure:"threads"` + Workers int `toml:"workers" mapstructure:"workers"` + Temp tempConfig `toml:"temp" mapstructure:"temp"` Log logConfig `toml:"log" mapstructure:"log"` DB dbConfig `toml:"db" mapstructure:"db"` @@ -72,6 +75,9 @@ func Init() { viper.AddConfigPath(".") viper.SetConfigType("toml") + viper.SetDefault("threads", 3) + viper.SetDefault("workers", 3) + viper.SetDefault("temp.base_path", "cache/") viper.SetDefault("temp.cache_ttl", 3600) diff --git a/core/core.go b/core/core.go index 3c19dab..fd2a09f 100644 --- a/core/core.go +++ b/core/core.go @@ -24,6 +24,7 @@ func processPendingTask(task types.Task) error { bot.Client.EditMessage(task.ChatID, task.ReplyMessageID, "正在下载文件...") dest, err := message.Download(&telegram.DownloadOptions{ FileName: common.GetCacheFilePath(task.FileName), + Threads: config.Cfg.Threads, // ProgressCallback: func(totalBytes, downloadedBytes int64) {}, }) if err != nil { @@ -90,8 +91,8 @@ func worker(queue *queue.TaskQueue, semaphore chan struct{}) { func Run() { logger.L.Info("Start processing tasks...") - semaphore := make(chan struct{}, 3) - for i := 0; i < 3; i++ { + semaphore := make(chan struct{}, config.Cfg.Workers) + for i := 0; i < config.Cfg.Workers; i++ { go worker(queue.Queue, semaphore) } }