Skip to content

Commit

Permalink
feat: add retry configuration for file upload and validate config values
Browse files Browse the repository at this point in the history
  • Loading branch information
krau committed Jan 19, 2025
1 parent e1bc80a commit 79fba91
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 20 deletions.
7 changes: 5 additions & 2 deletions config/viper.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)

Expand All @@ -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)
}
}
42 changes: 24 additions & 18 deletions core/reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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) {
Expand Down

0 comments on commit 79fba91

Please sign in to comment.