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:
Haopeng Huo
@@ -6,6 +6,7 @@ import (
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/duke-git/lancet/v2/retry"
|
||||
@@ -60,6 +61,7 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
} else {
|
||||
logger.Info("Batch file task completed successfully")
|
||||
}
|
||||
t.finishItems(err)
|
||||
t.Progress.OnDone(ctx, t, err)
|
||||
return err
|
||||
}
|
||||
@@ -98,7 +100,7 @@ func (t *Task) processElements(ctx context.Context, elems []*TaskElement) error
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for _, elem := range elems {
|
||||
eg.Go(func() error {
|
||||
if err := t.markProcessing(elem); err != nil {
|
||||
if err := t.markProcessing(ctx, elem); err != nil {
|
||||
return err
|
||||
}
|
||||
defer t.unmarkProcessing(elem.ID)
|
||||
@@ -121,7 +123,7 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
eg.SetLimit(config.C().Workers)
|
||||
for _, elem := range group.elems {
|
||||
eg.Go(func() error {
|
||||
if err := t.markProcessing(elem); err != nil {
|
||||
if err := t.markProcessing(ctx, elem); err != nil {
|
||||
return err
|
||||
}
|
||||
defer t.unmarkProcessing(elem.ID)
|
||||
@@ -144,11 +146,13 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
for _, elem := range group.elems {
|
||||
file, err := os.Open(elem.localPath)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
return fmt.Errorf("failed to open cache file: %w", err)
|
||||
}
|
||||
stat, err := file.Stat()
|
||||
if err != nil {
|
||||
file.Close()
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
return fmt.Errorf("failed to get cache file stat: %w", err)
|
||||
}
|
||||
openFiles = append(openFiles, file)
|
||||
@@ -161,19 +165,67 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error {
|
||||
PreserveCaption: elem.preserveCaption,
|
||||
})
|
||||
}
|
||||
for index, item := range items {
|
||||
t.recordDownloadComplete(group.elems[index].ID, item.Size)
|
||||
}
|
||||
return t.saveBatchItems(ctx, group, items)
|
||||
}
|
||||
|
||||
func (t *Task) saveBatchItems(ctx context.Context, group executionGroup, items []storagetypes.BatchItem) error {
|
||||
t.startUpload(ctx)
|
||||
if progressSaver, ok := group.batchSaver.(storage.StorageBatchProgressSaver); ok {
|
||||
err := progressSaver.SaveBatchWithProgress(ctx, items, func(index int, uploaded, total int64) {
|
||||
if index < 0 || index >= len(group.elems) {
|
||||
return
|
||||
}
|
||||
t.uploadCallback(ctx, group.elems[index].ID)(uploaded, total)
|
||||
})
|
||||
if err != nil {
|
||||
for _, elem := range group.elems {
|
||||
t.markItemFailed(elem.ID, FailureStageBatchUpload, err)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to save batch: %w", err)
|
||||
}
|
||||
for index, elem := range group.elems {
|
||||
t.uploadCallback(ctx, elem.ID)(items[index].Size, items[index].Size)
|
||||
t.markItemCompleted(elem.ID)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return nil
|
||||
}
|
||||
for i := range items {
|
||||
items[i].Reader = ioutil.NewProgressReader(
|
||||
items[i].Reader,
|
||||
items[i].Size,
|
||||
t.uploadCallback(ctx, group.elems[i].ID),
|
||||
)
|
||||
}
|
||||
if err := group.batchSaver.SaveBatch(ctx, items); err != nil {
|
||||
for _, elem := range group.elems {
|
||||
t.markItemFailed(elem.ID, FailureStageBatchUpload, err)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to save batch: %w", err)
|
||||
}
|
||||
for index, elem := range group.elems {
|
||||
t.uploadCallback(ctx, elem.ID)(items[index].Size, items[index].Size)
|
||||
t.markItemCompleted(elem.ID)
|
||||
}
|
||||
t.notifyStateChange(ctx)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *Task) markProcessing(elem *TaskElement) error {
|
||||
func (t *Task) markProcessing(ctx context.Context, elem *TaskElement) error {
|
||||
t.processingMu.Lock()
|
||||
defer t.processingMu.Unlock()
|
||||
if t.processing[elem.ID] != nil {
|
||||
t.processingMu.Unlock()
|
||||
return fmt.Errorf("element with ID %s is already being processed", elem.ID)
|
||||
}
|
||||
t.processing[elem.ID] = elem
|
||||
t.processingMu.Unlock()
|
||||
t.markItemActive(elem.ID, elem.stream, time.Now())
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -188,9 +240,12 @@ func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||
logger.Info("Starting file download")
|
||||
localFile, err := fsutil.CreateFile(elem.localPath)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to create local file: %w", err)
|
||||
}
|
||||
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
@@ -203,9 +258,13 @@ func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||
_, downloadErr := tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
|
||||
closeErr := localFile.Close()
|
||||
if downloadErr != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageDownload, downloadErr)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to download file: %w", downloadErr)
|
||||
}
|
||||
if closeErr != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, closeErr)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to close cache file: %w", closeErr)
|
||||
}
|
||||
logger.Info("File downloaded successfully")
|
||||
@@ -214,6 +273,8 @@ func (t *Task) downloadElement(ctx context.Context, elem *TaskElement) error {
|
||||
elem.Path += ext
|
||||
}
|
||||
}
|
||||
t.markItemDownloaded(elem.ID)
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -224,9 +285,15 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
defer pr.Close()
|
||||
errg, uploadCtx := errgroup.WithContext(ctx)
|
||||
errg.Go(func() error {
|
||||
return elem.Storage.Save(uploadCtx, pr, elem.Path)
|
||||
err := elem.Storage.Save(uploadCtx, pr, elem.Path)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageUpload, err)
|
||||
t.notifyStateChange(ctx)
|
||||
}
|
||||
return err
|
||||
})
|
||||
wr := ioutil.NewProgressWriter(pw, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
@@ -242,6 +309,8 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
_, err := tdler.NewDownloader(elem.File).Stream(uploadCtx, wr)
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to download file: %v", err)
|
||||
t.markItemFailed(elem.ID, FailureStageDownload, err)
|
||||
t.notifyStateChange(ctx)
|
||||
pw.CloseWithError(err)
|
||||
}
|
||||
return err
|
||||
@@ -249,12 +318,17 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
if err := errg.Wait(); err != nil {
|
||||
return fmt.Errorf("failed to download file in stream mode: %w", err)
|
||||
}
|
||||
t.recordDownloadComplete(elem.ID, 0)
|
||||
t.markItemCompleted(elem.ID)
|
||||
t.notifyStateChange(ctx)
|
||||
logger.Info("File downloaded successfully in stream mode")
|
||||
return nil
|
||||
}
|
||||
logger.Info("Starting file download")
|
||||
localFile, err := fsutil.CreateFile(elem.localPath)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to create local file: %w", err)
|
||||
}
|
||||
defer func() {
|
||||
@@ -263,6 +337,7 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
}
|
||||
}()
|
||||
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||
t.recordItemDownload(elem.ID, int64(n), time.Now())
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
@@ -274,6 +349,8 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
})
|
||||
_, err = tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageDownload, err)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to download file: %w", err)
|
||||
}
|
||||
logger.Info("File downloaded successfully")
|
||||
@@ -286,21 +363,50 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
var fileStat os.FileInfo
|
||||
fileStat, err = os.Stat(elem.localPath)
|
||||
if err != nil {
|
||||
t.markItemFailed(elem.ID, FailureStageCache, err)
|
||||
t.notifyStateChange(ctx)
|
||||
return fmt.Errorf("failed to get file stat: %w", err)
|
||||
}
|
||||
t.recordDownloadComplete(elem.ID, fileStat.Size())
|
||||
vctx := context.WithValue(ctx, ctxkey.ContentLength, fileStat.Size())
|
||||
t.startUpload(vctx)
|
||||
onProgress := t.uploadCallback(vctx, elem.ID)
|
||||
attempt := 0
|
||||
retryLimit := int(config.C().Retry)
|
||||
lastFailureStage := FailureStageUpload
|
||||
err = retry.Retry(func() error {
|
||||
attempt++
|
||||
var file *os.File
|
||||
file, err = os.Open(elem.localPath)
|
||||
if err != nil {
|
||||
lastFailureStage = FailureStageCache
|
||||
t.markItemRetry(elem.ID, lastFailureStage, attempt, retryLimit, err)
|
||||
t.notifyStateChange(vctx)
|
||||
return fmt.Errorf("failed to open cache file: %w", err)
|
||||
}
|
||||
defer file.Close()
|
||||
if err = elem.Storage.Save(vctx, file, elem.Path); err != nil {
|
||||
onProgress(0, fileStat.Size())
|
||||
if progressSaver, ok := elem.Storage.(storage.StorageProgressSaver); ok {
|
||||
err = progressSaver.SaveWithProgress(vctx, file, elem.Path, onProgress)
|
||||
} else {
|
||||
err = elem.Storage.Save(vctx, ioutil.NewProgressReader(file, fileStat.Size(), onProgress), elem.Path)
|
||||
}
|
||||
if err != nil {
|
||||
logger.Errorf("Failed to save file: %s, retrying...", err)
|
||||
lastFailureStage = t.itemFailureStage(elem.ID)
|
||||
t.markItemRetry(elem.ID, lastFailureStage, attempt, retryLimit, err)
|
||||
t.notifyStateChange(vctx)
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}, retry.Context(vctx), retry.RetryTimes(uint(config.C().Retry)))
|
||||
if err == nil {
|
||||
onProgress(fileStat.Size(), fileStat.Size())
|
||||
t.markItemCompleted(elem.ID)
|
||||
t.notifyStateChange(vctx)
|
||||
} else {
|
||||
t.markItemFailed(elem.ID, lastFailureStage, err)
|
||||
t.notifyStateChange(vctx)
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
360
core/tasks/batchtfile/item_progress.go
Normal file
360
core/tasks/batchtfile/item_progress.go
Normal file
@@ -0,0 +1,360 @@
|
||||
package batchtfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
transferSpeedWindow = 5 * time.Second
|
||||
transferSamplePeriod = 250 * time.Millisecond
|
||||
)
|
||||
|
||||
// ItemPhase describes the current lifecycle stage of one batch item.
|
||||
type ItemPhase uint8
|
||||
|
||||
const (
|
||||
ItemPhaseWaiting ItemPhase = iota
|
||||
ItemPhaseDownloading
|
||||
ItemPhaseTransferring
|
||||
ItemPhaseDownloaded
|
||||
ItemPhaseUploading
|
||||
ItemPhaseRetrying
|
||||
ItemPhaseConfirming
|
||||
ItemPhaseCompleted
|
||||
ItemPhaseFailed
|
||||
ItemPhaseStopped
|
||||
)
|
||||
|
||||
// FailureStage identifies the operation that failed for one batch item.
|
||||
type FailureStage uint8
|
||||
|
||||
const (
|
||||
FailureStageNone FailureStage = iota
|
||||
FailureStageDownload
|
||||
FailureStageCache
|
||||
FailureStageUpload
|
||||
FailureStageConfirm
|
||||
FailureStageBatchUpload
|
||||
FailureStageInternal
|
||||
)
|
||||
|
||||
// TaskItemProgress is an immutable progress snapshot for one batch item.
|
||||
type TaskItemProgress struct {
|
||||
Index int
|
||||
ID string
|
||||
Name string
|
||||
Size int64
|
||||
Downloaded int64
|
||||
Uploaded int64
|
||||
DownloadSpeed float64
|
||||
UploadSpeed float64
|
||||
Phase ItemPhase
|
||||
FailureStage FailureStage
|
||||
RetryAttempt int
|
||||
RetryLimit int
|
||||
Error string
|
||||
}
|
||||
|
||||
type transferSample struct {
|
||||
at time.Time
|
||||
bytes int64
|
||||
}
|
||||
|
||||
type transferMeter struct {
|
||||
samples []transferSample
|
||||
latest transferSample
|
||||
hasData bool
|
||||
}
|
||||
|
||||
func (m *transferMeter) record(now time.Time, transferred int64) {
|
||||
if m.hasData && transferred < m.latest.bytes {
|
||||
m.reset()
|
||||
}
|
||||
if m.hasData && now.Before(m.latest.at) {
|
||||
now = m.latest.at
|
||||
}
|
||||
sample := transferSample{at: now, bytes: transferred}
|
||||
m.latest = sample
|
||||
m.hasData = true
|
||||
if len(m.samples) == 0 {
|
||||
m.samples = append(m.samples, sample)
|
||||
return
|
||||
}
|
||||
if now.Sub(m.samples[len(m.samples)-1].at) >= transferSamplePeriod {
|
||||
m.samples = append(m.samples, sample)
|
||||
}
|
||||
cutoff := now.Add(-transferSpeedWindow)
|
||||
for len(m.samples) > 1 && m.samples[0].at.Before(cutoff) {
|
||||
m.samples = m.samples[1:]
|
||||
}
|
||||
}
|
||||
|
||||
func (m *transferMeter) speed() float64 {
|
||||
if !m.hasData || len(m.samples) == 0 {
|
||||
return 0
|
||||
}
|
||||
first := m.samples[0]
|
||||
last := m.latest
|
||||
elapsed := last.at.Sub(first.at).Seconds()
|
||||
if elapsed <= 0 || last.bytes <= first.bytes {
|
||||
return 0
|
||||
}
|
||||
return float64(last.bytes-first.bytes) / elapsed
|
||||
}
|
||||
|
||||
func (m *transferMeter) reset() {
|
||||
m.samples = m.samples[:0]
|
||||
m.latest = transferSample{}
|
||||
m.hasData = false
|
||||
}
|
||||
|
||||
type itemProgressState struct {
|
||||
index int
|
||||
id string
|
||||
name string
|
||||
expectedSize int64
|
||||
actualSize int64
|
||||
downloaded int64
|
||||
uploaded int64
|
||||
phase ItemPhase
|
||||
failureStage FailureStage
|
||||
retryAttempt int
|
||||
retryLimit int
|
||||
err string
|
||||
downloadMeter transferMeter
|
||||
uploadMeter transferMeter
|
||||
}
|
||||
|
||||
func newItemProgressStates(elems []TaskElement) ([]itemProgressState, map[string]int) {
|
||||
states := make([]itemProgressState, 0, len(elems))
|
||||
index := make(map[string]int, len(elems))
|
||||
for i, elem := range elems {
|
||||
name := ""
|
||||
size := int64(0)
|
||||
if elem.File != nil {
|
||||
name = elem.File.Name()
|
||||
size = elem.File.Size()
|
||||
}
|
||||
states = append(states, itemProgressState{
|
||||
index: i + 1,
|
||||
id: elem.ID,
|
||||
name: name,
|
||||
expectedSize: size,
|
||||
phase: ItemPhaseWaiting,
|
||||
})
|
||||
index[elem.ID] = i
|
||||
}
|
||||
return states, index
|
||||
}
|
||||
|
||||
func (t *Task) updateItem(id string, update func(*itemProgressState)) bool {
|
||||
t.itemMu.Lock()
|
||||
defer t.itemMu.Unlock()
|
||||
index, ok := t.itemIndex[id]
|
||||
if !ok || index < 0 || index >= len(t.itemStates) {
|
||||
return false
|
||||
}
|
||||
update(&t.itemStates[index])
|
||||
return true
|
||||
}
|
||||
|
||||
func (t *Task) markItemActive(id string, stream bool, now time.Time) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
if stream {
|
||||
item.phase = ItemPhaseTransferring
|
||||
} else {
|
||||
item.phase = ItemPhaseDownloading
|
||||
}
|
||||
item.failureStage = FailureStageNone
|
||||
item.err = ""
|
||||
item.downloadMeter.record(now, item.downloaded)
|
||||
if stream {
|
||||
item.uploadMeter.record(now, item.uploaded)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) recordItemDownload(id string, n int64, now time.Time) {
|
||||
if n <= 0 {
|
||||
return
|
||||
}
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
item.downloaded += n
|
||||
item.downloadMeter.record(now, item.downloaded)
|
||||
if item.phase == ItemPhaseTransferring {
|
||||
item.uploaded += n
|
||||
item.uploadMeter.record(now, item.uploaded)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) markItemDownloaded(id string) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
item.phase = ItemPhaseDownloaded
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) recordItemDownloaded(id string, actualSize int64) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
if actualSize > 0 {
|
||||
item.actualSize = actualSize
|
||||
}
|
||||
if item.actualSize == 0 {
|
||||
item.actualSize = item.downloaded
|
||||
}
|
||||
if item.phase != ItemPhaseTransferring {
|
||||
item.phase = ItemPhaseDownloaded
|
||||
item.uploadMeter.reset()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) recordItemUpload(id string, uploaded, total int64, now time.Time) bool {
|
||||
becameConfirming := false
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
if uploaded < item.uploaded {
|
||||
if item.phase != ItemPhaseRetrying {
|
||||
return
|
||||
}
|
||||
item.uploadMeter.reset()
|
||||
}
|
||||
if total > 0 {
|
||||
item.actualSize = total
|
||||
}
|
||||
item.uploaded = uploaded
|
||||
item.uploadMeter.record(now, uploaded)
|
||||
if total > 0 && uploaded >= total {
|
||||
becameConfirming = item.phase != ItemPhaseConfirming
|
||||
item.phase = ItemPhaseConfirming
|
||||
return
|
||||
}
|
||||
item.phase = ItemPhaseUploading
|
||||
item.failureStage = FailureStageNone
|
||||
item.err = ""
|
||||
})
|
||||
return becameConfirming
|
||||
}
|
||||
|
||||
func (t *Task) markItemRetry(id string, stage FailureStage, attempt, limit int, err error) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
item.phase = ItemPhaseRetrying
|
||||
item.failureStage = stage
|
||||
item.retryAttempt = attempt
|
||||
item.retryLimit = limit
|
||||
item.err = compactError(err)
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) markItemFailed(id string, stage FailureStage, err error) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
if item.phase == ItemPhaseFailed || item.phase == ItemPhaseCompleted {
|
||||
return
|
||||
}
|
||||
if errors.Is(err, context.Canceled) {
|
||||
item.phase = ItemPhaseStopped
|
||||
return
|
||||
}
|
||||
item.phase = ItemPhaseFailed
|
||||
item.failureStage = stage
|
||||
item.err = compactError(err)
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) markItemCompleted(id string) {
|
||||
t.updateItem(id, func(item *itemProgressState) {
|
||||
item.phase = ItemPhaseCompleted
|
||||
item.failureStage = FailureStageNone
|
||||
item.err = ""
|
||||
item.retryAttempt = 0
|
||||
item.retryLimit = 0
|
||||
if item.actualSize == 0 {
|
||||
item.actualSize = max(item.downloaded, item.uploaded)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) finishItems(err error) {
|
||||
if err == nil {
|
||||
return
|
||||
}
|
||||
t.itemMu.Lock()
|
||||
defer t.itemMu.Unlock()
|
||||
for i := range t.itemStates {
|
||||
item := &t.itemStates[i]
|
||||
if item.phase == ItemPhaseCompleted || item.phase == ItemPhaseFailed {
|
||||
continue
|
||||
}
|
||||
item.phase = ItemPhaseStopped
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) itemFailureStage(id string) FailureStage {
|
||||
t.itemMu.RLock()
|
||||
defer t.itemMu.RUnlock()
|
||||
index, ok := t.itemIndex[id]
|
||||
if !ok || index < 0 || index >= len(t.itemStates) {
|
||||
return FailureStageUpload
|
||||
}
|
||||
if t.itemStates[index].phase == ItemPhaseConfirming {
|
||||
return FailureStageConfirm
|
||||
}
|
||||
return FailureStageUpload
|
||||
}
|
||||
|
||||
func (t *Task) Items() []TaskItemProgress {
|
||||
t.itemMu.RLock()
|
||||
defer t.itemMu.RUnlock()
|
||||
items := make([]TaskItemProgress, 0, len(t.itemStates))
|
||||
for i := range t.itemStates {
|
||||
item := &t.itemStates[i]
|
||||
size := item.actualSize
|
||||
if size == 0 {
|
||||
size = item.expectedSize
|
||||
}
|
||||
items = append(items, TaskItemProgress{
|
||||
Index: item.index,
|
||||
ID: item.id,
|
||||
Name: item.name,
|
||||
Size: size,
|
||||
Downloaded: item.downloaded,
|
||||
Uploaded: item.uploaded,
|
||||
DownloadSpeed: item.downloadMeter.speed(),
|
||||
UploadSpeed: item.uploadMeter.speed(),
|
||||
Phase: item.phase,
|
||||
FailureStage: item.failureStage,
|
||||
RetryAttempt: item.retryAttempt,
|
||||
RetryLimit: item.retryLimit,
|
||||
Error: item.err,
|
||||
})
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
func (t *Task) ActualTotalSize() int64 {
|
||||
items := t.Items()
|
||||
var total int64
|
||||
for _, item := range items {
|
||||
total += item.Size
|
||||
}
|
||||
return total
|
||||
}
|
||||
|
||||
type stateProgressTracker interface {
|
||||
OnStateChange(ctx context.Context, info TaskInfo)
|
||||
}
|
||||
|
||||
func (t *Task) notifyStateChange(ctx context.Context) {
|
||||
if tracker, ok := t.Progress.(stateProgressTracker); ok {
|
||||
tracker.OnStateChange(ctx, t)
|
||||
}
|
||||
}
|
||||
|
||||
func compactError(err error) string {
|
||||
if err == nil {
|
||||
return ""
|
||||
}
|
||||
return strings.Join(strings.Fields(err.Error()), " ")
|
||||
}
|
||||
@@ -4,20 +4,19 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"path"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/duke-git/lancet/v2/slice"
|
||||
"github.com/gotd/td/telegram/message/entity"
|
||||
"github.com/gotd/td/telegram/message/styling"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/dlutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/tgutil"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
)
|
||||
|
||||
type ProgressTracker interface {
|
||||
@@ -27,159 +26,471 @@ type ProgressTracker interface {
|
||||
}
|
||||
|
||||
type Progress struct {
|
||||
MessageID int
|
||||
ChatID int64
|
||||
start time.Time
|
||||
lastUpdatePercent atomic.Int32
|
||||
skippedFiles []string
|
||||
MessageID int
|
||||
ChatID int64
|
||||
updateMu sync.Mutex
|
||||
lastUpdateAt time.Time
|
||||
lastText string
|
||||
done bool
|
||||
skippedFiles []string
|
||||
}
|
||||
|
||||
type renderedBatchMessage struct {
|
||||
Text string
|
||||
Entities []tg.MessageEntityClass
|
||||
Err error
|
||||
}
|
||||
|
||||
const (
|
||||
progressRenderInterval = time.Second
|
||||
maxVisibleActiveItems = 5
|
||||
progressBarWidth = 10
|
||||
maxDisplayNameRunes = 36
|
||||
maxDisplayErrorRunes = 240
|
||||
)
|
||||
|
||||
func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
p.start = time.Now()
|
||||
p.lastUpdatePercent.Store(0)
|
||||
log.FromContext(ctx).Debugf("Batch task progress tracking started for message %d in chat %d", p.MessageID, p.ChatID)
|
||||
entityBuilder := entity.Builder{}
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressBatchStartPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB (%d个文件)", float64(info.TotalSize())/(1024*1024), info.Count())),
|
||||
); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
|
||||
return
|
||||
}
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{
|
||||
Rows: []tg.KeyboardButtonRow{
|
||||
{
|
||||
Buttons: []tg.KeyboardButtonClass{
|
||||
tgutil.BuildCancelButton(info.TaskID()),
|
||||
},
|
||||
},
|
||||
}},
|
||||
)
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
return
|
||||
}
|
||||
p.render(ctx, info, true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo) {
|
||||
if !shouldUpdateProgress(info.TotalSize(), info.Downloaded(), int(p.lastUpdatePercent.Load())) {
|
||||
p.render(ctx, info, false)
|
||||
}
|
||||
|
||||
func (p *Progress) OnStateChange(ctx context.Context, info TaskInfo) {
|
||||
p.render(ctx, info, true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnUploadStart(ctx context.Context, info TaskInfo, _ int64) {
|
||||
p.render(ctx, info, true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnUploadProgress(ctx context.Context, info TaskInfo, _, _ int64) {
|
||||
p.render(ctx, info, false)
|
||||
}
|
||||
|
||||
func (p *Progress) render(ctx context.Context, info TaskInfo, priority bool) {
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
if p.done {
|
||||
return
|
||||
}
|
||||
percent := int((info.Downloaded() * 100) / info.TotalSize())
|
||||
if p.lastUpdatePercent.Load() == int32(percent) {
|
||||
now := time.Now()
|
||||
if !priority && !p.lastUpdateAt.IsZero() && now.Sub(p.lastUpdateAt) < progressRenderInterval {
|
||||
return
|
||||
}
|
||||
p.lastUpdatePercent.Store(int32(percent))
|
||||
log.FromContext(ctx).Debugf("Progress update: %s, %d/%d", info.TaskID(), info.Downloaded(), info.TotalSize())
|
||||
entityBuilder := entity.Builder{}
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressBatchProcessingPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB (%d个文件)", float64(info.TotalSize())/(1024*1024), info.Count())),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressProcessingListPrefix, nil)),
|
||||
func() styling.StyledTextOption {
|
||||
var lines []string
|
||||
for _, elem := range info.Processing() {
|
||||
lines = append(lines, fmt.Sprintf(" - %s (%.2f MB)", elem.FileName(), float64(elem.FileSize())/(1024*1024)))
|
||||
}
|
||||
if len(lines) == 0 {
|
||||
lines = append(lines, i18n.T(i18nk.BotMsgProgressProcessingNone, nil))
|
||||
}
|
||||
return styling.Plain(slice.Join(lines, "\n"))
|
||||
}(),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressAvgSpeedPrefix, nil)),
|
||||
styling.Bold(fmt.Sprintf("%.2f MB/s", dlutil.GetSpeed(info.Downloaded(), p.start)/(1024*1024))),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressCurrentProgressPrefix, nil)),
|
||||
styling.Bold(fmt.Sprintf("%.2f%%", float64(info.Downloaded())/float64(info.TotalSize())*100)),
|
||||
); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
|
||||
message := buildBatchProgressMessage(info, p.skippedFiles, visibleActiveItems())
|
||||
if message.Err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to render batch progress message: %v", message.Err)
|
||||
return
|
||||
}
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{
|
||||
Rows: []tg.KeyboardButtonRow{
|
||||
{
|
||||
Buttons: []tg.KeyboardButtonClass{
|
||||
tgutil.BuildCancelButton(info.TaskID()),
|
||||
},
|
||||
},
|
||||
}},
|
||||
)
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
if message.Text == p.lastText {
|
||||
return
|
||||
}
|
||||
p.lastText = message.Text
|
||||
p.lastUpdateAt = now
|
||||
p.editMessage(ctx, info.TaskID(), message, true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnDone(ctx context.Context, info TaskInfo, err error) {
|
||||
if err != nil {
|
||||
log.FromContext(ctx).Errorf("Batch task %s failed: %s", info.TaskID(), err)
|
||||
} else {
|
||||
log.FromContext(ctx).Debugf("Batch task %s completed successfully", info.TaskID())
|
||||
}
|
||||
entityBuilder := entity.Builder{}
|
||||
var stylingErr error
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, context.Canceled) {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressTaskCanceled, nil)),
|
||||
)
|
||||
} else {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressTaskFailedWithError, map[string]any{
|
||||
"Error": "",
|
||||
})),
|
||||
styling.Code(err.Error()),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressBatchDonePrefix, nil)),
|
||||
styling.Code(strconv.Itoa(info.Count())),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressTotalSizePrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB", float64(info.TotalSize())/(1024*1024))),
|
||||
func() styling.StyledTextOption {
|
||||
if len(p.skippedFiles) == 0 {
|
||||
return styling.Plain("")
|
||||
}
|
||||
return styling.Plain("\n\n" + i18n.T(i18nk.BotMsgCommonInfoConflictFilesSkipped, map[string]any{
|
||||
"Skipped": strings.Join(p.skippedFiles, "\n"),
|
||||
}))
|
||||
}(),
|
||||
)
|
||||
}
|
||||
|
||||
if stylingErr != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", stylingErr)
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
if p.done {
|
||||
return
|
||||
}
|
||||
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
p.done = true
|
||||
message := buildBatchDoneMessage(info, p.skippedFiles, err)
|
||||
if message.Err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to render final batch progress message: %v", message.Err)
|
||||
return
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
p.lastText = message.Text
|
||||
p.editMessage(ctx, info.TaskID(), message, false)
|
||||
}
|
||||
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
func (p *Progress) editMessage(ctx context.Context, taskID string, message renderedBatchMessage, cancellable bool) {
|
||||
if message.Err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to render batch progress message: %v", message.Err)
|
||||
return
|
||||
}
|
||||
req := buildBatchEditMessageRequest(p.MessageID, taskID, message, cancellable)
|
||||
if ext := tgutil.ExtFromContext(ctx); ext != nil {
|
||||
if _, err := ext.EditMessage(p.ChatID, req); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to edit batch progress message: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func buildBatchEditMessageRequest(messageID int, taskID string, message renderedBatchMessage, cancellable bool) *tg.MessagesEditMessageRequest {
|
||||
req := &tg.MessagesEditMessageRequest{ID: messageID}
|
||||
req.SetMessage(message.Text)
|
||||
if len(message.Entities) > 0 {
|
||||
req.SetEntities(message.Entities)
|
||||
}
|
||||
if cancellable {
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{Rows: []tg.KeyboardButtonRow{{
|
||||
Buttons: []tg.KeyboardButtonClass{tgutil.BuildCancelButton(taskID)},
|
||||
}}})
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func buildBatchProgressText(info TaskInfo, skipped []string, activeLimit int) string {
|
||||
return buildBatchProgressMessage(info, skipped, activeLimit).Text
|
||||
}
|
||||
|
||||
func buildBatchProgressMessage(info TaskInfo, skipped []string, activeLimit int) renderedBatchMessage {
|
||||
items := info.Items()
|
||||
completed, waiting, downloaded, failed := itemCounts(items)
|
||||
downloadSpeed, uploadSpeed := aggregateSpeeds(items)
|
||||
if activeLimit < 1 {
|
||||
activeLimit = 1
|
||||
}
|
||||
|
||||
total := len(items) + len(skipped)
|
||||
downloadSpeedText := formatSpeed(downloadSpeed)
|
||||
uploadSpeedText := formatSpeed(uploadSpeed)
|
||||
header := localizedProgressMarkup(i18nk.BotMsgProgressBatchStatusHeader, map[string]any{
|
||||
"Total": total,
|
||||
"Completed": completed,
|
||||
"Downloaded": downloaded,
|
||||
"Waiting": waiting,
|
||||
"DownloadSpeed": downloadSpeedText,
|
||||
"UploadSpeed": uploadSpeedText,
|
||||
})
|
||||
|
||||
var markup strings.Builder
|
||||
markup.WriteString(header)
|
||||
|
||||
visibleItems, hiddenTransfers, summarizedConfirming := visibleBatchItems(items, activeLimit)
|
||||
for _, item := range visibleItems {
|
||||
markup.WriteString("\n\n")
|
||||
markup.WriteString(formatActiveItemMarkup(item, len(items)))
|
||||
}
|
||||
|
||||
if hiddenTransfers > 0 {
|
||||
markup.WriteString("\n\n")
|
||||
markup.WriteString(localizedProgressMarkup(i18nk.BotMsgProgressBatchSummaryHiddenActive, map[string]any{"Count": hiddenTransfers}))
|
||||
}
|
||||
if summarizedConfirming > 0 {
|
||||
markup.WriteString("\n\n")
|
||||
markup.WriteString(localizedProgressMarkup(i18nk.BotMsgProgressBatchSummaryConfirming, map[string]any{"Count": summarizedConfirming}))
|
||||
}
|
||||
if failed > 0 {
|
||||
markup.WriteString("\n")
|
||||
markup.WriteString(localizedProgressMarkup(i18nk.BotMsgProgressBatchSummaryFailed, map[string]any{"Count": failed}))
|
||||
}
|
||||
if len(skipped) > 0 {
|
||||
markup.WriteString("\n")
|
||||
markup.WriteString(localizedProgressMarkup(i18nk.BotMsgProgressBatchSummarySkipped, map[string]any{"Count": len(skipped)}))
|
||||
}
|
||||
return completeBatchMessage(markup.String())
|
||||
}
|
||||
|
||||
func buildBatchDoneMarkup(info TaskInfo, skipped []string, err error) string {
|
||||
items := info.Items()
|
||||
totalSize := info.ActualTotalSize()
|
||||
if totalSize == 0 {
|
||||
totalSize = info.TotalSize()
|
||||
}
|
||||
if err == nil {
|
||||
if len(skipped) > 0 {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchDoneWithSkipped, map[string]any{
|
||||
"Success": len(items),
|
||||
"Skipped": len(skipped),
|
||||
"Size": dlutil.FormatSize(totalSize),
|
||||
})
|
||||
}
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchDone, map[string]any{
|
||||
"Count": len(items),
|
||||
"Size": dlutil.FormatSize(totalSize),
|
||||
})
|
||||
}
|
||||
completed, _, _, failed := itemCounts(items)
|
||||
incomplete := max(len(items)-completed-failed, 0)
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchCanceled, map[string]any{
|
||||
"Total": len(items) + len(skipped),
|
||||
"Completed": completed,
|
||||
"Incomplete": incomplete,
|
||||
"Skipped": len(skipped),
|
||||
})
|
||||
}
|
||||
|
||||
failedItems := make([]TaskItemProgress, 0, failed)
|
||||
for _, item := range items {
|
||||
if item.Phase == ItemPhaseFailed {
|
||||
failedItems = append(failedItems, item)
|
||||
}
|
||||
}
|
||||
if len(failedItems) > 1 && failedItems[0].FailureStage == FailureStageBatchUpload {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchFailedGroup, map[string]any{
|
||||
"Affected": len(failedItems),
|
||||
"Reason": displayError(firstError(failedItems), err),
|
||||
"Completed": completed,
|
||||
"Failed": len(failedItems),
|
||||
"Incomplete": incomplete,
|
||||
})
|
||||
}
|
||||
if len(failedItems) == 0 {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchFailedTask, map[string]any{
|
||||
"Reason": displayError("", err),
|
||||
"Completed": completed,
|
||||
"Incomplete": incomplete,
|
||||
})
|
||||
}
|
||||
item := failedItems[0]
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchFailedItem, map[string]any{
|
||||
"Index": item.Index,
|
||||
"Name": truncateFilename(item.Name, maxDisplayNameRunes),
|
||||
"Stage": failureStageLabel(item.FailureStage),
|
||||
"Progress": failureProgress(item),
|
||||
"Speed": failureSpeed(item),
|
||||
"Reason": displayError(item.Error, err),
|
||||
"Completed": completed,
|
||||
"Failed": failed,
|
||||
"Incomplete": incomplete,
|
||||
})
|
||||
}
|
||||
|
||||
func buildBatchDoneMessage(info TaskInfo, skipped []string, err error) renderedBatchMessage {
|
||||
return completeBatchMessage(buildBatchDoneMarkup(info, skipped, err))
|
||||
}
|
||||
|
||||
func formatActiveItemMarkup(item TaskItemProgress, total int) string {
|
||||
data := map[string]any{
|
||||
"Index": item.Index,
|
||||
"Total": total,
|
||||
"Name": truncateFilename(item.Name, maxDisplayNameRunes),
|
||||
"Speed": formatSpeed(itemSpeed(item)),
|
||||
"Progress": itemPercent(item),
|
||||
"Bar": textProgressBar(itemPercent(item)),
|
||||
"Current": dlutil.FormatSize(itemBytes(item)),
|
||||
"Size": dlutil.FormatSize(item.Size),
|
||||
"Attempt": min(max(item.RetryAttempt, 1), max(item.RetryLimit, 1)),
|
||||
"Limit": max(item.RetryLimit, 1),
|
||||
"Reason": truncateRunes(item.Error, maxDisplayErrorRunes),
|
||||
}
|
||||
switch item.Phase {
|
||||
case ItemPhaseDownloading:
|
||||
if item.Size <= 0 {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemDownloadingUnknown, data)
|
||||
}
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemDownloading, data)
|
||||
case ItemPhaseTransferring:
|
||||
if item.Size <= 0 {
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemTransferringUnknown, data)
|
||||
}
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemTransferring, data)
|
||||
case ItemPhaseUploading:
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemUploading, data)
|
||||
case ItemPhaseRetrying:
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemRetrying, data)
|
||||
case ItemPhaseConfirming:
|
||||
return localizedProgressMarkup(i18nk.BotMsgProgressBatchItemConfirming, data)
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func localizedProgressMarkup(key i18nk.Key, data map[string]any) string {
|
||||
return i18n.T(key, tgutil.EscapeHTMLTemplateData(data))
|
||||
}
|
||||
|
||||
func visibleBatchItems(items []TaskItemProgress, limit int) (visible []TaskItemProgress, hiddenTransfers, summarizedConfirming int) {
|
||||
visible = make([]TaskItemProgress, 0, limit)
|
||||
transferCount := 0
|
||||
confirmingCount := 0
|
||||
for _, item := range items {
|
||||
switch {
|
||||
case isTransferPhase(item.Phase):
|
||||
transferCount++
|
||||
if len(visible) < limit {
|
||||
visible = append(visible, item)
|
||||
}
|
||||
case item.Phase == ItemPhaseConfirming:
|
||||
confirmingCount++
|
||||
}
|
||||
}
|
||||
hiddenTransfers = transferCount - len(visible)
|
||||
if confirmingCount == 1 && len(visible) < limit {
|
||||
for _, item := range items {
|
||||
if item.Phase == ItemPhaseConfirming {
|
||||
visible = append(visible, item)
|
||||
return visible, hiddenTransfers, 0
|
||||
}
|
||||
}
|
||||
}
|
||||
return visible, hiddenTransfers, confirmingCount
|
||||
}
|
||||
|
||||
func completeBatchMessage(markup string) renderedBatchMessage {
|
||||
text, entities, err := tgutil.RenderHTML(markup)
|
||||
return renderedBatchMessage{Text: text, Entities: entities, Err: err}
|
||||
}
|
||||
|
||||
func itemCounts(items []TaskItemProgress) (completed, waiting, downloaded, failed int) {
|
||||
for _, item := range items {
|
||||
switch item.Phase {
|
||||
case ItemPhaseCompleted:
|
||||
completed++
|
||||
case ItemPhaseWaiting:
|
||||
waiting++
|
||||
case ItemPhaseDownloaded:
|
||||
downloaded++
|
||||
case ItemPhaseFailed:
|
||||
failed++
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func aggregateSpeeds(items []TaskItemProgress) (download, upload float64) {
|
||||
for _, item := range items {
|
||||
switch item.Phase {
|
||||
case ItemPhaseDownloading:
|
||||
download += item.DownloadSpeed
|
||||
case ItemPhaseTransferring:
|
||||
download += item.DownloadSpeed
|
||||
upload += item.UploadSpeed
|
||||
case ItemPhaseUploading:
|
||||
upload += item.UploadSpeed
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func isTransferPhase(phase ItemPhase) bool {
|
||||
switch phase {
|
||||
case ItemPhaseDownloading, ItemPhaseTransferring, ItemPhaseUploading, ItemPhaseRetrying:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func itemBytes(item TaskItemProgress) int64 {
|
||||
switch item.Phase {
|
||||
case ItemPhaseDownloading, ItemPhaseTransferring:
|
||||
return item.Downloaded
|
||||
default:
|
||||
return item.Uploaded
|
||||
}
|
||||
}
|
||||
|
||||
func itemSpeed(item TaskItemProgress) float64 {
|
||||
switch item.Phase {
|
||||
case ItemPhaseDownloading, ItemPhaseTransferring:
|
||||
return item.DownloadSpeed
|
||||
case ItemPhaseUploading:
|
||||
return item.UploadSpeed
|
||||
case ItemPhaseRetrying:
|
||||
return item.UploadSpeed
|
||||
default:
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func itemPercent(item TaskItemProgress) int {
|
||||
if item.Size <= 0 {
|
||||
return 0
|
||||
}
|
||||
return int(min(itemBytes(item), item.Size) * 100 / item.Size)
|
||||
}
|
||||
|
||||
func textProgressBar(percent int) string {
|
||||
percent = min(max(percent, 0), 100)
|
||||
filled := percent * progressBarWidth / 100
|
||||
return strings.Repeat("🟩", filled) + strings.Repeat("⬜️", progressBarWidth-filled)
|
||||
}
|
||||
|
||||
func formatSpeed(speed float64) string {
|
||||
if speed <= 0 {
|
||||
return "0 B/s"
|
||||
}
|
||||
return dlutil.FormatSize(int64(speed)) + "/s"
|
||||
}
|
||||
|
||||
func truncateFilename(name string, limit int) string {
|
||||
if utf8.RuneCountInString(name) <= limit {
|
||||
return name
|
||||
}
|
||||
ext := path.Ext(name)
|
||||
if utf8.RuneCountInString(ext) >= limit-2 {
|
||||
return truncateRunes(name, limit-1) + "…"
|
||||
}
|
||||
base := strings.TrimSuffix(name, ext)
|
||||
baseLimit := limit - utf8.RuneCountInString(ext) - 1
|
||||
return truncateRunes(base, baseLimit) + "…" + ext
|
||||
}
|
||||
|
||||
func truncateRunes(value string, limit int) string {
|
||||
if limit <= 0 {
|
||||
return ""
|
||||
}
|
||||
runes := []rune(value)
|
||||
if len(runes) <= limit {
|
||||
return value
|
||||
}
|
||||
return string(runes[:limit])
|
||||
}
|
||||
|
||||
func displayError(itemError string, fallback error) string {
|
||||
if itemError == "" && fallback != nil {
|
||||
itemError = compactError(fallback)
|
||||
}
|
||||
return truncateRunes(itemError, maxDisplayErrorRunes)
|
||||
}
|
||||
|
||||
func firstError(items []TaskItemProgress) string {
|
||||
for _, item := range items {
|
||||
if item.Error != "" {
|
||||
return item.Error
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func failureStageLabel(stage FailureStage) string {
|
||||
switch stage {
|
||||
case FailureStageDownload:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageDownload, nil)
|
||||
case FailureStageCache:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageCache, nil)
|
||||
case FailureStageUpload:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageUpload, nil)
|
||||
case FailureStageConfirm:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageConfirm, nil)
|
||||
case FailureStageBatchUpload:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageBatchUpload, nil)
|
||||
default:
|
||||
return i18n.T(i18nk.BotMsgProgressBatchFailureStageInternal, nil)
|
||||
}
|
||||
}
|
||||
|
||||
func failureProgress(item TaskItemProgress) string {
|
||||
if item.Size <= 0 {
|
||||
return dlutil.FormatSize(failureBytes(item))
|
||||
}
|
||||
return fmt.Sprintf("%d%%", min(failureBytes(item), item.Size)*100/item.Size)
|
||||
}
|
||||
|
||||
func failureSpeed(item TaskItemProgress) string {
|
||||
if item.FailureStage == FailureStageDownload || item.FailureStage == FailureStageCache {
|
||||
return formatSpeed(item.DownloadSpeed)
|
||||
}
|
||||
return formatSpeed(item.UploadSpeed)
|
||||
}
|
||||
|
||||
func failureBytes(item TaskItemProgress) int64 {
|
||||
if item.FailureStage == FailureStageDownload || item.FailureStage == FailureStageCache {
|
||||
return item.Downloaded
|
||||
}
|
||||
return item.Uploaded
|
||||
}
|
||||
|
||||
func visibleActiveItems() int {
|
||||
return min(max(config.C().Workers, 1), maxVisibleActiveItems)
|
||||
}
|
||||
|
||||
func NewProgressTracker(messageID int, chatID int64) ProgressTracker {
|
||||
|
||||
315
core/tasks/batchtfile/progress_regression_test.go
Normal file
315
core/tasks/batchtfile/progress_regression_test.go
Normal file
@@ -0,0 +1,315 @@
|
||||
package batchtfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||
)
|
||||
|
||||
type progressRegressionRecorder struct {
|
||||
mu sync.Mutex
|
||||
startTotal int64
|
||||
notifications []int64
|
||||
}
|
||||
|
||||
func (*progressRegressionRecorder) OnStart(context.Context, TaskInfo) {}
|
||||
func (*progressRegressionRecorder) OnProgress(context.Context, TaskInfo) {}
|
||||
func (*progressRegressionRecorder) OnDone(context.Context, TaskInfo, error) {}
|
||||
|
||||
func (r *progressRegressionRecorder) OnUploadStart(_ context.Context, _ TaskInfo, total int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.startTotal = total
|
||||
}
|
||||
|
||||
func (r *progressRegressionRecorder) OnUploadProgress(_ context.Context, _ TaskInfo, uploaded, _ int64) {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
r.notifications = append(r.notifications, uploaded)
|
||||
}
|
||||
|
||||
type orderedProgressRegressionRecorder struct {
|
||||
firstEntered chan struct{}
|
||||
releaseFirst chan struct{}
|
||||
secondEntered chan struct{}
|
||||
mu sync.Mutex
|
||||
notifications []int64
|
||||
}
|
||||
|
||||
func (*orderedProgressRegressionRecorder) OnStart(context.Context, TaskInfo) {}
|
||||
func (*orderedProgressRegressionRecorder) OnProgress(context.Context, TaskInfo) {}
|
||||
func (*orderedProgressRegressionRecorder) OnDone(context.Context, TaskInfo, error) {}
|
||||
func (*orderedProgressRegressionRecorder) OnUploadStart(context.Context, TaskInfo, int64) {
|
||||
}
|
||||
|
||||
func (r *orderedProgressRegressionRecorder) OnUploadProgress(_ context.Context, _ TaskInfo, uploaded, _ int64) {
|
||||
if uploaded == 100 {
|
||||
close(r.firstEntered)
|
||||
<-r.releaseFirst
|
||||
}
|
||||
if uploaded == 200 {
|
||||
close(r.secondEntered)
|
||||
}
|
||||
r.mu.Lock()
|
||||
r.notifications = append(r.notifications, uploaded)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func TestBatchProgressShowsTransferSpeedAndSize(t *testing.T) {
|
||||
useProgressRegressionLocale(t)
|
||||
task := newProgressRegressionTask(nil,
|
||||
progressRegressionFile{"downloading", 1000},
|
||||
progressRegressionFile{"uploading", 1000},
|
||||
progressRegressionFile{"waiting", 1000},
|
||||
)
|
||||
started := time.Unix(100, 0)
|
||||
task.markItemActive("downloading", false, started)
|
||||
task.recordItemDownload("downloading", 500, started.Add(time.Second))
|
||||
task.recordItemDownloaded("uploading", 1000)
|
||||
task.recordItemUpload("uploading", 0, 1000, started.Add(time.Second))
|
||||
task.recordItemUpload("uploading", 250, 1000, started.Add(2*time.Second))
|
||||
|
||||
message := buildBatchProgressMessage(task, nil, 2)
|
||||
if message.Err != nil {
|
||||
t.Fatalf("buildBatchProgressMessage() failed: %v", message.Err)
|
||||
}
|
||||
assertProgressRegressionContains(t, message.Text,
|
||||
"状态：✅ 0 ｜ 📥 0 ｜ ⏳ 1",
|
||||
"⬇️ 1/3 下载中",
|
||||
"速度：500 B/s",
|
||||
"大小：500 B / 1000 B",
|
||||
"⬆️ 2/3 上传中",
|
||||
"速度：250 B/s",
|
||||
"大小：250 B / 1000 B",
|
||||
)
|
||||
bold, _, blockquote, _ := batchEntityCounts(message.Entities)
|
||||
if bold != 3 || blockquote != 2 {
|
||||
t.Fatalf("entity counts = bold:%d blockquote:%d, want bold:3 blockquote:2", bold, blockquote)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchProgressLimitsRowsWithoutHidingActiveUpload(t *testing.T) {
|
||||
useProgressRegressionLocale(t)
|
||||
task := newProgressRegressionTask(nil,
|
||||
progressRegressionFile{"confirm-01", 100},
|
||||
progressRegressionFile{"confirm-02", 100},
|
||||
progressRegressionFile{"uploading", 100},
|
||||
progressRegressionFile{"downloading", 100},
|
||||
)
|
||||
started := time.Unix(100, 0)
|
||||
task.recordItemUpload("confirm-01", 100, 100, started)
|
||||
task.recordItemUpload("confirm-02", 100, 100, started)
|
||||
task.recordItemUpload("uploading", 40, 100, started.Add(time.Second))
|
||||
task.markItemActive("downloading", false, started)
|
||||
|
||||
message := buildBatchProgressText(task, nil, 2)
|
||||
assertProgressRegressionContains(t, message,
|
||||
"uploading.bin",
|
||||
"downloading.bin",
|
||||
"☁️ 已上传，等待整组发送：2",
|
||||
)
|
||||
if strings.Contains(message, "confirm-01.bin") || strings.Contains(message, "confirm-02.bin") {
|
||||
t.Fatalf("confirmation rows displaced active transfers:\n%s", message)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchProgressTemplateOwnsStylesAndEscapesValues(t *testing.T) {
|
||||
useProgressRegressionLocale(t)
|
||||
fileID := `<b>A&B</b>`
|
||||
task := newProgressRegressionTask(nil, progressRegressionFile{fileID, 100})
|
||||
task.markItemRetry(fileID, FailureStageUpload, 1, 3, errors.New(`<i>remote & failed</i>`))
|
||||
|
||||
message := buildBatchProgressMessage(task, nil, 1)
|
||||
if message.Err != nil {
|
||||
t.Fatalf("buildBatchProgressMessage() failed: %v", message.Err)
|
||||
}
|
||||
assertProgressRegressionContains(t, message.Text,
|
||||
`<b>A&B</b>.bin`,
|
||||
`<i>remote & failed</i>`,
|
||||
)
|
||||
bold, _, blockquote, italic := batchEntityCounts(message.Entities)
|
||||
if bold != 2 || blockquote != 1 || italic != 0 {
|
||||
t.Fatalf("entity counts = bold:%d blockquote:%d italic:%d", bold, blockquote, italic)
|
||||
}
|
||||
|
||||
i18n.Init("en")
|
||||
english := buildBatchProgressMessage(task, nil, 1)
|
||||
if english.Err != nil {
|
||||
t.Fatalf("English batch template failed: %v", english.Err)
|
||||
}
|
||||
assertProgressRegressionContains(t, english.Text, "📦 Processing", "Retrying upload", `<b>A&B</b>.bin`)
|
||||
}
|
||||
|
||||
func TestDownloadProgressContinuesAfterUploadStarts(t *testing.T) {
|
||||
useProgressRegressionLocale(t)
|
||||
progress := new(Progress)
|
||||
task := newProgressRegressionTask(progress,
|
||||
progressRegressionFile{"uploading", 100},
|
||||
progressRegressionFile{"downloading", 100},
|
||||
)
|
||||
progress.OnStart(t.Context(), task)
|
||||
task.recordDownloadComplete("uploading", 100)
|
||||
task.uploadCallback(t.Context(), "uploading")(50, 100)
|
||||
|
||||
started := time.Unix(100, 0)
|
||||
task.markItemActive("downloading", false, started)
|
||||
task.recordItemDownload("downloading", 50, started.Add(time.Second))
|
||||
progress.updateMu.Lock()
|
||||
progress.lastUpdateAt = time.Now().Add(-progressRenderInterval)
|
||||
progress.updateMu.Unlock()
|
||||
progress.OnProgress(t.Context(), task)
|
||||
|
||||
progress.updateMu.Lock()
|
||||
text := progress.lastText
|
||||
progress.updateMu.Unlock()
|
||||
assertProgressRegressionContains(t, text,
|
||||
"uploading.bin",
|
||||
"🟩🟩🟩🟩🟩⬜️⬜️⬜️⬜️⬜️ 50%",
|
||||
"总速度：⬇️ 50 B/s ｜ ⬆️ 0 B/s",
|
||||
"🔄 另有 1 个文件正在处理",
|
||||
)
|
||||
}
|
||||
|
||||
func TestBatchUploadIgnoresOutOfOrderBytesAndAllowsRetryReset(t *testing.T) {
|
||||
recorder := new(progressRegressionRecorder)
|
||||
task := newProgressRegressionTask(recorder, progressRegressionFile{"file", 100})
|
||||
task.recordDownloadComplete("file", 100)
|
||||
callback := task.uploadCallback(t.Context(), "file")
|
||||
callback(80, 100)
|
||||
callback(10, 100)
|
||||
|
||||
if got := task.Items()[0].Uploaded; got != 80 {
|
||||
t.Fatalf("out-of-order callback regressed item to %d, want 80", got)
|
||||
}
|
||||
task.markItemRetry("file", FailureStageUpload, 1, 3, context.DeadlineExceeded)
|
||||
callback(0, 100)
|
||||
callback(10, 100)
|
||||
if got := task.Items()[0].Uploaded; got != 10 {
|
||||
t.Fatalf("retry did not reset item progress: got %d, want 10", got)
|
||||
}
|
||||
|
||||
recorder.mu.Lock()
|
||||
defer recorder.mu.Unlock()
|
||||
for index := 1; index < len(recorder.notifications); index++ {
|
||||
if recorder.notifications[index] < recorder.notifications[index-1] {
|
||||
t.Fatalf("aggregate progress regressed: %v", recorder.notifications)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadProgressNotificationsRemainOrdered(t *testing.T) {
|
||||
recorder := &orderedProgressRegressionRecorder{
|
||||
firstEntered: make(chan struct{}),
|
||||
releaseFirst: make(chan struct{}),
|
||||
secondEntered: make(chan struct{}),
|
||||
}
|
||||
task := newProgressRegressionTask(recorder,
|
||||
progressRegressionFile{"first", 100},
|
||||
progressRegressionFile{"second", 100},
|
||||
)
|
||||
task.recordDownloadComplete("first", 100)
|
||||
task.recordDownloadComplete("second", 100)
|
||||
first := task.uploadCallback(t.Context(), "first")
|
||||
second := task.uploadCallback(t.Context(), "second")
|
||||
|
||||
var wait sync.WaitGroup
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
first(100, 100)
|
||||
}()
|
||||
<-recorder.firstEntered
|
||||
wait.Add(1)
|
||||
go func() {
|
||||
defer wait.Done()
|
||||
second(100, 100)
|
||||
}()
|
||||
|
||||
overtook := false
|
||||
select {
|
||||
case <-recorder.secondEntered:
|
||||
overtook = true
|
||||
case <-time.After(100 * time.Millisecond):
|
||||
}
|
||||
close(recorder.releaseFirst)
|
||||
wait.Wait()
|
||||
if overtook {
|
||||
t.Fatal("later aggregate notification overtook the first callback")
|
||||
}
|
||||
|
||||
recorder.mu.Lock()
|
||||
defer recorder.mu.Unlock()
|
||||
if got := recorder.notifications; len(got) != 2 || got[0] != 100 || got[1] != 200 {
|
||||
t.Fatalf("upload notifications = %v, want [100 200]", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBatchUploadUsesActualSizeWhenMetadataIsUnknown(t *testing.T) {
|
||||
recorder := new(progressRegressionRecorder)
|
||||
task := newProgressRegressionTask(recorder, progressRegressionFile{"photo", 0})
|
||||
task.recordDownloadComplete("photo", 25)
|
||||
task.uploadCallback(t.Context(), "photo")(25, 25)
|
||||
|
||||
recorder.mu.Lock()
|
||||
defer recorder.mu.Unlock()
|
||||
if recorder.startTotal != 25 {
|
||||
t.Fatalf("upload start total = %d, want actual size 25", recorder.startTotal)
|
||||
}
|
||||
if got := task.ActualTotalSize(); got != 25 {
|
||||
t.Fatalf("actual total size = %d, want 25", got)
|
||||
}
|
||||
}
|
||||
|
||||
type progressRegressionFile struct {
|
||||
id string
|
||||
size int64
|
||||
}
|
||||
|
||||
func newProgressRegressionTask(progress ProgressTracker, files ...progressRegressionFile) *Task {
|
||||
elems := make([]TaskElement, 0, len(files))
|
||||
for _, file := range files {
|
||||
elems = append(elems, TaskElement{
|
||||
ID: file.id,
|
||||
File: tfile.NewTGFile(nil, nil, file.size, file.id+".bin"),
|
||||
})
|
||||
}
|
||||
return NewBatchTGFileTask("progress-regression", context.Background(), elems, progress, true)
|
||||
}
|
||||
|
||||
func useProgressRegressionLocale(t *testing.T) {
|
||||
t.Helper()
|
||||
i18n.Init("zh-Hans")
|
||||
t.Cleanup(func() { i18n.Init("zh-Hans") })
|
||||
}
|
||||
|
||||
func assertProgressRegressionContains(t *testing.T, value string, wants ...string) {
|
||||
t.Helper()
|
||||
for _, want := range wants {
|
||||
if !strings.Contains(value, want) {
|
||||
t.Fatalf("text does not contain %q:\n%s", want, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func batchEntityCounts(entities []tg.MessageEntityClass) (bold, code, blockquote, italic int) {
|
||||
for _, messageEntity := range entities {
|
||||
switch messageEntity.(type) {
|
||||
case *tg.MessageEntityBold:
|
||||
bold++
|
||||
case *tg.MessageEntityCode:
|
||||
code++
|
||||
case *tg.MessageEntityBlockquote:
|
||||
blockquote++
|
||||
case *tg.MessageEntityItalic:
|
||||
italic++
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
@@ -31,16 +31,23 @@ type TaskElement struct {
|
||||
}
|
||||
|
||||
type Task struct {
|
||||
ID string
|
||||
ctx context.Context
|
||||
elems []TaskElement
|
||||
Progress ProgressTracker
|
||||
IgnoreErrors bool // if true, errors during processing will be ignored
|
||||
downloaded atomic.Int64
|
||||
totalSize int64
|
||||
processing map[string]TaskElementInfo
|
||||
processingMu sync.RWMutex
|
||||
failed map[string]error // [TODO] errors for each element
|
||||
ID string
|
||||
ctx context.Context
|
||||
elems []TaskElement
|
||||
Progress ProgressTracker
|
||||
IgnoreErrors bool // if true, errors during processing will be ignored
|
||||
downloaded atomic.Int64
|
||||
totalSize int64
|
||||
uploadTotalSize atomic.Int64
|
||||
processing map[string]TaskElementInfo
|
||||
processingMu sync.RWMutex
|
||||
itemStates []itemProgressState
|
||||
itemIndex map[string]int
|
||||
itemMu sync.RWMutex
|
||||
uploadOnce sync.Once
|
||||
uploadMu sync.Mutex
|
||||
uploaded map[string]int64
|
||||
failed map[string]error // [TODO] errors for each element
|
||||
}
|
||||
|
||||
// Title implements core.Exectable.
|
||||
@@ -109,6 +116,7 @@ func NewBatchTGFileTask(
|
||||
progress ProgressTracker,
|
||||
ignoreErrors bool,
|
||||
) *Task {
|
||||
itemStates, itemIndex := newItemProgressStates(files)
|
||||
task := &Task{
|
||||
ID: id,
|
||||
ctx: ctx,
|
||||
@@ -123,6 +131,9 @@ func NewBatchTGFileTask(
|
||||
return total
|
||||
}(),
|
||||
processing: make(map[string]TaskElementInfo),
|
||||
itemStates: itemStates,
|
||||
itemIndex: itemIndex,
|
||||
uploaded: make(map[string]int64),
|
||||
IgnoreErrors: ignoreErrors,
|
||||
processingMu: sync.RWMutex{},
|
||||
failed: make(map[string]error),
|
||||
|
||||
@@ -27,8 +27,10 @@ type TaskInfo interface {
|
||||
TaskID() string
|
||||
TotalSize() int64
|
||||
Downloaded() int64
|
||||
ActualTotalSize() int64
|
||||
Count() int
|
||||
Processing() []TaskElementInfo
|
||||
Items() []TaskItemProgress
|
||||
}
|
||||
|
||||
func (t *Task) TaskID() string {
|
||||
|
||||
70
core/tasks/batchtfile/upload_progress.go
Normal file
70
core/tasks/batchtfile/upload_progress.go
Normal file
@@ -0,0 +1,70 @@
|
||||
package batchtfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
)
|
||||
|
||||
// UploadProgressTracker optionally extends a batch progress tracker with a
|
||||
// distinct aggregate upload phase.
|
||||
type UploadProgressTracker interface {
|
||||
OnUploadStart(ctx context.Context, info TaskInfo, total int64)
|
||||
OnUploadProgress(ctx context.Context, info TaskInfo, uploaded, total int64)
|
||||
}
|
||||
|
||||
func (t *Task) startUpload(ctx context.Context) {
|
||||
tracker, ok := t.Progress.(UploadProgressTracker)
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
t.uploadMu.Lock()
|
||||
defer t.uploadMu.Unlock()
|
||||
t.uploadOnce.Do(func() {
|
||||
tracker.OnUploadStart(ctx, t, t.uploadTotalSize.Load())
|
||||
})
|
||||
}
|
||||
|
||||
func (t *Task) uploadCallback(ctx context.Context, id string) func(uploaded, total int64) {
|
||||
return func(uploaded, total int64) {
|
||||
tracker, ok := t.Progress.(UploadProgressTracker)
|
||||
if !ok || uploaded < 0 {
|
||||
return
|
||||
}
|
||||
t.startUpload(ctx)
|
||||
if total > 0 && uploaded > total {
|
||||
uploaded = total
|
||||
}
|
||||
|
||||
t.uploadMu.Lock()
|
||||
defer t.uploadMu.Unlock()
|
||||
becameConfirming := t.recordItemUpload(id, uploaded, total, time.Now())
|
||||
if t.uploaded == nil {
|
||||
t.uploaded = make(map[string]int64)
|
||||
}
|
||||
previous, tracked := t.uploaded[id]
|
||||
if !tracked || uploaded > previous {
|
||||
t.uploaded[id] = uploaded
|
||||
}
|
||||
var aggregate int64
|
||||
for _, current := range t.uploaded {
|
||||
aggregate += current
|
||||
}
|
||||
uploadTotal := t.uploadTotalSize.Load()
|
||||
if aggregate > uploadTotal {
|
||||
aggregate = uploadTotal
|
||||
}
|
||||
tracker.OnUploadProgress(ctx, t, aggregate, uploadTotal)
|
||||
if becameConfirming {
|
||||
t.notifyStateChange(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *Task) recordDownloadComplete(id string, uploadSize int64) {
|
||||
t.uploadMu.Lock()
|
||||
defer t.uploadMu.Unlock()
|
||||
if uploadSize > 0 {
|
||||
t.uploadTotalSize.Add(uploadSize)
|
||||
}
|
||||
t.recordItemDownloaded(id, uploadSize)
|
||||
}
|
||||
@@ -3,6 +3,7 @@ package tfile
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"path"
|
||||
|
||||
@@ -10,10 +11,12 @@ import (
|
||||
"github.com/duke-git/lancet/v2/retry"
|
||||
"github.com/krau/SaveAny-Bot/common/tdler"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/ioutil"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
"github.com/krau/SaveAny-Bot/pkg/storagetypes"
|
||||
tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
)
|
||||
|
||||
func (t *Task) Execute(ctx context.Context) error {
|
||||
@@ -68,7 +71,25 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
return fmt.Errorf("failed to open cache file: %w", err)
|
||||
}
|
||||
defer file.Close()
|
||||
if err = t.Storage.Save(vctx, file, t.Path); err != nil {
|
||||
uploadProgress, tracksUpload := t.Progress.(UploadProgressTracker)
|
||||
if !tracksUpload {
|
||||
if err = t.Storage.Save(vctx, file, t.Path); err != nil {
|
||||
return fmt.Errorf("failed to save file: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
uploadProgress.OnUploadStart(vctx, t, fileStat.Size())
|
||||
onProgress := func(uploaded, total int64) {
|
||||
uploadProgress.OnUploadProgress(vctx, t, uploaded, total)
|
||||
}
|
||||
if progressSaver, ok := t.Storage.(storage.StorageProgressSaver); ok {
|
||||
err = progressSaver.SaveWithProgress(vctx, file, t.Path, onProgress)
|
||||
} else {
|
||||
var reader io.Reader = ioutil.NewProgressReader(file, fileStat.Size(), onProgress)
|
||||
err = t.Storage.Save(vctx, reader, t.Path)
|
||||
}
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to save file: %w", err)
|
||||
}
|
||||
return nil
|
||||
|
||||
@@ -4,12 +4,12 @@ import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/gotd/td/telegram/message/entity"
|
||||
"github.com/gotd/td/telegram/message/styling"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n/i18nk"
|
||||
@@ -23,153 +23,319 @@ type ProgressTracker interface {
|
||||
OnDone(ctx context.Context, info TaskInfo, err error)
|
||||
}
|
||||
|
||||
// UploadProgressTracker optionally extends a task progress tracker with a
|
||||
// distinct upload phase. Keeping it separate preserves compatibility with
|
||||
// custom download-only trackers.
|
||||
type UploadProgressTracker interface {
|
||||
OnUploadStart(ctx context.Context, info TaskInfo, total int64)
|
||||
OnUploadProgress(ctx context.Context, info TaskInfo, uploaded, total int64)
|
||||
}
|
||||
|
||||
type Progress struct {
|
||||
MessageID int
|
||||
ChatID int64
|
||||
start time.Time
|
||||
lastUpdatePercent atomic.Int32
|
||||
lastUpdateAt atomic.Int64
|
||||
updateMu sync.Mutex
|
||||
uploadAttempt int
|
||||
uploadedBytes int64
|
||||
actualSize int64
|
||||
hasActualSize bool
|
||||
}
|
||||
|
||||
const (
|
||||
uploadProgressMinInterval = time.Second
|
||||
uploadProgressMaxInterval = 3 * time.Second
|
||||
singleProgressBarWidth = 10
|
||||
maxSingleErrorRunes = 240
|
||||
)
|
||||
|
||||
type singleProgressPhase int
|
||||
|
||||
const (
|
||||
singlePhaseDownloading singleProgressPhase = iota
|
||||
singlePhaseUploading
|
||||
singlePhaseRetrying
|
||||
)
|
||||
|
||||
type renderedSingleMessage struct {
|
||||
Text string
|
||||
Entities []tg.MessageEntityClass
|
||||
Err error
|
||||
}
|
||||
|
||||
func (p *Progress) OnStart(ctx context.Context, info TaskInfo) {
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
p.start = time.Now()
|
||||
p.lastUpdatePercent.Store(0)
|
||||
p.lastUpdateAt.Store(0)
|
||||
p.uploadAttempt = 0
|
||||
p.uploadedBytes = 0
|
||||
p.actualSize = 0
|
||||
p.hasActualSize = false
|
||||
log.FromContext(ctx).Debugf("Progress tracking started for message %d in chat %d", p.MessageID, p.ChatID)
|
||||
entityBuilder := entity.Builder{}
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressFileStartPrefix, nil)),
|
||||
styling.Code(info.FileName()),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("[%s]:%s", info.StorageName(), info.StoragePath())),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressFileSizePrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB", float64(info.FileSize())/(1024*1024))),
|
||||
); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
|
||||
return
|
||||
}
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{
|
||||
Rows: []tg.KeyboardButtonRow{
|
||||
{
|
||||
Buttons: []tg.KeyboardButtonClass{
|
||||
tgutil.BuildCancelButton(info.TaskID()),
|
||||
},
|
||||
},
|
||||
}},
|
||||
)
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
return
|
||||
}
|
||||
p.editMessage(ctx, info.TaskID(), buildSingleProgressMessage(info, singlePhaseDownloading, 0, info.FileSize(), 0, 0), true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnProgress(ctx context.Context, info TaskInfo, downloaded, total int64) {
|
||||
if !shouldUpdateProgress(total, downloaded, int(p.lastUpdatePercent.Load())) {
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
now := time.Now()
|
||||
elapsed := uploadProgressMaxInterval
|
||||
if lastUpdateAt := p.lastUpdateAt.Load(); lastUpdateAt > 0 {
|
||||
elapsed = now.Sub(time.Unix(0, lastUpdateAt))
|
||||
}
|
||||
if !shouldUpdateSingleDownloadProgress(total, downloaded, int(p.lastUpdatePercent.Load()), elapsed) {
|
||||
return
|
||||
}
|
||||
percent := int32((downloaded * 100) / total)
|
||||
if p.lastUpdatePercent.Load() == percent {
|
||||
return
|
||||
if total > 0 {
|
||||
percent := int32((downloaded * 100) / total)
|
||||
if p.lastUpdatePercent.Load() == percent {
|
||||
return
|
||||
}
|
||||
p.lastUpdatePercent.Store(percent)
|
||||
}
|
||||
p.lastUpdatePercent.Store(percent)
|
||||
p.lastUpdateAt.Store(now.UnixNano())
|
||||
log.FromContext(ctx).Debugf("Progress update: %s, %d/%d", info.FileName(), downloaded, total)
|
||||
entityBuilder := entity.Builder{}
|
||||
var entities []tg.MessageEntityClass
|
||||
if err := styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressFileProcessingPrefix, nil)),
|
||||
styling.Code(info.FileName()),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("[%s]:%s", info.StorageName(), info.StoragePath())),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressFileSizePrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("%.2f MB", float64(total)/(1024*1024))),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressAvgSpeedPrefix, nil)),
|
||||
styling.Bold(fmt.Sprintf("%.2f MB/s", dlutil.GetSpeed(downloaded, p.start)/(1024*1024))),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressCurrentProgressPrefix, nil)),
|
||||
styling.Bold(fmt.Sprintf("%.2f%%", float64(downloaded)/float64(total)*100)),
|
||||
); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", err)
|
||||
p.editMessage(ctx, info.TaskID(), buildSingleProgressMessage(
|
||||
info,
|
||||
singlePhaseDownloading,
|
||||
downloaded,
|
||||
total,
|
||||
dlutil.GetSpeed(downloaded, p.start),
|
||||
0,
|
||||
), true)
|
||||
}
|
||||
|
||||
func shouldUpdateSingleDownloadProgress(total, downloaded int64, lastPercent int, elapsed time.Duration) bool {
|
||||
if total > 0 {
|
||||
return shouldUpdateProgress(total, downloaded, lastPercent)
|
||||
}
|
||||
return downloaded > 0 && elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
|
||||
func (p *Progress) OnUploadStart(ctx context.Context, info TaskInfo, total int64) {
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
p.start = time.Now()
|
||||
p.lastUpdatePercent.Store(0)
|
||||
p.lastUpdateAt.Store(p.start.UnixNano())
|
||||
p.uploadAttempt++
|
||||
p.uploadedBytes = 0
|
||||
p.actualSize = max(total, 0)
|
||||
p.hasActualSize = true
|
||||
log.FromContext(ctx).Debugf("Upload progress tracking started: %s", info.FileName())
|
||||
phase := singleUploadPhase(p.uploadAttempt)
|
||||
p.editMessage(ctx, info.TaskID(), buildSingleProgressMessage(info, phase, 0, total, 0, p.uploadAttempt), true)
|
||||
}
|
||||
|
||||
func (p *Progress) OnUploadProgress(ctx context.Context, info TaskInfo, uploaded, total int64) {
|
||||
if total <= 0 || uploaded <= 0 {
|
||||
return
|
||||
}
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
if uploaded > total {
|
||||
uploaded = total
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{
|
||||
Rows: []tg.KeyboardButtonRow{
|
||||
{
|
||||
Buttons: []tg.KeyboardButtonClass{
|
||||
tgutil.BuildCancelButton(info.TaskID()),
|
||||
},
|
||||
},
|
||||
}},
|
||||
)
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
if uploaded < p.uploadedBytes {
|
||||
return
|
||||
}
|
||||
p.uploadedBytes = uploaded
|
||||
|
||||
now := time.Now()
|
||||
lastUpdateAt := time.Unix(0, p.lastUpdateAt.Load())
|
||||
lastPercent := int(p.lastUpdatePercent.Load())
|
||||
if !shouldUpdateUploadProgress(total, uploaded, lastPercent, now.Sub(lastUpdateAt)) {
|
||||
return
|
||||
}
|
||||
|
||||
percent := int32((uploaded * 100) / total)
|
||||
p.lastUpdatePercent.Store(percent)
|
||||
p.lastUpdateAt.Store(now.UnixNano())
|
||||
log.FromContext(ctx).Debugf("Upload progress update: %s, %d/%d", info.FileName(), uploaded, total)
|
||||
p.editMessage(ctx, info.TaskID(), buildSingleProgressMessage(
|
||||
info,
|
||||
singleUploadPhase(p.uploadAttempt),
|
||||
uploaded,
|
||||
total,
|
||||
dlutil.GetSpeed(uploaded, p.start),
|
||||
p.uploadAttempt,
|
||||
), true)
|
||||
}
|
||||
|
||||
func shouldUpdateUploadProgress(total, uploaded int64, lastPercent int, elapsed time.Duration) bool {
|
||||
if total <= 0 || uploaded <= 0 {
|
||||
return false
|
||||
}
|
||||
if uploaded >= total {
|
||||
return lastPercent < 100 && elapsed >= uploadProgressMinInterval
|
||||
}
|
||||
percent := int((uploaded * 100) / total)
|
||||
if percent < lastPercent {
|
||||
return false
|
||||
}
|
||||
if elapsed < uploadProgressMinInterval {
|
||||
return false
|
||||
}
|
||||
if percent == lastPercent {
|
||||
return elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
return shouldUpdateProgress(total, uploaded, lastPercent) || elapsed >= uploadProgressMaxInterval
|
||||
}
|
||||
|
||||
func singleUploadPhase(attempt int) singleProgressPhase {
|
||||
if attempt > 1 {
|
||||
return singlePhaseRetrying
|
||||
}
|
||||
return singlePhaseUploading
|
||||
}
|
||||
|
||||
func (p *Progress) OnDone(ctx context.Context, info TaskInfo, err error) {
|
||||
p.updateMu.Lock()
|
||||
defer p.updateMu.Unlock()
|
||||
if err != nil {
|
||||
log.FromContext(ctx).Errorf("Progress error for file [%s]: %v", info.FileName(), err)
|
||||
} else {
|
||||
log.FromContext(ctx).Debugf("Progress done for file [%s]", info.FileName())
|
||||
}
|
||||
|
||||
entityBuilder := entity.Builder{}
|
||||
var stylingErr error
|
||||
p.editMessage(ctx, info.TaskID(), buildSingleDoneMessage(info, p.doneSize(info), err), false)
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, context.Canceled) {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressTaskCanceled, nil)),
|
||||
styling.Plain("\n"),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressFileNamePrefix, nil)),
|
||||
styling.Code(info.FileName()),
|
||||
)
|
||||
} else {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressDownloadFailedPrefix, nil)),
|
||||
styling.Code(info.FileName()),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressErrorPrefix, nil)),
|
||||
styling.Bold(err.Error()),
|
||||
)
|
||||
}
|
||||
} else {
|
||||
stylingErr = styling.Perform(&entityBuilder,
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressDownloadDonePrefix, nil)),
|
||||
styling.Code(info.FileName()),
|
||||
styling.Plain(i18n.T(i18nk.BotMsgProgressSavePathPrefix, nil)),
|
||||
styling.Code(fmt.Sprintf("[%s]:%s", info.StorageName(), info.StoragePath())),
|
||||
)
|
||||
func (p *Progress) doneSize(info TaskInfo) int64 {
|
||||
if p.hasActualSize {
|
||||
return p.actualSize
|
||||
}
|
||||
return max(info.FileSize(), 0)
|
||||
}
|
||||
|
||||
if stylingErr != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to build entities: %s", stylingErr)
|
||||
func (p *Progress) editMessage(ctx context.Context, taskID string, message renderedSingleMessage, cancellable bool) {
|
||||
if message.Err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to render file progress message: %v", message.Err)
|
||||
return
|
||||
}
|
||||
|
||||
text, entities := entityBuilder.Complete()
|
||||
req := &tg.MessagesEditMessageRequest{
|
||||
ID: p.MessageID,
|
||||
req := buildSingleEditMessageRequest(p.MessageID, taskID, message, cancellable)
|
||||
if ext := tgutil.ExtFromContext(ctx); ext != nil {
|
||||
if _, err := ext.EditMessage(p.ChatID, req); err != nil {
|
||||
log.FromContext(ctx).Errorf("Failed to edit file progress message: %v", err)
|
||||
}
|
||||
}
|
||||
req.SetMessage(text)
|
||||
req.SetEntities(entities)
|
||||
}
|
||||
|
||||
ext := tgutil.ExtFromContext(ctx)
|
||||
if ext != nil {
|
||||
ext.EditMessage(p.ChatID, req)
|
||||
func buildSingleEditMessageRequest(messageID int, taskID string, message renderedSingleMessage, cancellable bool) *tg.MessagesEditMessageRequest {
|
||||
req := &tg.MessagesEditMessageRequest{ID: messageID}
|
||||
req.SetMessage(message.Text)
|
||||
if len(message.Entities) > 0 {
|
||||
req.SetEntities(message.Entities)
|
||||
}
|
||||
if cancellable {
|
||||
req.SetReplyMarkup(&tg.ReplyInlineMarkup{Rows: []tg.KeyboardButtonRow{{
|
||||
Buttons: []tg.KeyboardButtonClass{tgutil.BuildCancelButton(taskID)},
|
||||
}}})
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func buildSingleProgressMessage(
|
||||
info TaskInfo,
|
||||
phase singleProgressPhase,
|
||||
current int64,
|
||||
total int64,
|
||||
speed float64,
|
||||
attempt int,
|
||||
) renderedSingleMessage {
|
||||
if current < 0 {
|
||||
current = 0
|
||||
}
|
||||
if total > 0 && current > total {
|
||||
current = total
|
||||
}
|
||||
percent := singleProgressPercent(current, total)
|
||||
destination := fmt.Sprintf("[%s]:%s", info.StorageName(), info.StoragePath())
|
||||
data := map[string]any{
|
||||
"Name": info.FileName(),
|
||||
"Bar": singleProgressBar(percent),
|
||||
"Progress": percent,
|
||||
"Speed": singleProgressSpeed(speed),
|
||||
"Current": dlutil.FormatSize(current),
|
||||
"Size": dlutil.FormatSize(total),
|
||||
"Destination": destination,
|
||||
"Attempt": max(attempt, 1),
|
||||
}
|
||||
|
||||
var key i18nk.Key
|
||||
switch phase {
|
||||
case singlePhaseUploading:
|
||||
key = i18nk.BotMsgProgressSingleUploading
|
||||
case singlePhaseRetrying:
|
||||
key = i18nk.BotMsgProgressSingleUploadRetrying
|
||||
default:
|
||||
key = i18nk.BotMsgProgressSingleDownloading
|
||||
if total <= 0 {
|
||||
key = i18nk.BotMsgProgressSingleDownloadingUnknown
|
||||
}
|
||||
}
|
||||
markup := i18n.T(i18nk.BotMsgProgressSingleStatusHeader, nil) + "\n\n" + localizedProgressMarkup(key, data)
|
||||
return completeSingleMessage(markup)
|
||||
}
|
||||
|
||||
func buildSingleDoneMessage(info TaskInfo, size int64, err error) renderedSingleMessage {
|
||||
data := map[string]any{
|
||||
"Name": info.FileName(),
|
||||
"Size": dlutil.FormatSize(max(size, 0)),
|
||||
"Destination": fmt.Sprintf("[%s]:%s", info.StorageName(), info.StoragePath()),
|
||||
}
|
||||
var key i18nk.Key
|
||||
switch {
|
||||
case err == nil:
|
||||
key = i18nk.BotMsgProgressSingleDone
|
||||
case errors.Is(err, context.Canceled):
|
||||
key = i18nk.BotMsgProgressSingleCanceled
|
||||
default:
|
||||
data["Reason"] = truncateSingleError(err.Error())
|
||||
key = i18nk.BotMsgProgressSingleFailed
|
||||
}
|
||||
return completeSingleMessage(localizedProgressMarkup(key, data))
|
||||
}
|
||||
|
||||
func localizedProgressMarkup(key i18nk.Key, data map[string]any) string {
|
||||
return i18n.T(key, tgutil.EscapeHTMLTemplateData(data))
|
||||
}
|
||||
|
||||
func completeSingleMessage(markup string) renderedSingleMessage {
|
||||
text, entities, err := tgutil.RenderHTML(markup)
|
||||
return renderedSingleMessage{Text: text, Entities: entities, Err: err}
|
||||
}
|
||||
|
||||
func singleProgressPercent(current, total int64) int {
|
||||
if total <= 0 {
|
||||
return 0
|
||||
}
|
||||
return int(min(max(current, 0), total) * 100 / total)
|
||||
}
|
||||
|
||||
func singleProgressBar(percent int) string {
|
||||
percent = min(max(percent, 0), 100)
|
||||
filled := percent * singleProgressBarWidth / 100
|
||||
return strings.Repeat("🟩", filled) + strings.Repeat("⬜️", singleProgressBarWidth-filled)
|
||||
}
|
||||
|
||||
func singleProgressSpeed(speed float64) string {
|
||||
if speed <= 0 {
|
||||
return "0 B/s"
|
||||
}
|
||||
return dlutil.FormatSize(int64(speed)) + "/s"
|
||||
}
|
||||
|
||||
func truncateSingleError(value string) string {
|
||||
runes := []rune(value)
|
||||
if len(runes) <= maxSingleErrorRunes {
|
||||
return value
|
||||
}
|
||||
return string(runes[:maxSingleErrorRunes])
|
||||
}
|
||||
|
||||
type ProgressOption func(*Progress)
|
||||
|
||||
180
core/tasks/tfile/progress_test.go
Normal file
180
core/tasks/tfile/progress_test.go
Normal file
@@ -0,0 +1,180 @@
|
||||
package tfile
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/krau/SaveAny-Bot/common/i18n"
|
||||
)
|
||||
|
||||
type progressTestTaskInfo struct{}
|
||||
|
||||
func (progressTestTaskInfo) TaskID() string { return "task" }
|
||||
func (progressTestTaskInfo) FileName() string { return "file.bin" }
|
||||
func (progressTestTaskInfo) FileSize() int64 { return 100 << 20 }
|
||||
func (progressTestTaskInfo) StoragePath() string { return "file.bin" }
|
||||
func (progressTestTaskInfo) StorageName() string { return "test" }
|
||||
|
||||
func TestShouldUpdateUploadProgress(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
total int64
|
||||
uploaded int64
|
||||
lastPercent int
|
||||
elapsed time.Duration
|
||||
want bool
|
||||
}{
|
||||
{name: "invalid total", total: 0, uploaded: 1, want: false},
|
||||
{name: "no uploaded bytes", total: 100, uploaded: 0, want: false},
|
||||
{name: "percentage threshold", total: 100 << 20, uploaded: 10 << 20, elapsed: uploadProgressMinInterval, want: true},
|
||||
{name: "percentage threshold rate limited", total: 100 << 20, uploaded: 10 << 20, elapsed: uploadProgressMinInterval - time.Millisecond, want: false},
|
||||
{name: "maximum time threshold", total: 100 << 20, uploaded: 1 << 20, elapsed: uploadProgressMaxInterval, want: true},
|
||||
{name: "below thresholds", total: 100 << 20, uploaded: 1 << 20, elapsed: uploadProgressMaxInterval - time.Millisecond, want: false},
|
||||
{name: "completion", total: 100, uploaded: 100, lastPercent: 99, elapsed: uploadProgressMinInterval, want: true},
|
||||
{name: "completion rate limited", total: 100, uploaded: 100, lastPercent: 99, elapsed: uploadProgressMinInterval - time.Millisecond, want: false},
|
||||
{name: "completion already reported", total: 100, uploaded: 100, lastPercent: 100, elapsed: uploadProgressMinInterval, want: false},
|
||||
{name: "out of order callback", total: 100, uploaded: 40, lastPercent: 60, elapsed: uploadProgressMaxInterval, want: false},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := shouldUpdateUploadProgress(tt.total, tt.uploaded, tt.lastPercent, tt.elapsed)
|
||||
if got != tt.want {
|
||||
t.Fatalf("shouldUpdateUploadProgress() = %v, want %v", got, tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestUploadProgressConcurrentCallbacks(t *testing.T) {
|
||||
progress := new(Progress)
|
||||
ctx := context.Background()
|
||||
info := progressTestTaskInfo{}
|
||||
const total = int64(100 << 20)
|
||||
|
||||
progress.OnUploadStart(ctx, info, total)
|
||||
progress.lastUpdateAt.Store(time.Now().Add(-uploadProgressMaxInterval).UnixNano())
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for uploaded := int64(1 << 20); uploaded <= total; uploaded += 1 << 20 {
|
||||
uploaded := uploaded
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
progress.OnUploadProgress(ctx, info, uploaded, total)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
percent := progress.lastUpdatePercent.Load()
|
||||
if percent <= 0 || percent > 100 {
|
||||
t.Fatalf("last upload percentage = %d, want a value in (0, 100]", percent)
|
||||
}
|
||||
if progress.uploadedBytes != total {
|
||||
t.Fatalf("maximum uploaded bytes = %d, want %d", progress.uploadedBytes, total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSingleUploadRetryKeepsRichProgressLayout(t *testing.T) {
|
||||
i18n.Init("zh-Hans")
|
||||
t.Cleanup(func() { i18n.Init("zh-Hans") })
|
||||
|
||||
message := buildSingleProgressMessage(
|
||||
progressTestTaskInfo{},
|
||||
singleUploadPhase(2),
|
||||
25<<20,
|
||||
100<<20,
|
||||
5<<20,
|
||||
2,
|
||||
)
|
||||
if message.Err != nil {
|
||||
t.Fatalf("buildSingleProgressMessage() failed: %v", message.Err)
|
||||
}
|
||||
for _, want := range []string{
|
||||
"🔁 上传重试",
|
||||
"🟩🟩⬜️⬜️⬜️⬜️⬜️⬜️⬜️⬜️ 25%",
|
||||
"尝试次数：2",
|
||||
"速度：5.00 MB/s",
|
||||
"大小：25.00 MB / 100.00 MB",
|
||||
} {
|
||||
if !strings.Contains(message.Text, want) {
|
||||
t.Fatalf("retry progress does not contain %q:\n%s", want, message.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestSingleProgressTemplateOwnsStylesAndEscapesValues(t *testing.T) {
|
||||
i18n.Init("en")
|
||||
t.Cleanup(func() { i18n.Init("zh-Hans") })
|
||||
info := htmlProgressTestTaskInfo{}
|
||||
|
||||
message := buildSingleProgressMessage(info, singlePhaseDownloading, 50, 100, 25, 0)
|
||||
if message.Err != nil {
|
||||
t.Fatalf("buildSingleProgressMessage() failed: %v", message.Err)
|
||||
}
|
||||
for _, want := range []string{
|
||||
`<b>A&B</b>.bin`,
|
||||
`[store<&>]:dir/<i>x</i>&`,
|
||||
"Speed: 25 B/s",
|
||||
} {
|
||||
if !strings.Contains(message.Text, want) {
|
||||
t.Fatalf("progress text does not contain %q:\n%s", want, message.Text)
|
||||
}
|
||||
}
|
||||
bold, code, blockquote, italic := singleEntityCounts(message.Entities)
|
||||
if bold != 2 || code != 6 || blockquote != 1 || italic != 0 {
|
||||
t.Fatalf("progress entity counts = bold:%d code:%d blockquote:%d italic:%d", bold, code, blockquote, italic)
|
||||
}
|
||||
|
||||
failure := buildSingleDoneMessage(info, 100, errors.New(`<i>remote & failed</i>`))
|
||||
if failure.Err != nil {
|
||||
t.Fatalf("buildSingleDoneMessage() failed: %v", failure.Err)
|
||||
}
|
||||
if !strings.Contains(failure.Text, `<i>remote & failed</i>`) {
|
||||
t.Fatalf("failure reason was not preserved literally:\n%s", failure.Text)
|
||||
}
|
||||
bold, code, blockquote, italic = singleEntityCounts(failure.Entities)
|
||||
if bold != 1 || code != 2 || blockquote != 0 || italic != 0 {
|
||||
t.Fatalf("failure entity counts = bold:%d code:%d blockquote:%d italic:%d", bold, code, blockquote, italic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSingleDoneSizeUsesActualUploadSize(t *testing.T) {
|
||||
progress := new(Progress)
|
||||
info := progressTestTaskInfo{}
|
||||
progress.OnStart(context.Background(), info)
|
||||
progress.OnUploadStart(context.Background(), info, 2048)
|
||||
|
||||
if got := progress.doneSize(info); got != 2048 {
|
||||
t.Fatalf("done size = %d, want actual upload size 2048", got)
|
||||
}
|
||||
}
|
||||
|
||||
type htmlProgressTestTaskInfo struct{}
|
||||
|
||||
func (htmlProgressTestTaskInfo) TaskID() string { return "html-task" }
|
||||
func (htmlProgressTestTaskInfo) FileName() string { return `<b>A&B</b>.bin` }
|
||||
func (htmlProgressTestTaskInfo) FileSize() int64 { return 100 }
|
||||
func (htmlProgressTestTaskInfo) StoragePath() string { return `dir/<i>x</i>&` }
|
||||
func (htmlProgressTestTaskInfo) StorageName() string { return `store<&>` }
|
||||
|
||||
func singleEntityCounts(entities []tg.MessageEntityClass) (bold, code, blockquote, italic int) {
|
||||
for _, messageEntity := range entities {
|
||||
switch messageEntity.(type) {
|
||||
case *tg.MessageEntityBold:
|
||||
bold++
|
||||
case *tg.MessageEntityCode:
|
||||
code++
|
||||
case *tg.MessageEntityBlockquote:
|
||||
blockquote++
|
||||
case *tg.MessageEntityItalic:
|
||||
italic++
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
Reference in New Issue
Block a user