From 8d0851f37a8183ece89b320d64b869b985bd1145 Mon Sep 17 00:00:00 2001
From: krau <71133316+krau@users.noreply.github.com>
Date: Tue, 7 Jan 2025 10:30:33 +0800
Subject: [PATCH] feat: add retry configuration for failed  save task

---
 config.example.toml |  1 +
 config/viper.go     |  1 +
 core/core.go        | 18 +++++++++++++++---
 3 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/config.example.toml b/config.example.toml
index a2e6d59..ac0540a 100644
--- 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
index 0a5ab85..41d9b70 100644
--- 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
index 497450f..e51f0b7 100644
--- 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
 }