mirror of
https://github.com/krau/SaveAny-Bot.git
synced 2026-06-29 03:01:24 +08:00
feat: implement task event system for progress tracking and reporting (#220)
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/tasktype"
|
||||
"github.com/krau/SaveAny-Bot/pkg/queue"
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
)
|
||||
|
||||
var queueInstance *queue.TaskQueue[Executable]
|
||||
@@ -30,11 +31,14 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan
|
||||
break // queue closed and empty
|
||||
}
|
||||
exe := qtask.Data
|
||||
taskCtx := qtask.Context()
|
||||
logger.Infof("Processing task: %s", exe.TaskID())
|
||||
if err := ExecCommandString(qtask.Context(), execHooks.TaskBeforeStart); err != nil {
|
||||
taskevent.Emit(taskCtx, taskevent.Event{TaskID: exe.TaskID(), Phase: taskevent.PhaseStart})
|
||||
if err := ExecCommandString(taskCtx, execHooks.TaskBeforeStart); err != nil {
|
||||
logger.Errorf("Failed to execute before start hook for task %s: %v", exe.TaskID(), err)
|
||||
}
|
||||
if err := exe.Execute(qtask.Context()); err != nil {
|
||||
err = exe.Execute(taskCtx)
|
||||
if err != nil {
|
||||
if errors.Is(err, context.Canceled) {
|
||||
logger.Infof("Task %s was canceled", exe.TaskID())
|
||||
if err := ExecCommandString(ctx, execHooks.TaskCancel); err != nil {
|
||||
@@ -52,6 +56,7 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan
|
||||
logger.Errorf("Failed to execute success hook for task %s: %v", exe.TaskID(), err)
|
||||
}
|
||||
}
|
||||
taskevent.Emit(taskCtx, taskevent.Event{TaskID: exe.TaskID(), Phase: taskevent.PhaseDone, Err: err})
|
||||
qe.Done(qtask.ID)
|
||||
<-semaphore
|
||||
}
|
||||
|
||||
@@ -6,12 +6,14 @@ import (
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"time"
|
||||
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/aria2"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
)
|
||||
|
||||
// Execute implements core.Executable.
|
||||
@@ -77,6 +79,12 @@ func (t *Task) waitForDownload(ctx context.Context) error {
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnProgress(ctx, t, status)
|
||||
}
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: parseInt64(status.TotalLength),
|
||||
DownloadedBytes: parseInt64(status.CompletedLength),
|
||||
})
|
||||
|
||||
// Check if download is complete
|
||||
if status.IsDownloadComplete() {
|
||||
@@ -248,3 +256,16 @@ func (t *Task) cancelAria2Download() {
|
||||
logger.Debugf("Failed to remove download result for %s: %v", t.GID, err)
|
||||
}
|
||||
}
|
||||
|
||||
// parseInt64 parses an aria2 status string (decimal bytes) into int64,
|
||||
// returning 0 on failure so it can be used directly in progress events.
|
||||
func parseInt64(s string) int64 {
|
||||
if s == "" {
|
||||
return 0
|
||||
}
|
||||
n, err := strconv.ParseInt(s, 10, 64)
|
||||
if err != nil {
|
||||
return 0
|
||||
}
|
||||
return n
|
||||
}
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"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/taskevent"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
@@ -62,8 +63,14 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
return elem.Storage.Save(uploadCtx, pr, elem.Path)
|
||||
})
|
||||
wr := ioutil.NewProgressWriter(pw, func(n int) {
|
||||
t.downloaded.Add(int64(n))
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: t.totalSize,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
})
|
||||
errg.Go(func() error {
|
||||
defer pw.Close()
|
||||
@@ -92,8 +99,14 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
}
|
||||
}()
|
||||
wrAt := ioutil.NewProgressWriterAt(localFile, func(n int) {
|
||||
t.downloaded.Add(int64(n))
|
||||
downloaded := t.downloaded.Add(int64(n))
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: t.totalSize,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
})
|
||||
_, err = tdler.NewDownloader(elem.File).Parallel(ctx, wrAt)
|
||||
if err != nil {
|
||||
|
||||
@@ -15,6 +15,7 @@ import (
|
||||
"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/taskevent"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
@@ -143,10 +144,16 @@ func (t *Task) processLink(ctx context.Context, file *File) error {
|
||||
}
|
||||
}()
|
||||
wr := ioutil.NewProgressWriter(cacheFile, func(n int) {
|
||||
t.downloadedBytes.Add(int64(n))
|
||||
downloaded := t.downloadedBytes.Add(int64(n))
|
||||
if t.Progress != nil {
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
}
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: t.totalBytes,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
})
|
||||
|
||||
copyResultCh := make(chan error, 1)
|
||||
|
||||
@@ -16,6 +16,7 @@ import (
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
"github.com/krau/SaveAny-Bot/pkg/parser"
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
@@ -107,10 +108,16 @@ func (t *Task) processResource(ctx context.Context, resource parser.Resource) er
|
||||
}
|
||||
}()
|
||||
wr := ioutil.NewProgressWriter(cacheFile, func(n int) {
|
||||
t.downloadedBytes.Add(int64(n))
|
||||
downloaded := t.downloadedBytes.Add(int64(n))
|
||||
if t.progress != nil {
|
||||
t.progress.OnProgress(ctx, t)
|
||||
}
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: t.totalBytes,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
})
|
||||
|
||||
copyResultCh := make(chan error, 1)
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/duke-git/lancet/v2/retry"
|
||||
"github.com/krau/SaveAny-Bot/common/utils/fsutil"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
|
||||
@@ -27,8 +28,14 @@ func (t *Task) Execute(ctx context.Context) error {
|
||||
logger.Errorf("Error processing picture %s: %v", pic, err)
|
||||
return fmt.Errorf("failed to process picture %s: %w", pic, err)
|
||||
}
|
||||
t.downloaded.Add(1)
|
||||
downloaded := t.downloaded.Add(1)
|
||||
t.progress.OnProgress(gctx, t)
|
||||
taskevent.Emit(gctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalFiles: t.totalpics,
|
||||
DownloadedFiles: int(downloaded),
|
||||
})
|
||||
return nil
|
||||
})
|
||||
}
|
||||
|
||||
@@ -4,6 +4,8 @@ import (
|
||||
"context"
|
||||
"io"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
)
|
||||
|
||||
type ProgressWriterAt struct {
|
||||
@@ -20,9 +22,16 @@ func (w *ProgressWriterAt) WriteAt(p []byte, off int64) (int, error) {
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
downloaded := w.downloaded.Add(int64(at))
|
||||
if w.progress != nil {
|
||||
w.progress.OnProgress(w.ctx, w.info, w.downloaded.Add(int64(at)), w.total)
|
||||
w.progress.OnProgress(w.ctx, w.info, downloaded, w.total)
|
||||
}
|
||||
taskevent.Emit(w.ctx, taskevent.Event{
|
||||
TaskID: w.info.TaskID(),
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: w.total,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
return at, nil
|
||||
}
|
||||
|
||||
@@ -56,9 +65,16 @@ func (w *ProgressWriter) Write(p []byte) (int, error) {
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
downloaded := w.downloaded.Add(int64(at))
|
||||
if w.progress != nil {
|
||||
w.progress.OnProgress(w.ctx, w.info, w.downloaded.Add(int64(at)), w.total)
|
||||
w.progress.OnProgress(w.ctx, w.info, downloaded, w.total)
|
||||
}
|
||||
taskevent.Emit(w.ctx, taskevent.Event{
|
||||
TaskID: w.info.TaskID(),
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: w.total,
|
||||
DownloadedBytes: downloaded,
|
||||
})
|
||||
return at, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"github.com/charmbracelet/log"
|
||||
"github.com/krau/SaveAny-Bot/config"
|
||||
"github.com/krau/SaveAny-Bot/pkg/enums/ctxkey"
|
||||
"github.com/krau/SaveAny-Bot/pkg/taskevent"
|
||||
"github.com/krau/SaveAny-Bot/storage"
|
||||
"golang.org/x/sync/errgroup"
|
||||
)
|
||||
@@ -116,6 +117,12 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error {
|
||||
|
||||
t.uploaded.Add(size)
|
||||
t.Progress.OnProgress(ctx, t)
|
||||
taskevent.Emit(ctx, taskevent.Event{
|
||||
TaskID: t.ID,
|
||||
Phase: taskevent.PhaseProgress,
|
||||
TotalBytes: t.totalSize,
|
||||
DownloadedBytes: t.uploaded.Load(),
|
||||
})
|
||||
|
||||
logger.Info("File uploaded successfully")
|
||||
return nil
|
||||
|
||||
Reference in New Issue
Block a user