feat: implement task event system for progress tracking and reporting (#220)

This commit is contained in:
Krau
2026-06-25 21:37:36 +08:00
committed by GitHub
parent 9c2e70ed43
commit f02860ff3f
17 changed files with 397 additions and 214 deletions

View File

@@ -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
}

View File

@@ -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
}

View File

@@ -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 {

View File

@@ -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)

View File

@@ -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)

View File

@@ -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
})
}

View File

@@ -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
}

View File

@@ -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