diff --git a/config/viper.go b/config/viper.go index 41d9b70..42b5622 100644 --- a/config/viper.go +++ b/config/viper.go @@ -82,6 +82,7 @@ func Init() { viper.SetConfigType("toml") viper.SetDefault("workers", 3) + viper.SetDefault("retry", 3) viper.SetDefault("temp.base_path", "cache/") viper.SetDefault("temp.cache_ttl", 3600) @@ -92,8 +93,6 @@ func Init() { viper.SetDefault("db.path", "data/saveany.db") - viper.SetDefault("telegram.api", "https://api.telegram.org") - viper.SetDefault("storage.alist.base_path", "/") viper.SetDefault("storage.alist.token_exp", 3600) @@ -107,4 +106,8 @@ func Init() { fmt.Println("Error unmarshalling config file, ", err) os.Exit(1) } + if Cfg.Workers < 1 || Cfg.Retry < 1 { + fmt.Println("Invalid workers or retry value") + os.Exit(1) + } } diff --git a/core/reader.go b/core/reader.go index ebef27a..14a5324 100644 --- a/core/reader.go +++ b/core/reader.go @@ -4,9 +4,11 @@ import ( "context" "fmt" "io" + "strings" "github.com/celestix/gotgproto" "github.com/gotd/td/tg" + "github.com/krau/SaveAny-Bot/config" ) type telegramReader struct { @@ -90,25 +92,29 @@ func NewTelegramReader( } func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) { - - req := &tg.UploadGetFileRequest{ - Offset: offset, - Limit: int(limit), - Location: r.location, - } - - res, err := r.client.API().UploadGetFile(r.ctx, req) - - if err != nil { - return nil, err - } - - switch result := res.(type) { - case *tg.UploadFile: - return result.Bytes, nil - default: - return nil, fmt.Errorf("unexpected type %T", r) + var lastError error + for i := 0; i < config.Cfg.Retry; i++ { + req := &tg.UploadGetFileRequest{ + Offset: offset, + Limit: int(limit), + Location: r.location, + } + res, err := r.client.API().UploadGetFile(r.ctx, req) + if err != nil { + if strings.Contains(err.Error(), tg.ErrTimeout) { + lastError = err + continue + } + return nil, err + } + switch result := res.(type) { + case *tg.UploadFile: + return result.Bytes, nil + default: + return nil, fmt.Errorf("unexpected type %T", r) + } } + return nil, lastError } func (r *telegramReader) partStream() func() ([]byte, error) {