feat: add retry configuration for file upload and validate config values

This commit is contained in:
krau
2025-01-19 14:25:56 +08:00
parent e1bc80ab7d
commit 79fba918cf
2 changed files with 29 additions and 20 deletions

View File

@@ -82,6 +82,7 @@ func Init() {
viper.SetConfigType("toml") viper.SetConfigType("toml")
viper.SetDefault("workers", 3) viper.SetDefault("workers", 3)
viper.SetDefault("retry", 3)
viper.SetDefault("temp.base_path", "cache/") viper.SetDefault("temp.base_path", "cache/")
viper.SetDefault("temp.cache_ttl", 3600) viper.SetDefault("temp.cache_ttl", 3600)
@@ -92,8 +93,6 @@ func Init() {
viper.SetDefault("db.path", "data/saveany.db") 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.base_path", "/")
viper.SetDefault("storage.alist.token_exp", 3600) viper.SetDefault("storage.alist.token_exp", 3600)
@@ -107,4 +106,8 @@ func Init() {
fmt.Println("Error unmarshalling config file, ", err) fmt.Println("Error unmarshalling config file, ", err)
os.Exit(1) os.Exit(1)
} }
if Cfg.Workers < 1 || Cfg.Retry < 1 {
fmt.Println("Invalid workers or retry value")
os.Exit(1)
}
} }

View File

@@ -4,9 +4,11 @@ import (
"context" "context"
"fmt" "fmt"
"io" "io"
"strings"
"github.com/celestix/gotgproto" "github.com/celestix/gotgproto"
"github.com/gotd/td/tg" "github.com/gotd/td/tg"
"github.com/krau/SaveAny-Bot/config"
) )
type telegramReader struct { type telegramReader struct {
@@ -90,25 +92,29 @@ func NewTelegramReader(
} }
func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) { func (r *telegramReader) chunk(offset int64, limit int64) ([]byte, error) {
var lastError error
req := &tg.UploadGetFileRequest{ for i := 0; i < config.Cfg.Retry; i++ {
Offset: offset, req := &tg.UploadGetFileRequest{
Limit: int(limit), Offset: offset,
Location: r.location, Limit: int(limit),
} Location: r.location,
}
res, err := r.client.API().UploadGetFile(r.ctx, req) res, err := r.client.API().UploadGetFile(r.ctx, req)
if err != nil {
if err != nil { if strings.Contains(err.Error(), tg.ErrTimeout) {
return nil, err lastError = err
} continue
}
switch result := res.(type) { return nil, err
case *tg.UploadFile: }
return result.Bytes, nil switch result := res.(type) {
default: case *tg.UploadFile:
return nil, fmt.Errorf("unexpected type %T", r) return result.Bytes, nil
default:
return nil, fmt.Errorf("unexpected type %T", r)
}
} }
return nil, lastError
} }
func (r *telegramReader) partStream() func() ([]byte, error) { func (r *telegramReader) partStream() func() ([]byte, error) {