feat: add retry configuration for failed save task

krau · Jan 7, 2025 · 8d0851f · 8d0851f
1 parent 4b5cabc
commit 8d0851f
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 3 deletions.
diff --git a/config.example.toml b/config.example.toml
@@ -1,4 +1,5 @@
 workers = 4 # 同时下载文件数
+retry = 3   # 下载失败重试次数
 
 [telegram]
 token = ""                                    # Bot Token

diff --git a/config/viper.go b/config/viper.go
@@ -9,6 +9,7 @@ import (
 
 type Config struct {
 	Workers int `toml:"workers" mapstructure:"workers"`
+	Retry   int `toml:"retry" mapstructure:"retry"` // Retry times for failed tasks
 
 	Temp     tempConfig     `toml:"temp" mapstructure:"temp"`
 	Log      logConfig      `toml:"log" mapstructure:"log"`

diff --git a/core/core.go b/core/core.go
@@ -100,9 +100,21 @@ func processPendingTask(task *types.Task) error {
 		Message: "下载完成, 正在转存文件...",
 		ID:      task.ReplyMessageID,
 	})
-
-	if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil {
-		return fmt.Errorf("Failed to save file: %w", err)
+	if config.Cfg.Retry <= 0 {
+		if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil {
+			return fmt.Errorf("Failed to save file: %w", err)
+		}
+	} else {
+		for i := 0; i < config.Cfg.Retry; i++ {
+			if err := storage.Save(task.Storage, task.Ctx, dest.Name(), task.StoragePath); err != nil {
+				logger.L.Errorf("Failed to save file: %s, retrying...", err)
+				if i == config.Cfg.Retry-1 {
+					return fmt.Errorf("Failed to save file: %w", err)
+				}
+			} else {
+				break
+			}
+		}
 	}
 	return nil
 }