mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-08-14 17:03:57 +08:00
feat: show upload progress for batch file tasks (#228)
* feat: show upload progress for file tasks * feat: show upload progress for batch file tasks * style: distinguish download and upload phases * style: mark successful task completion * fix: preserve storage save error context * fix: serialize single-file progress updates * fix: stabilize batch upload progress reporting Correct batch upload totals, completion state, actual file sizes, and synchronized progress snapshots. Add regression coverage for concurrent updates and phase transitions. * feat: show per-file transfer progress Format single-file and batch download and upload states with Telegram entities, blockquotes, speeds, transferred sizes, progress bars, concise counters, and accurate confirmation handling. Cover upload retries plus final, error, and cancellation messages with regression tests. * chore: remove transfer progress tests * test: restore critical transfer progress coverage * test: cover interleaved batch transfers * fix: declare progress styles in locale templates
This commit is contained in:
@@ -38,6 +38,31 @@ type StorageBatchSaver interface {
|
||||
SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error
|
||||
}
|
||||
|
||||
// StorageBatchProgressSaver reports confirmed upload progress for each item in
|
||||
// a logical batch. The item index matches the items slice passed to
|
||||
// SaveBatchWithProgress.
|
||||
type StorageBatchProgressSaver interface {
|
||||
StorageBatchSaver
|
||||
SaveBatchWithProgress(
|
||||
ctx context.Context,
|
||||
items []storagetypes.BatchItem,
|
||||
onProgress func(index int, uploaded, total int64),
|
||||
) error
|
||||
}
|
||||
|
||||
// StorageProgressSaver reports bytes after the backend has accepted them for
|
||||
// upload. Backends with native progress support should implement this instead
|
||||
// of relying on progress inferred from reads of the input stream.
|
||||
type StorageProgressSaver interface {
|
||||
Storage
|
||||
SaveWithProgress(
|
||||
ctx context.Context,
|
||||
reader io.Reader,
|
||||
storagePath string,
|
||||
onProgress func(uploaded, total int64),
|
||||
) error
|
||||
}
|
||||
|
||||
// StorageListable 表示支持列举目录内容的存储
|
||||
type StorageListable interface {
|
||||
Storage
|
||||
@@ -52,6 +77,9 @@ type StorageReadable interface {
|
||||
|
||||
var Storages = make(map[string]Storage)
|
||||
|
||||
var _ StorageProgressSaver = (*telegram.Telegram)(nil)
|
||||
var _ StorageBatchProgressSaver = (*telegram.Telegram)(nil)
|
||||
|
||||
type StorageConstructor func() Storage
|
||||
|
||||
var storageConstructors = map[storenum.StorageType]StorageConstructor{
|
||||
|
||||
60
storage/telegram/progress.go
Normal file
60
storage/telegram/progress.go
Normal file
@@ -0,0 +1,60 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
"github.com/gotd/td/telegram/uploader"
|
||||
)
|
||||
|
||||
var _ uploader.Progress = (*uploadProgress)(nil)
|
||||
|
||||
type uploadProgress struct {
|
||||
mu sync.Mutex
|
||||
onProgress func(uploaded, total int64)
|
||||
total int64
|
||||
uploaded int64
|
||||
byID map[int64]int64
|
||||
}
|
||||
|
||||
func newUploadProgress(total int64, onProgress func(uploaded, total int64)) *uploadProgress {
|
||||
return &uploadProgress{
|
||||
onProgress: onProgress,
|
||||
total: total,
|
||||
byID: make(map[int64]int64),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *uploadProgress) Chunk(ctx context.Context, state uploader.ProgressState) error {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
p.mu.Lock()
|
||||
previous := p.byID[state.ID]
|
||||
if state.Uploaded <= previous {
|
||||
p.mu.Unlock()
|
||||
return nil
|
||||
}
|
||||
p.byID[state.ID] = state.Uploaded
|
||||
p.uploaded += state.Uploaded - previous
|
||||
uploaded := p.uploaded
|
||||
total := p.total
|
||||
if total <= 0 {
|
||||
total = state.Total
|
||||
}
|
||||
p.mu.Unlock()
|
||||
|
||||
if p.onProgress != nil {
|
||||
p.onProgress(uploaded, total)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (p *uploadProgress) reset(total int64) {
|
||||
p.mu.Lock()
|
||||
p.total = total
|
||||
p.uploaded = 0
|
||||
p.byID = make(map[int64]int64)
|
||||
p.mu.Unlock()
|
||||
}
|
||||
80
storage/telegram/progress_test.go
Normal file
80
storage/telegram/progress_test.go
Normal file
@@ -0,0 +1,80 @@
|
||||
package telegram
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"github.com/gotd/td/telegram/uploader"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
)
|
||||
|
||||
func TestUploadProgressAggregatesUploaderParts(t *testing.T) {
|
||||
type update struct {
|
||||
uploaded int64
|
||||
total int64
|
||||
}
|
||||
var updates []update
|
||||
progress := newUploadProgress(100, func(uploaded, total int64) {
|
||||
updates = append(updates, update{uploaded: uploaded, total: total})
|
||||
})
|
||||
|
||||
states := []uploader.ProgressState{
|
||||
{ID: 1, Uploaded: 20, Total: 60},
|
||||
{ID: 1, Uploaded: 20, Total: 60},
|
||||
{ID: 1, Uploaded: 60, Total: 60},
|
||||
{ID: 2, Uploaded: 10, Total: 40},
|
||||
{ID: 2, Uploaded: 40, Total: 40},
|
||||
}
|
||||
for _, state := range states {
|
||||
if err := progress.Chunk(context.Background(), state); err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
want := []update{{20, 100}, {60, 100}, {70, 100}, {100, 100}}
|
||||
if len(updates) != len(want) {
|
||||
t.Fatalf("got %d updates, want %d", len(updates), len(want))
|
||||
}
|
||||
for i := range want {
|
||||
if updates[i] != want[i] {
|
||||
t.Fatalf("update %d = %+v, want %+v", i, updates[i], want[i])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadProgressResetForSplitFiles(t *testing.T) {
|
||||
var uploaded, total int64
|
||||
progress := newUploadProgress(100, func(current, size int64) {
|
||||
uploaded, total = current, size
|
||||
})
|
||||
if err := progress.Chunk(context.Background(), uploader.ProgressState{ID: 1, Uploaded: 100, Total: 100}); err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
|
||||
progress.reset(120)
|
||||
if err := progress.Chunk(context.Background(), uploader.ProgressState{ID: 2, Uploaded: 30, Total: 60}); err != nil {
|
||||
t.Fatalf("Chunk() after reset failed: %v", err)
|
||||
}
|
||||
if uploaded != 30 || total != 120 {
|
||||
t.Fatalf("progress after reset = %d/%d, want 30/120", uploaded, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchItemUploadProgressPreservesItemIndex(t *testing.T) {
|
||||
var gotIndex int
|
||||
var gotUploaded, gotTotal int64
|
||||
progress := batchItemUploadProgress(batchMediaItem{
|
||||
index: 4,
|
||||
item: storagetypes.BatchItem{Size: 100},
|
||||
}, func(index int, uploaded, total int64) {
|
||||
gotIndex = index
|
||||
gotUploaded = uploaded
|
||||
gotTotal = total
|
||||
})
|
||||
if err := progress.Chunk(context.Background(), uploader.ProgressState{ID: 1, Uploaded: 25, Total: 100}); err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
if gotIndex != 4 || gotUploaded != 25 || gotTotal != 100 {
|
||||
t.Fatalf("batch progress = index %d, %d/%d; want index 4, 25/100", gotIndex, gotUploaded, gotTotal)
|
||||
}
|
||||
}
|
||||
@@ -52,6 +52,7 @@ type preparedMedia struct {
|
||||
|
||||
type batchMediaItem struct {
|
||||
item storagetypes.BatchItem
|
||||
index int
|
||||
chatID int64
|
||||
albumEligible bool
|
||||
useSingleSave bool
|
||||
@@ -87,6 +88,22 @@ func (t *Telegram) Exists(ctx context.Context, storagePath string) bool {
|
||||
}
|
||||
|
||||
func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) error {
|
||||
return t.save(ctx, r, storagePath, nil)
|
||||
}
|
||||
|
||||
// SaveWithProgress saves a file while reporting Telegram-confirmed upload
|
||||
// progress after each uploaded part.
|
||||
func (t *Telegram) SaveWithProgress(
|
||||
ctx context.Context,
|
||||
r io.Reader,
|
||||
storagePath string,
|
||||
onProgress func(uploaded, total int64),
|
||||
) error {
|
||||
size := contentLength(ctx)
|
||||
return t.save(ctx, r, storagePath, newUploadProgress(size, onProgress))
|
||||
}
|
||||
|
||||
func (t *Telegram) save(ctx context.Context, r io.Reader, storagePath string, progress *uploadProgress) error {
|
||||
storagePath = path.Clean(storagePath)
|
||||
captionOverride := sourceCaptionOverride(ctx)
|
||||
tctx := tgutil.ExtFromContext(ctx)
|
||||
@@ -114,7 +131,7 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
||||
}
|
||||
}
|
||||
}
|
||||
upler := t.newUploader(tctx, size)
|
||||
upler := t.newUploader(tctx, size, progress)
|
||||
peer := tryGetInputPeer(tctx, chatID)
|
||||
if peer == nil || peer.Zero() {
|
||||
return fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
||||
@@ -150,7 +167,7 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
||||
for _, part := range parts {
|
||||
log.FromContext(ctx).Infof("Prepared lossless video part %s (%d bytes)", part.Name, part.Size)
|
||||
}
|
||||
return t.uploadLosslessVideoParts(ctx, tctx, storagePath, parts, captionOverride)
|
||||
return t.uploadLosslessVideoParts(ctx, tctx, storagePath, parts, captionOverride, progress)
|
||||
}
|
||||
if _, seekErr := rs.Seek(0, io.SeekStart); seekErr != nil {
|
||||
return fmt.Errorf("failed to seek large video before ZIP fallback: %w", seekErr)
|
||||
@@ -158,13 +175,13 @@ func (t *Telegram) Save(ctx context.Context, r io.Reader, storagePath string) er
|
||||
}
|
||||
}
|
||||
}
|
||||
return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize)
|
||||
return t.splitUpload(tctx, r, filename, upler, peer, size, splitSize, progress)
|
||||
}
|
||||
|
||||
if err := t.limiter.Wait(ctx); err != nil {
|
||||
return fmt.Errorf("rate limit failed: %w", err)
|
||||
}
|
||||
prepared, err := t.prepareMedia(ctx, tctx, r, storagePath, size, nil)
|
||||
prepared, err := t.prepareMedia(ctx, tctx, r, storagePath, size, nil, progress)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -229,10 +246,14 @@ func (t *Telegram) target(tctx *ext.Context, storagePath string) (string, int64)
|
||||
return filename, chatID
|
||||
}
|
||||
|
||||
func (t *Telegram) newUploader(tctx *ext.Context, size int64) *uploader.Uploader {
|
||||
return uploader.NewUploader(tctx.Raw).
|
||||
func (t *Telegram) newUploader(tctx *ext.Context, size int64, progress *uploadProgress) *uploader.Uploader {
|
||||
upler := uploader.NewUploader(tctx.Raw).
|
||||
WithPartSize(tglimit.MaxUploadPartSize).
|
||||
WithThreads(dlutil.BestThreads(size, config.C().Threads))
|
||||
if progress != nil {
|
||||
upler = upler.WithProgress(progress)
|
||||
}
|
||||
return upler
|
||||
}
|
||||
|
||||
func mediaCaption(filename string, override *string) []message.StyledTextOption {
|
||||
@@ -245,10 +266,18 @@ func mediaCaption(filename string, override *string) []message.StyledTextOption
|
||||
return []message.StyledTextOption{styling.Plain(*override)}
|
||||
}
|
||||
|
||||
func (t *Telegram) prepareMedia(ctx context.Context, tctx *ext.Context, r io.Reader, storagePath string, size int64, captionOverride *string) (*preparedMedia, error) {
|
||||
func (t *Telegram) prepareMedia(
|
||||
ctx context.Context,
|
||||
tctx *ext.Context,
|
||||
r io.Reader,
|
||||
storagePath string,
|
||||
size int64,
|
||||
captionOverride *string,
|
||||
progress *uploadProgress,
|
||||
) (*preparedMedia, error) {
|
||||
storagePath = path.Clean(storagePath)
|
||||
filename, chatID := t.target(tctx, storagePath)
|
||||
upler := t.newUploader(tctx, size)
|
||||
upler := t.newUploader(tctx, size, progress)
|
||||
peer := tryGetInputPeer(tctx, chatID)
|
||||
if peer == nil || peer.Zero() {
|
||||
return nil, fmt.Errorf("failed to get input peer for chat ID %d", chatID)
|
||||
@@ -342,21 +371,40 @@ func (t *Telegram) prepareMedia(ctx context.Context, tctx *ext.Context, r io.Rea
|
||||
|
||||
// SaveBatch preserves each source photo/video group as a Telegram album.
|
||||
func (t *Telegram) SaveBatch(ctx context.Context, items []storagetypes.BatchItem) error {
|
||||
return t.saveBatch(ctx, items, nil)
|
||||
}
|
||||
|
||||
// SaveBatchWithProgress preserves source media groups while reporting native
|
||||
// Telegram upload progress for each input item.
|
||||
func (t *Telegram) SaveBatchWithProgress(
|
||||
ctx context.Context,
|
||||
items []storagetypes.BatchItem,
|
||||
onProgress func(index int, uploaded, total int64),
|
||||
) error {
|
||||
return t.saveBatch(ctx, items, onProgress)
|
||||
}
|
||||
|
||||
func (t *Telegram) saveBatch(
|
||||
ctx context.Context,
|
||||
items []storagetypes.BatchItem,
|
||||
onProgress func(index int, uploaded, total int64),
|
||||
) error {
|
||||
tctx := tgutil.ExtFromContext(ctx)
|
||||
if tctx == nil {
|
||||
return fmt.Errorf("failed to get telegram context")
|
||||
}
|
||||
|
||||
inspected := make([]batchMediaItem, 0, len(items))
|
||||
for _, item := range items {
|
||||
for index, item := range items {
|
||||
mediaItem, err := t.inspectBatchItem(tctx, item)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
mediaItem.index = index
|
||||
inspected = append(inspected, mediaItem)
|
||||
}
|
||||
for _, group := range planMediaGroups(inspected) {
|
||||
if err := t.saveMediaGroup(ctx, tctx, group); err != nil {
|
||||
if err := t.saveMediaGroup(ctx, tctx, group, onProgress); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -412,10 +460,28 @@ func planMediaGroups(items []batchMediaItem) [][]batchMediaItem {
|
||||
return groups
|
||||
}
|
||||
|
||||
func (t *Telegram) saveMediaGroup(ctx context.Context, tctx *ext.Context, group []batchMediaItem) error {
|
||||
func batchItemUploadProgress(
|
||||
mediaItem batchMediaItem,
|
||||
onProgress func(index int, uploaded, total int64),
|
||||
) *uploadProgress {
|
||||
if onProgress == nil {
|
||||
return nil
|
||||
}
|
||||
return newUploadProgress(mediaItem.item.Size, func(uploaded, total int64) {
|
||||
onProgress(mediaItem.index, uploaded, total)
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Telegram) saveMediaGroup(
|
||||
ctx context.Context,
|
||||
tctx *ext.Context,
|
||||
group []batchMediaItem,
|
||||
onProgress func(index int, uploaded, total int64),
|
||||
) error {
|
||||
return retry.Retry(func() error {
|
||||
if len(group) == 1 && group[0].useSingleSave {
|
||||
item := group[0].item
|
||||
mediaItem := group[0]
|
||||
item := mediaItem.item
|
||||
if _, err := item.Reader.Seek(0, io.SeekStart); err != nil {
|
||||
return fmt.Errorf("failed to seek batch item: %w", err)
|
||||
}
|
||||
@@ -423,7 +489,12 @@ func (t *Telegram) saveMediaGroup(ctx context.Context, tctx *ext.Context, group
|
||||
if item.PreserveCaption {
|
||||
itemCtx = storagetypes.WithSourceCaption(itemCtx, item.Caption)
|
||||
}
|
||||
return t.Save(itemCtx, item.Reader, item.StoragePath)
|
||||
if onProgress == nil {
|
||||
return t.Save(itemCtx, item.Reader, item.StoragePath)
|
||||
}
|
||||
return t.SaveWithProgress(itemCtx, item.Reader, item.StoragePath, func(uploaded, total int64) {
|
||||
onProgress(mediaItem.index, uploaded, total)
|
||||
})
|
||||
}
|
||||
if err := t.limiter.Wait(ctx); err != nil {
|
||||
return fmt.Errorf("rate limit failed: %w", err)
|
||||
@@ -439,7 +510,8 @@ func (t *Telegram) saveMediaGroup(ctx context.Context, tctx *ext.Context, group
|
||||
if item.PreserveCaption {
|
||||
captionOverride = &item.Caption
|
||||
}
|
||||
media, err := t.prepareMedia(ctx, tctx, item.Reader, item.StoragePath, item.Size, captionOverride)
|
||||
progress := batchItemUploadProgress(mediaItem, onProgress)
|
||||
media, err := t.prepareMedia(ctx, tctx, item.Reader, item.StoragePath, item.Size, captionOverride, progress)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -466,7 +538,15 @@ func (t *Telegram) CannotStream() string {
|
||||
return "Telegram storage must use a ReaderSeeker"
|
||||
}
|
||||
|
||||
func (t *Telegram) splitUpload(ctx *ext.Context, r io.Reader, filename string, upler *uploader.Uploader, peer tg.InputPeerClass, fileSize, splitSize int64) error {
|
||||
func (t *Telegram) splitUpload(
|
||||
ctx *ext.Context,
|
||||
r io.Reader,
|
||||
filename string,
|
||||
upler *uploader.Uploader,
|
||||
peer tg.InputPeerClass,
|
||||
fileSize, splitSize int64,
|
||||
progress *uploadProgress,
|
||||
) error {
|
||||
tempId := xid.New().String()
|
||||
outputBase := filepath.Join(config.C().Temp.BasePath, tempId, strings.Split(filename, ".")[0])
|
||||
defer func() {
|
||||
@@ -483,6 +563,17 @@ func (t *Telegram) splitUpload(ctx *ext.Context, r io.Reader, filename string, u
|
||||
return fmt.Errorf("failed to glob split files: %w", err)
|
||||
}
|
||||
inputFiles := make([]tg.InputFileClass, 0, len(matched))
|
||||
if progress != nil {
|
||||
var uploadSize int64
|
||||
for _, partPath := range matched {
|
||||
partInfo, err := os.Stat(partPath)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to stat split part %s: %w", partPath, err)
|
||||
}
|
||||
uploadSize += partInfo.Size()
|
||||
}
|
||||
progress.reset(uploadSize)
|
||||
}
|
||||
for _, partPath := range matched {
|
||||
// 串行上传, 不然容易被tg风控
|
||||
err = func() error {
|
||||
|
||||
@@ -293,6 +293,7 @@ func (t *Telegram) uploadLosslessVideoParts(
|
||||
storagePath string,
|
||||
parts []losslessVideoPart,
|
||||
sourceCaption *string,
|
||||
progress *uploadProgress,
|
||||
) error {
|
||||
if len(parts) == 0 {
|
||||
return fmt.Errorf("no lossless video parts to upload")
|
||||
@@ -304,6 +305,7 @@ func (t *Telegram) uploadLosslessVideoParts(
|
||||
maxLosslessVideoParts,
|
||||
)
|
||||
}
|
||||
resetLosslessVideoUploadProgress(progress, parts)
|
||||
|
||||
prepared := make([]preparedMedia, 0, len(parts))
|
||||
for index, part := range parts {
|
||||
@@ -318,6 +320,7 @@ func (t *Telegram) uploadLosslessVideoParts(
|
||||
partStoragePath(storagePath, part.Name),
|
||||
part.Size,
|
||||
videoPartCaption(sourceCaption, index),
|
||||
progress,
|
||||
)
|
||||
closeErr := partFile.Close()
|
||||
if prepareErr != nil {
|
||||
@@ -346,6 +349,17 @@ func (t *Telegram) uploadLosslessVideoParts(
|
||||
return nil
|
||||
}
|
||||
|
||||
func resetLosslessVideoUploadProgress(progress *uploadProgress, parts []losslessVideoPart) {
|
||||
if progress == nil {
|
||||
return
|
||||
}
|
||||
var total int64
|
||||
for _, part := range parts {
|
||||
total += part.Size
|
||||
}
|
||||
progress.reset(total)
|
||||
}
|
||||
|
||||
func videoPartCaption(sourceCaption *string, index int) *string {
|
||||
if sourceCaption == nil {
|
||||
return nil
|
||||
|
||||
@@ -7,6 +7,8 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/gotd/td/telegram/uploader"
|
||||
)
|
||||
|
||||
func TestInitialSegmentDuration(t *testing.T) {
|
||||
@@ -76,6 +78,25 @@ func TestVideoPartCaption(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestResetLosslessVideoUploadProgressUsesCombinedPartSize(t *testing.T) {
|
||||
var uploaded, total int64
|
||||
progress := newUploadProgress(999, func(current, size int64) {
|
||||
uploaded = current
|
||||
total = size
|
||||
})
|
||||
resetLosslessVideoUploadProgress(progress, []losslessVideoPart{
|
||||
{Size: 100},
|
||||
{Size: 250},
|
||||
})
|
||||
|
||||
if err := progress.Chunk(t.Context(), uploader.ProgressState{ID: 1, Uploaded: 50, Total: 100}); err != nil {
|
||||
t.Fatalf("Chunk() failed: %v", err)
|
||||
}
|
||||
if uploaded != 50 || total != 350 {
|
||||
t.Fatalf("lossless video progress = %d/%d, want 50/350", uploaded, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplitLosslessVideoRetriesOversizedPart(t *testing.T) {
|
||||
tempDir := t.TempDir()
|
||||
inputPath := filepath.Join(tempDir, "source.mov")
|
||||
|
||||
Reference in New Issue
Block a user