Skip to content

Commit

Permalink
feat: configurable download threads and workers
Browse files Browse the repository at this point in the history
  • Loading branch information
krau committed Oct 12, 2024
1 parent b5af39c commit d5f2144
Show file tree
Hide file tree
Showing 2 changed files with 9 additions and 2 deletions.
6 changes: 6 additions & 0 deletions config/viper.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down Expand Up @@ -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)

Expand Down
5 changes: 3 additions & 2 deletions core/core.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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)
}
}

0 comments on commit d5f2144

Please sign in to comment.