diff --git a/cmd/run.go b/cmd/run.go index 33020ba..90581b8 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -17,6 +17,7 @@ import ( "github.com/krau/SaveAny-Bot/common/cache" "github.com/krau/SaveAny-Bot/common/i18n" "github.com/krau/SaveAny-Bot/common/utils/fsutil" + "github.com/krau/SaveAny-Bot/common/utils/tgutil" "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/core" "github.com/krau/SaveAny-Bot/database" @@ -63,7 +64,12 @@ func Run(cmd *cobra.Command, _ []string) { } return nil }) - core.RecoverTasks(ctx) + // 恢复任务携带 ext 上下文, 让进度编辑/取消按钮在恢复后继续工作。 + recoverCtx := context.Background() + if ectx := bot.ExtContext(); ectx != nil { + recoverCtx = tgutil.ExtWithContext(recoverCtx, ectx) + } + core.RecoverTasks(recoverCtx) core.Run(ctx) diff --git a/core/core.go b/core/core.go index 975377b..fb03db1 100644 --- a/core/core.go +++ b/core/core.go @@ -74,7 +74,9 @@ func worker(ctx context.Context, qe *queue.TaskQueue[Executable], semaphore chan } taskevent.Emit(taskCtx, taskevent.Event{TaskID: exe.TaskID(), Phase: taskevent.PhaseDone, Err: err}) qe.Done(qtask.ID) - if err := database.DeleteTask(ctx, exe.TaskID()); err != nil { + // 用独立 ctx 删除: 优雅关停时 run ctx 已被取消, 会留下已完成任务的行, + // 导致重启后重复执行 (重复上传)。 + if err := database.DeleteTask(context.Background(), exe.TaskID()); err != nil { logger.Errorf("Failed to delete persisted task %s: %v", exe.TaskID(), err) } <-semaphore diff --git a/core/persist.go b/core/persist.go index 9379039..99976ee 100644 --- a/core/persist.go +++ b/core/persist.go @@ -3,6 +3,9 @@ package core import ( "context" "sync" + "time" + + "fmt" "github.com/charmbracelet/log" "github.com/gotd/td/telegram/downloader" @@ -75,11 +78,28 @@ func persistTask(ctx context.Context, task Executable) error { }) } +// UpdateTaskPayload atomically mutates the persisted payload of a running +// task (e.g. recording per-element upload progress for recovery). +func UpdateTaskPayload(ctx context.Context, id string, mutate func(payload []byte) ([]byte, error)) error { + row, err := database.GetTask(ctx, id) + if err != nil { + return err + } + updated, err := mutate(row.Payload) + if err != nil { + return fmt.Errorf("mutate payload: %w", err) + } + return database.UpdateTaskPayload(ctx, id, updated) +} + // RecoverTasks re-enqueues tasks that were unfinished when the process last -// exited. Must be called after storages are loaded and before Run. Tasks of -// types without a registered codec are dropped with a warning. +// exited. Must be called after storages are loaded and before Run. Tasks +// that cannot be recovered are marked failed and kept for visibility. func RecoverTasks(ctx context.Context) { logger := log.FromContext(ctx) + if err := database.DeleteStaleFailedTasks(ctx, 24*time.Hour); err != nil { + logger.Warnf("Failed to clean stale failed tasks: %v", err) + } tasks, err := database.GetUnfinishedTasks(ctx) if err != nil { logger.Errorf("Failed to load unfinished tasks: %v", err) @@ -88,27 +108,38 @@ func RecoverTasks(ctx context.Context) { for _, t := range tasks { codec, ok := TaskCodecFor(tasktype.TaskType(t.Type)) if !ok { - logger.Warnf("Dropping unrecoverable task %s (type %s): no codec registered", t.ID, t.Type) - if err := database.DeleteTask(ctx, t.ID); err != nil { - logger.Errorf("Failed to delete task %s: %v", t.ID, err) - } + logger.Warnf("Task %s (type %s) cannot be recovered: no codec registered", t.ID, t.Type) + markRecoverFailed(ctx, t, "no codec registered") continue } task, err := codec.Unmarshal(t.Payload) if err != nil { - logger.Errorf("Dropping task %s: failed to rebuild: %v", t.ID, err) - if err := database.DeleteTask(ctx, t.ID); err != nil { - logger.Errorf("Failed to delete task %s: %v", t.ID, err) - } + logger.Errorf("Task %s cannot be recovered: failed to rebuild: %v", t.ID, err) + markRecoverFailed(ctx, t, err.Error()) + continue + } + if initQueue().Contains(task.TaskID()) { + // Already live in the queue (e.g. submitted via API during + // startup); keep the row as-is. + logger.Infof("Task %s already queued, keeping row", t.ID) continue } if err := AddTask(ctx, task); err != nil { - logger.Errorf("Dropping task %s: failed to re-enqueue: %v", t.ID, err) - if err := database.DeleteTask(ctx, t.ID); err != nil { - logger.Errorf("Failed to delete task %s: %v", t.ID, err) - } + logger.Errorf("Task %s cannot be recovered: failed to re-enqueue: %v", t.ID, err) + markRecoverFailed(ctx, t, err.Error()) continue } + // Upsert cleared the original creation time; restore it so + // GetUnfinishedTasks ordering stays stable across restarts. + if err := database.RestoreTaskCreatedAt(ctx, t.ID, t.CreatedAt); err != nil { + logger.Warnf("Failed to restore created_at for task %s: %v", t.ID, err) + } logger.Infof("Recovered task %s (%s)", t.ID, t.Type) } } + +func markRecoverFailed(ctx context.Context, t database.Task, reason string) { + if err := database.UpdateTaskStatus(ctx, t.ID, database.TaskStatusFailed, reason); err != nil { + log.FromContext(ctx).Errorf("Failed to mark task %s as failed: %v", t.ID, err) + } +} diff --git a/core/persist_test.go b/core/persist_test.go index 3db8784..85da92a 100644 --- a/core/persist_test.go +++ b/core/persist_test.go @@ -52,7 +52,7 @@ func initRecoveryEnv(t *testing.T) context.Context { return context.Background() } -func TestRecoverTasksReenqueuesAndDropsUnknown(t *testing.T) { +func TestRecoverTasksReenqueuesAndMarksUnknownFailed(t *testing.T) { ctx := initRecoveryEnv(t) if err := database.CreateTask(ctx, &database.Task{ @@ -96,9 +96,20 @@ func TestRecoverTasksReenqueuesAndDropsUnknown(t *testing.T) { t.Fatalf("recovered task status = %s, want queued", task.Status) } } + // The unrecoverable task must be kept and marked failed, not silently deleted. + drop, err := database.GetTask(ctx, "drop-1") + if err != nil { + t.Fatalf("dropped task row missing: %v", err) + } + if drop.Status != string(database.TaskStatusFailed) { + t.Fatalf("dropped task status = %s, want failed", drop.Status) + } + if drop.Error == "" { + t.Fatalf("dropped task has no failure reason") + } } -func TestRecoverTasksDropsInvalidPayload(t *testing.T) { +func TestRecoverTasksMarksInvalidPayloadFailed(t *testing.T) { ctx := initRecoveryEnv(t) if err := database.CreateTask(ctx, &database.Task{ @@ -108,7 +119,7 @@ func TestRecoverTasksDropsInvalidPayload(t *testing.T) { } RecoverTasks(ctx) - // bad-1 must not be enqueued nor remain in the database. + // bad-1 must not be enqueued; its row is kept as failed. for _, info := range GetQueuedTasks(ctx) { if info.ID == "bad-1" { t.Fatalf("task with invalid payload was enqueued") @@ -121,4 +132,31 @@ func TestRecoverTasksDropsInvalidPayload(t *testing.T) { if count != 0 { t.Fatalf("unfinished rows = %d, want 0", count) } + bad, err := database.GetTask(ctx, "bad-1") + if err != nil { + t.Fatalf("failed task row missing: %v", err) + } + if bad.Status != string(database.TaskStatusFailed) { + t.Fatalf("bad task status = %s, want failed", bad.Status) + } +} + +func TestRecoverTasksSkipsAlreadyQueued(t *testing.T) { + ctx := initRecoveryEnv(t) + + // A task submitted during startup is both persisted and in the queue. + task := &stubTask{id: "live-1"} + if err := AddTask(ctx, task); err != nil { + t.Fatal(err) + } + RecoverTasks(ctx) + + // The row must survive with its original status. + row, err := database.GetTask(ctx, "live-1") + if err != nil { + t.Fatalf("row missing for queued task: %v", err) + } + if row.Status != string(database.TaskStatusQueued) { + t.Fatalf("row status = %s, want queued", row.Status) + } } diff --git a/core/tasks/batchtfile/codec.go b/core/tasks/batchtfile/codec.go index 489ada7..5785eb1 100644 --- a/core/tasks/batchtfile/codec.go +++ b/core/tasks/batchtfile/codec.go @@ -8,6 +8,8 @@ import ( "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/core" + tftask "github.com/krau/SaveAny-Bot/core/tasks/tfile" + "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" "github.com/krau/SaveAny-Bot/pkg/enums/tasktype" tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile" "github.com/krau/SaveAny-Bot/storage" @@ -24,27 +26,86 @@ type elementPayload struct { } type taskPayload struct { + Kind string `json:"kind"` // "batch" ID string `json:"id"` Elements []elementPayload `json:"elements"` ChatID int64 `json:"chat_id"` MessageID int `json:"message_id"` IgnoreErrors bool `json:"ignore_errors"` + Overwrite bool `json:"overwrite"` + // Done lists element IDs whose upload completed; they are skipped on recovery. + Done []string `json:"done"` } -type taskCodec struct{} +// tgfilesCodec is the single codec registered for TaskTypeTgfiles: it +// dispatches between single-file and batch tasks by concrete type on marshal +// and by payload shape on unmarshal. Registering one codec per task class +// under the shared TaskTypeTgfiles key would let the last init() win and +// silently disable persistence for the other class. +type tgfilesCodec struct{} func init() { - core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, taskCodec{}) + core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, tgfilesCodec{}) } -func (taskCodec) Marshal(task core.Executable) ([]byte, error) { +func (tgfilesCodec) Marshal(task core.Executable) ([]byte, error) { + switch t := task.(type) { + case *tftask.Task: + return tftask.TaskCodec.Marshal(t) + case *Task: + return batchCodec{}.Marshal(t) + default: + return nil, fmt.Errorf("unexpected task type %T", task) + } +} + +// detectTaskKind returns "batch" or "file" for a persisted tgfiles payload. +// New payloads carry an explicit kind; legacy payloads are detected by shape. +func detectTaskKind(data []byte) (string, error) { + var shape struct { + Kind string `json:"kind"` + Elements []json.RawMessage `json:"elements"` + File json.RawMessage `json:"file"` + } + if err := json.Unmarshal(data, &shape); err != nil { + return "", fmt.Errorf("invalid task payload: %w", err) + } + switch { + case shape.Kind == "batch", shape.Kind == "" && shape.Elements != nil: + return "batch", nil + case shape.Kind == "file", shape.Kind == "" && shape.File != nil: + return "file", nil + default: + return "", fmt.Errorf("unrecognized task payload") + } +} + +func (tgfilesCodec) Unmarshal(data []byte) (core.Executable, error) { + kind, err := detectTaskKind(data) + if err != nil { + return nil, err + } + if kind == "batch" { + return batchCodec{}.Unmarshal(data) + } + return tftask.TaskCodec.Unmarshal(data) +} + +type batchCodec struct{} + +func (batchCodec) Marshal(task core.Executable) ([]byte, error) { t, ok := task.(*Task) if !ok { return nil, fmt.Errorf("unexpected task type %T", task) } p := taskPayload{ + Kind: "batch", ID: t.ID, IgnoreErrors: t.IgnoreErrors, + Done: t.completedElementIDs(), + } + if overwrite, ok := t.ctx.Value(ctxkey.OverwriteExisting).(bool); ok { + p.Overwrite = overwrite } for _, elem := range t.elems { filePayload, ok := tfilepkg.FilePayloadOf(elem.File) @@ -68,7 +129,7 @@ func (taskCodec) Marshal(task core.Executable) ([]byte, error) { return json.Marshal(p) } -func (taskCodec) Unmarshal(data []byte) (core.Executable, error) { +func (batchCodec) Unmarshal(data []byte) (core.Executable, error) { var p taskPayload if err := json.Unmarshal(data, &p); err != nil { return nil, fmt.Errorf("invalid task payload: %w", err) @@ -77,8 +138,15 @@ func (taskCodec) Unmarshal(data []byte) (core.Executable, error) { if dler == nil { return nil, fmt.Errorf("no downloader client available") } + done := make(map[string]struct{}, len(p.Done)) + for _, id := range p.Done { + done[id] = struct{}{} + } elems := make([]TaskElement, 0, len(p.Elements)) for _, ep := range p.Elements { + if _, ok := done[ep.ID]; ok { + continue // upload already completed; do not re-run + } stor, err := storage.GetStorageByName(context.Background(), ep.Storage) if err != nil { return nil, fmt.Errorf("storage %q: %w", ep.Storage, err) @@ -102,5 +170,7 @@ func (taskCodec) Unmarshal(data []byte) (core.Executable, error) { if p.ChatID != 0 { progress = NewProgressTracker(p.MessageID, p.ChatID) } - return NewBatchTGFileTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors), nil + task := NewBatchTGFileTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors) + task.overwrite = p.Overwrite + return task, nil } diff --git a/core/tasks/batchtfile/codec_test.go b/core/tasks/batchtfile/codec_test.go new file mode 100644 index 0000000..f445e57 --- /dev/null +++ b/core/tasks/batchtfile/codec_test.go @@ -0,0 +1,39 @@ +package batchtfile + +import ( + "testing" +) + +func TestDetectTaskKind(t *testing.T) { + tests := []struct { + name string + payload string + want string + wantErr bool + }{ + {"batch with kind", `{"kind":"batch","id":"1","elements":[]}`, "batch", false}, + {"file with kind", `{"kind":"file","id":"1","file":{}}`, "file", false}, + {"legacy batch by shape", `{"id":"1","elements":[]}`, "batch", false}, + {"legacy file by shape", `{"id":"1","file":{}}`, "file", false}, + {"legacy batch with element", `{"id":"1","elements":[{"id":"e"}]}`, "batch", false}, + {"no discriminator", `{"id":"1"}`, "", true}, + {"invalid json", `not json`, "", true}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := detectTaskKind([]byte(tt.payload)) + if tt.wantErr { + if err == nil { + t.Fatalf("expected error, got kind %q", got) + } + return + } + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if got != tt.want { + t.Fatalf("kind = %q, want %q", got, tt.want) + } + }) + } +} diff --git a/core/tasks/batchtfile/execute.go b/core/tasks/batchtfile/execute.go index 4b027cf..a65da1d 100644 --- a/core/tasks/batchtfile/execute.go +++ b/core/tasks/batchtfile/execute.go @@ -35,6 +35,9 @@ func (g executionGroup) usesBatchSaver() bool { func (t *Task) Execute(ctx context.Context) error { logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("batch_file[%s]", t.ID)) logger.Info("Starting batch file task") + if t.overwrite { + ctx = storage.WithOverwrite(ctx) + } if t.Progress != nil { t.Progress.OnStart(ctx, t) } @@ -227,6 +230,9 @@ func (t *Task) processBatch(ctx context.Context, group executionGroup) error { err := t.saveBatchItems(ctx, successElems, items) if err == nil { uploaded = true + for _, elem := range successElems { + t.persistElementDone(ctx, elem.ID) + } } return err } @@ -430,6 +436,7 @@ func (t *Task) processElement(ctx context.Context, elem TaskElement) error { onProgress(fileStat.Size(), fileStat.Size()) t.markItemCompleted(elem.ID) t.notifyStateChange(vctx) + t.persistElementDone(ctx, elem.ID) success = true } else { t.markItemFailed(elem.ID, lastFailureStage, err) diff --git a/core/tasks/batchtfile/task.go b/core/tasks/batchtfile/task.go index ca05f52..17e60d4 100644 --- a/core/tasks/batchtfile/task.go +++ b/core/tasks/batchtfile/task.go @@ -2,11 +2,13 @@ package batchtfile import ( "context" + "encoding/json" "fmt" "path/filepath" "sync" "sync/atomic" + "github.com/charmbracelet/log" "github.com/krau/SaveAny-Bot/common/utils/tgutil" "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/core" @@ -47,6 +49,7 @@ type Task struct { uploadOnce sync.Once uploadMu sync.Mutex uploaded map[string]int64 + overwrite bool // recovered: overwrite storage targets instead of uniquifying } // Title implements core.Exectable. @@ -58,6 +61,41 @@ func (t *Task) Type() tasktype.TaskType { return tasktype.TaskTypeTgfiles } +// completedElementIDs returns the element IDs whose upload finished, for +// persisting upload progress so recovery can skip them. +func (t *Task) completedElementIDs() []string { + t.itemMu.RLock() + defer t.itemMu.RUnlock() + var ids []string + for _, item := range t.itemStates { + if item.phase == ItemPhaseCompleted { + ids = append(ids, item.id) + } + } + return ids +} + +// persistElementDone records an element's completed upload in the persisted +// payload so a restart does not re-upload it. +func (t *Task) persistElementDone(ctx context.Context, elemID string) { + err := core.UpdateTaskPayload(ctx, t.ID, func(payload []byte) ([]byte, error) { + var p taskPayload + if err := json.Unmarshal(payload, &p); err != nil { + return nil, err + } + for _, id := range p.Done { + if id == elemID { + return payload, nil + } + } + p.Done = append(p.Done, elemID) + return json.Marshal(p) + }) + if err != nil { + log.FromContext(ctx).Warnf("Failed to persist element completion %s: %v", elemID, err) + } +} + func NewTaskElement( stor storage.Storage, path string, diff --git a/core/tasks/tfile/codec.go b/core/tasks/tfile/codec.go index f3cc206..55ab13c 100644 --- a/core/tasks/tfile/codec.go +++ b/core/tasks/tfile/codec.go @@ -8,25 +8,28 @@ import ( "github.com/krau/SaveAny-Bot/config" "github.com/krau/SaveAny-Bot/core" - "github.com/krau/SaveAny-Bot/pkg/enums/tasktype" + "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" tfilepkg "github.com/krau/SaveAny-Bot/pkg/tfile" "github.com/krau/SaveAny-Bot/storage" ) type taskPayload struct { + Kind string `json:"kind"` // "file" ID string `json:"id"` Storage string `json:"storage"` Path string `json:"path"` File tfilepkg.FilePayload `json:"file"` ChatID int64 `json:"chat_id"` MessageID int `json:"message_id"` + Overwrite bool `json:"overwrite"` + Caption string `json:"caption"` } type taskCodec struct{} -func init() { - core.RegisterTaskCodec(tasktype.TaskTypeTgfiles, taskCodec{}) -} +// TaskCodec serializes single-file tasks. It is registered together with the +// batch codec under TaskTypeTgfiles (see core/tasks/batchtfile/codec.go). +var TaskCodec core.TaskCodec = taskCodec{} func (taskCodec) Marshal(task core.Executable) ([]byte, error) { t, ok := task.(*Task) @@ -38,11 +41,18 @@ func (taskCodec) Marshal(task core.Executable) ([]byte, error) { return nil, fmt.Errorf("file %T is not serializable", t.File) } p := taskPayload{ + Kind: "file", ID: t.ID, Storage: t.Storage.Name(), Path: t.Path, File: filePayload, } + if overwrite, ok := t.Ctx.Value(ctxkey.OverwriteExisting).(bool); ok { + p.Overwrite = overwrite + } + if caption, ok := sourceCaption(t.File); ok { + p.Caption = caption + } if progress, ok := t.Progress.(*Progress); ok { p.ChatID = progress.ChatID p.MessageID = progress.MessageID @@ -81,5 +91,7 @@ func (taskCodec) Unmarshal(data []byte) (core.Executable, error) { Progress: progress, stream: false, // recovered tasks always download to cache first localPath: localPath, + overwrite: p.Overwrite, + caption: p.Caption, }, nil } diff --git a/core/tasks/tfile/execute.go b/core/tasks/tfile/execute.go index 3f8db41..460da9d 100644 --- a/core/tasks/tfile/execute.go +++ b/core/tasks/tfile/execute.go @@ -25,6 +25,9 @@ func (t *Task) Execute(ctx context.Context) (err error) { t.Progress.OnDone(ctx, t, err) } }() + if t.overwrite { + ctx = storage.WithOverwrite(ctx) + } if t.Progress != nil { t.Progress.OnStart(ctx, t) } @@ -47,7 +50,9 @@ func (t *Task) Execute(ctx context.Context) (err error) { return fmt.Errorf("failed to get file stat: %w", err) } vctx := context.WithValue(ctx, ctxkey.ContentLength, fileStat.Size()) - if caption, ok := sourceCaption(t.File); ok { + if t.caption != "" { + vctx = storagetypes.WithSourceCaption(vctx, t.caption) + } else if caption, ok := sourceCaption(t.File); ok { vctx = storagetypes.WithSourceCaption(vctx, caption) } err = retry.Retry(func() error { diff --git a/core/tasks/tfile/tftask.go b/core/tasks/tfile/tftask.go index b85585c..c6d205b 100644 --- a/core/tasks/tfile/tftask.go +++ b/core/tasks/tfile/tftask.go @@ -23,6 +23,8 @@ type Task struct { Progress ProgressTracker stream bool // true if the file should be downloaded in stream mode localPath string + overwrite bool // recovered: overwrite the storage target instead of uniquifying + caption string // recovered: source caption for the telegram backend } // Title implements core.Exectable. diff --git a/core/tasks/transfer/codec.go b/core/tasks/transfer/codec.go index 06740b8..499fc9d 100644 --- a/core/tasks/transfer/codec.go +++ b/core/tasks/transfer/codec.go @@ -6,11 +6,17 @@ import ( "fmt" "github.com/krau/SaveAny-Bot/core" + "github.com/krau/SaveAny-Bot/pkg/enums/ctxkey" "github.com/krau/SaveAny-Bot/pkg/enums/tasktype" "github.com/krau/SaveAny-Bot/pkg/storagetypes" "github.com/krau/SaveAny-Bot/storage" ) +func ctxOverwrite(ctx context.Context) bool { + overwrite, _ := ctx.Value(ctxkey.OverwriteExisting).(bool) + return overwrite +} + type elementPayload struct { ID string `json:"id"` SourceStorage string `json:"source_storage"` @@ -26,6 +32,7 @@ type taskPayload struct { ChatID int64 `json:"chat_id"` MessageID int `json:"message_id"` IgnoreErrors bool `json:"ignore_errors"` + Overwrite bool `json:"overwrite"` } type taskCodec struct{} @@ -42,6 +49,7 @@ func (taskCodec) Marshal(task core.Executable) ([]byte, error) { p := taskPayload{ ID: t.ID, IgnoreErrors: t.IgnoreErrors, + Overwrite: ctxOverwrite(t.ctx), } for _, elem := range t.elems { p.Elements = append(p.Elements, elementPayload{ @@ -88,5 +96,7 @@ func (taskCodec) Unmarshal(data []byte) (core.Executable, error) { if p.ChatID != 0 { progress = NewProgressTracker(p.MessageID, p.ChatID) } - return NewTransferTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors), nil + task := NewTransferTask(p.ID, context.Background(), elems, progress, p.IgnoreErrors) + task.overwrite = p.Overwrite + return task, nil } diff --git a/core/tasks/transfer/execute.go b/core/tasks/transfer/execute.go index b517f05..555ee81 100644 --- a/core/tasks/transfer/execute.go +++ b/core/tasks/transfer/execute.go @@ -21,6 +21,9 @@ import ( func (t *Task) Execute(ctx context.Context) error { logger := log.FromContext(ctx).WithPrefix(fmt.Sprintf("transfer[%s]", t.ID)) logger.Info("Starting transfer task") + if t.overwrite { + ctx = storage.WithOverwrite(ctx) + } if t.Progress != nil { t.Progress.OnStart(ctx, t) } diff --git a/core/tasks/transfer/task.go b/core/tasks/transfer/task.go index 85e2bde..fe75e11 100644 --- a/core/tasks/transfer/task.go +++ b/core/tasks/transfer/task.go @@ -35,6 +35,7 @@ type Task struct { processing map[string]TaskElementInfo processingMu sync.RWMutex failed map[string]error + overwrite bool // recovered: overwrite storage targets instead of uniquifying } // Title implements core.Executable. diff --git a/database/task.go b/database/task.go index c03c62d..6c110d3 100644 --- a/database/task.go +++ b/database/task.go @@ -58,6 +58,51 @@ func UpdateTaskStatus(ctx context.Context, id string, status TaskStatus, errMsg }).Error } +func GetTask(ctx context.Context, id string) (*Task, error) { + if db == nil { + return nil, errNotInitialized + } + var task Task + if err := db.WithContext(ctx).First(&task, "id = ?", id).Error; err != nil { + return nil, err + } + return &task, nil +} + +// UpdateTaskPayload replaces the payload of an existing task row. +func UpdateTaskPayload(ctx context.Context, id string, payload []byte) error { + if db == nil { + return errNotInitialized + } + return db.WithContext(ctx).Model(&Task{}). + Where("id = ?", id). + Updates(map[string]any{ + "payload": payload, + "updated_at": time.Now(), + }).Error +} + +// RestoreTaskCreatedAt restores the original creation time after a +// re-enqueue overwrote it. +func RestoreTaskCreatedAt(ctx context.Context, id string, createdAt time.Time) error { + if db == nil { + return errNotInitialized + } + return db.WithContext(ctx).Model(&Task{}). + Where("id = ?", id). + Update("created_at", createdAt).Error +} + +// DeleteStaleFailedTasks removes failed rows older than the given age. +func DeleteStaleFailedTasks(ctx context.Context, maxAge time.Duration) error { + if db == nil { + return errNotInitialized + } + return db.WithContext(ctx). + Where("status = ? AND updated_at < ?", string(TaskStatusFailed), time.Now().Add(-maxAge)). + Delete(&Task{}).Error +} + func DeleteTask(ctx context.Context, id string) error { if db == nil { return errNotInitialized diff --git a/pkg/queue/queue.go b/pkg/queue/queue.go index 47c248d..b0a284c 100644 --- a/pkg/queue/queue.go +++ b/pkg/queue/queue.go @@ -116,6 +116,22 @@ func (tq *TaskQueue[T]) ActiveLength() int { return count } +// Contains reports whether a task with the given ID is queued or running. +func (tq *TaskQueue[T]) Contains(taskID string) bool { + tq.mu.RLock() + defer tq.mu.RUnlock() + if _, ok := tq.runningTaskMap[taskID]; ok { + return true + } + for element := tq.tasks.Front(); element != nil; element = element.Next() { + task := element.Value.(*Task[T]) + if task.ID == taskID && !task.Cancelled() { + return true + } + } + return false +} + // RunningTasks returns the currently running tasks' info. func (tq *TaskQueue[T]) RunningTasks() []TaskInfo { tq.mu.RLock()