feat: 添加块大小配置项

krau · Oct 12, 2024 · edef778 · edef778
1 parent 0bd9b77
commit edef778
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/config/viper.go b/config/viper.go
@@ -8,8 +8,9 @@ import (
 )
 
 type Config struct {
-	Threads int `toml:"threads" mapstructure:"threads"`
-	Workers int `toml:"workers" mapstructure:"workers"`
+	Threads   int   `toml:"threads" mapstructure:"threads"`
+	Workers   int   `toml:"workers" mapstructure:"workers"`
+	ChunkSize int32 `toml:"chunk_size" mapstructure:"chunk_size"`
 
 	Temp     tempConfig     `toml:"temp" mapstructure:"temp"`
 	Log      logConfig      `toml:"log" mapstructure:"log"`
@@ -77,6 +78,7 @@ func Init() {
 
 	viper.SetDefault("threads", 3)
 	viper.SetDefault("workers", 3)
+	viper.SetDefault("chunk_size", 50*1024*1024)
 
 	viper.SetDefault("temp.base_path", "cache/")
 	viper.SetDefault("temp.cache_ttl", 3600)

diff --git a/core/core.go b/core/core.go
@@ -26,8 +26,9 @@ func processPendingTask(task types.Task) error {
 	logger.L.Debugf("Start downloading file: %s", task.FileName)
 	bot.Client.EditMessage(task.ChatID, task.ReplyMessageID, "正在下载文件...")
 	dest, err := message.Download(&telegram.DownloadOptions{
-		FileName: common.GetCacheFilePath(task.FileName),
-		Threads:  config.Cfg.Threads,
+		FileName:  common.GetCacheFilePath(task.FileName),
+		Threads:   config.Cfg.Threads,
+		ChunkSize: config.Cfg.ChunkSize,
 		// ProgressCallback: func(totalBytes, downloadedBytes int64) {},
 	})
 	if err != nil {