refactor(downloader): inspect runtime task state

This commit is contained in:
saltbo
2026-06-05 23:01:01 -04:00
parent e113c6c2cf
commit f4acfc15ef
8 changed files with 576 additions and 145 deletions
+2 -2
View File
@@ -33,7 +33,7 @@ func TestCreateObjectUsesRenameConflictStrategy(t *testing.T) {
}
}
func TestAssignedTasksFetchesRecoverableStatuses(t *testing.T) {
func TestAssignedTasksFetchesRunnableStatuses(t *testing.T) {
var statuses []string
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.URL.Path != "/api/download-tasks" {
@@ -55,7 +55,7 @@ func TestAssignedTasksFetchesRecoverableStatuses(t *testing.T) {
sort.Strings(statuses)
expected := []string{"assigned", "downloading", "interrupted", "uploading"}
if !reflect.DeepEqual(statuses, expected) {
t.Fatalf("expected recoverable statuses %v, got %v", expected, statuses)
t.Fatalf("expected runnable statuses %v, got %v", expected, statuses)
}
if len(tasks) != 4 {
t.Fatalf("expected four tasks, got %d", len(tasks))
+63 -33
View File
@@ -155,40 +155,17 @@ func (a Aria2) RestoreSeed(ctx context.Context, ref SeedRef) (*Seed, error) {
}, nil
}
func (a Aria2) Recover(ctx context.Context, task client.DownloadTask) (Result, bool, error) {
func (a Aria2) InspectTask(ctx context.Context, task client.DownloadTask) (TaskSnapshot, bool, error) {
aria, err := a.client(ctx)
if err == nil {
defer aria.Close()
status, ok, findErr := a.findTask(ctx, &aria, task)
if findErr != nil {
return Result{}, false, findErr
}
if ok && isAria2DownloadComplete(status) {
files, err := a.getAria2Files(ctx, &aria, status.GID)
if err != nil {
return Result{}, false, err
}
taskDir := aria2StatusTaskDir(status, filepath.Join(a.Dir, task.ID))
result, err := resultFromAria2Files(task, taskDir, status.BitTorrent.Info.Name, files)
if err == nil && a.RetainSeed && task.SourceType != "http" {
result.Seed = a.seedFromStatus(status, taskDir)
}
return result, err == nil, err
}
if ok {
return Result{}, false, fmt.Errorf(
"aria2 task %s is not a completed upload resume candidate: status=%s completed=%d total=%d",
status.GID,
status.Status,
int64(status.CompletedLength),
int64(status.TotalLength),
)
}
}
if err != nil {
return Result{}, false, err
return TaskSnapshot{}, false, err
}
return Result{}, false, nil
defer aria.Close()
status, ok, err := a.findTask(ctx, &aria, task)
if err != nil || !ok {
return TaskSnapshot{}, ok, err
}
return a.snapshotTask(ctx, &aria, task, status)
}
func (a Aria2) Download(ctx context.Context, task client.DownloadTask, progress Progress) (Result, error) {
@@ -203,7 +180,7 @@ func (a Aria2) Download(ctx context.Context, task client.DownloadTask, progress
}
defer aria.Close()
if shouldRecoverExistingAria2Task(task) {
if shouldAttachExistingAria2Task(task) {
status, ok, err := a.findTask(ctx, &aria, task)
if err != nil {
return Result{}, fmt.Errorf("find aria2 task: %w", err)
@@ -265,10 +242,63 @@ func (a Aria2) Download(ctx context.Context, task client.DownloadTask, progress
return result, nil
}
func shouldRecoverExistingAria2Task(task client.DownloadTask) bool {
func shouldAttachExistingAria2Task(task client.DownloadTask) bool {
return task.Status == "downloading" || task.Status == "uploading"
}
func (a Aria2) snapshotTask(
ctx context.Context,
aria **arigo.Client,
task client.DownloadTask,
status arigo.Status,
) (TaskSnapshot, bool, error) {
total := int64(status.TotalLength)
completed := int64(status.CompletedLength)
bps := int64(status.DownloadSpeed)
var totalPtr *int64
if total > 0 {
totalPtr = &total
}
peers := a.getAria2Peers(ctx, aria, status.GID)
snapshot := TaskSnapshot{
State: aria2TaskState(status),
Downloaded: completed,
Total: totalPtr,
Bps: bps,
Detail: aria2Detail(status, peers),
Error: status.ErrorMessage,
}
if snapshot.State != TaskStateCompleted {
return snapshot, true, nil
}
files, err := a.getAria2Files(ctx, aria, status.GID)
if err != nil {
return TaskSnapshot{}, false, err
}
taskDir := aria2StatusTaskDir(status, filepath.Join(a.Dir, task.ID))
result, err := resultFromAria2Files(task, taskDir, status.BitTorrent.Info.Name, files)
if err != nil {
return TaskSnapshot{}, false, err
}
if a.RetainSeed && task.SourceType != "http" {
result.Seed = a.seedFromStatus(status, taskDir)
}
snapshot.Result = &result
return snapshot, true, nil
}
func aria2TaskState(status arigo.Status) TaskState {
if isAria2DownloadComplete(status) {
return TaskStateCompleted
}
switch string(status.Status) {
case string(arigo.StatusError), string(arigo.StatusRemoved):
return TaskStateFailed
default:
return TaskStateDownloading
}
}
func addAria2Task(ctx context.Context, aria *arigo.Client, task client.DownloadTask, options *arigo.Options) (arigo.GID, error) {
if task.SourceType != "torrent_url" {
return aria.AddURI(arigo.URIs(task.SourceURI), options)
+19 -1
View File
@@ -45,11 +45,29 @@ type SeedSnapshot struct {
type Progress func(downloaded int64, total *int64, bps int64, detail *client.DownloadTaskDetail) error
type TaskState string
const (
TaskStateDownloading TaskState = "downloading"
TaskStateCompleted TaskState = "completed"
TaskStateFailed TaskState = "failed"
)
type TaskSnapshot struct {
State TaskState
Downloaded int64
Total *int64
Bps int64
Detail *client.DownloadTaskDetail
Result *Result
Error string
}
type Engine interface {
Name() string
Capabilities() []string
Check(ctx context.Context) error
Recover(ctx context.Context, task client.DownloadTask) (Result, bool, error)
InspectTask(ctx context.Context, task client.DownloadTask) (TaskSnapshot, bool, error)
Download(ctx context.Context, task client.DownloadTask, progress Progress) (Result, error)
}
+109 -7
View File
@@ -12,6 +12,7 @@ import (
"testing"
"github.com/Braurbeki/arigo"
qbittorrent "github.com/autobrr/go-qbittorrent"
"github.com/cenkalti/rpc2"
"github.com/saltbo/zpan/downloader/internal/client"
)
@@ -116,7 +117,40 @@ func TestHTTPDownloadResumesExistingFile(t *testing.T) {
}
}
func TestHTTPRecoverDoesNotTrustLocalTaskDirectory(t *testing.T) {
func TestHTTPInspectTaskUsesCompletedCheckpoint(t *testing.T) {
dir := t.TempDir()
taskDir := filepath.Join(dir, "task-1")
if err := os.MkdirAll(taskDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(taskDir, "payload.bin"), []byte("payload"), 0o644); err != nil {
t.Fatal(err)
}
total := int64(7)
snapshot, found, err := (HTTP{Dir: dir}).InspectTask(context.Background(), client.DownloadTask{
ID: "task-1",
SourceType: "http",
SourceURI: "https://example.com/payload.bin",
DownloadedBytes: 7,
TotalBytes: &total,
})
if err != nil {
t.Fatal(err)
}
if !found {
t.Fatal("expected completed http checkpoint to be found")
}
if snapshot.State != TaskStateCompleted {
t.Fatalf("expected completed state, got %#v", snapshot)
}
if snapshot.Result == nil || snapshot.Result.Path != filepath.Join(taskDir, "payload.bin") || snapshot.Result.Size != 7 {
t.Fatalf("unexpected completed result: %#v", snapshot.Result)
}
}
func TestHTTPInspectTaskDoesNotTrustLocalFileWithoutCompletedCheckpoint(t *testing.T) {
dir := t.TempDir()
taskDir := filepath.Join(dir, "task-1")
if err := os.MkdirAll(taskDir, 0o755); err != nil {
@@ -126,7 +160,7 @@ func TestHTTPRecoverDoesNotTrustLocalTaskDirectory(t *testing.T) {
t.Fatal(err)
}
result, recovered, err := (HTTP{Dir: dir}).Recover(context.Background(), client.DownloadTask{
snapshot, found, err := (HTTP{Dir: dir}).InspectTask(context.Background(), client.DownloadTask{
ID: "task-1",
SourceType: "http",
SourceURI: "https://example.com/payload.bin",
@@ -135,8 +169,32 @@ func TestHTTPRecoverDoesNotTrustLocalTaskDirectory(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if recovered {
t.Fatalf("expected local directory not to be treated as a completed runtime result, got %#v", result)
if found {
t.Fatalf("expected local file without checkpoint not to be trusted, got %#v", snapshot)
}
}
func TestHTTPInspectTaskRejectsSizeMismatch(t *testing.T) {
dir := t.TempDir()
taskDir := filepath.Join(dir, "task-1")
if err := os.MkdirAll(taskDir, 0o755); err != nil {
t.Fatal(err)
}
if err := os.WriteFile(filepath.Join(taskDir, "payload.bin"), []byte("payload"), 0o644); err != nil {
t.Fatal(err)
}
total := int64(8)
_, _, err := (HTTP{Dir: dir}).InspectTask(context.Background(), client.DownloadTask{
ID: "task-1",
SourceType: "http",
SourceURI: "https://example.com/payload.bin",
DownloadedBytes: 8,
TotalBytes: &total,
})
if err == nil || !strings.Contains(err.Error(), "size mismatch") {
t.Fatalf("expected size mismatch error, got %v", err)
}
}
@@ -294,17 +352,61 @@ func TestIsAria2DownloadCompleteTreatsActiveFullTorrentAsComplete(t *testing.T)
}
if !isAria2DownloadComplete(status) {
t.Fatal("expected active full torrent to be recoverable as a completed download")
t.Fatal("expected active full torrent to be treated as completed")
}
}
func TestAria2TaskState(t *testing.T) {
if got := aria2TaskState(arigo.Status{
Status: arigo.StatusActive,
TotalLength: 100,
CompletedLength: 100,
Files: []arigo.File{{Path: "/tmp/zpan/task-1/file.bin", Length: 100, CompletedLength: 100}},
}); got != TaskStateCompleted {
t.Fatalf("expected active full torrent to be completed, got %s", got)
}
if got := aria2TaskState(arigo.Status{Status: arigo.StatusActive, TotalLength: 100, CompletedLength: 10}); got != TaskStateDownloading {
t.Fatalf("expected active partial torrent to be downloading, got %s", got)
}
if got := aria2TaskState(arigo.Status{Status: arigo.StatusError}); got != TaskStateFailed {
t.Fatalf("expected error torrent to be failed, got %s", got)
}
}
func TestQBittorrentTaskState(t *testing.T) {
if got := qbittorrentTaskState(qbittorrent.Torrent{
State: qbittorrent.TorrentState("stalledUP"),
Progress: 1,
AmountLeft: 0,
TotalSize: 100,
}); got != TaskStateCompleted {
t.Fatalf("expected seeding torrent to be completed, got %s", got)
}
if got := qbittorrentTaskState(qbittorrent.Torrent{
State: qbittorrent.TorrentState("downloading"),
Progress: 0.5,
AmountLeft: 50,
TotalSize: 100,
}); got != TaskStateDownloading {
t.Fatalf("expected partial torrent to be downloading, got %s", got)
}
if got := qbittorrentTaskState(qbittorrent.Torrent{
State: qbittorrent.TorrentState("missingFiles"),
Progress: 0.5,
AmountLeft: 50,
TotalSize: 100,
}); got != TaskStateFailed {
t.Fatalf("expected missing files torrent to be failed, got %s", got)
}
}
func TestIsAria2InfoHashAlreadyRegistered(t *testing.T) {
err := errors.New("InfoHash 0546769f209ec059284b47f68659791a6f75ca8e is already registered.")
if !isAria2InfoHashAlreadyRegistered(err) {
t.Fatal("expected aria2 infohash conflict to be recoverable")
t.Fatal("expected aria2 infohash conflict to be attachable")
}
if isAria2InfoHashAlreadyRegistered(errors.New("aria2 download ended with status error")) {
t.Fatal("expected ordinary aria2 error to stay non-recoverable")
t.Fatal("expected ordinary aria2 error not to be treated as an infohash conflict")
}
}
+67 -5
View File
@@ -3,8 +3,10 @@ package engine
import (
"context"
"errors"
"fmt"
"io"
"net/http"
"net/url"
"os"
"path/filepath"
"strconv"
@@ -40,8 +42,36 @@ func (h HTTP) Check(ctx context.Context) error {
return os.Remove(path)
}
func (h HTTP) Recover(ctx context.Context, task client.DownloadTask) (Result, bool, error) {
return Result{}, false, nil
func (h HTTP) InspectTask(ctx context.Context, task client.DownloadTask) (TaskSnapshot, bool, error) {
if task.SourceType != "http" {
return TaskSnapshot{}, false, nil
}
size, ok := completedHTTPCheckpoint(task)
if !ok {
return TaskSnapshot{}, false, nil
}
path, name, err := h.outputPath(task)
if err != nil {
return TaskSnapshot{}, false, err
}
info, err := os.Stat(path)
if err != nil {
return TaskSnapshot{}, false, fmt.Errorf("inspect http completed file: %w", err)
}
if info.IsDir() {
return TaskSnapshot{}, false, fmt.Errorf("inspect http completed file: %s is a directory", path)
}
if info.Size() != size {
return TaskSnapshot{}, false, fmt.Errorf("inspect http completed file: size mismatch path=%s expected=%d actual=%d", path, size, info.Size())
}
result := Result{Path: path, Name: name, Size: size}
return TaskSnapshot{
State: TaskStateCompleted,
Downloaded: size,
Total: &size,
Detail: &client.DownloadTaskDetail{Engine: "builtin", Phase: "completed"},
Result: &result,
}, true, nil
}
func (h HTTP) Download(ctx context.Context, task client.DownloadTask, progress Progress) (Result, error) {
@@ -53,16 +83,18 @@ func (h HTTP) Download(ctx context.Context, task client.DownloadTask, progress P
return Result{}, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, task.SourceURI, nil)
path, name, err := h.outputPath(task)
if err != nil {
return Result{}, err
}
name := outputName(task, filenameFromURL(req.URL))
path := filepath.Join(taskDir, name)
existingSize, err := existingFileSize(path)
if err != nil {
return Result{}, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodGet, task.SourceURI, nil)
if err != nil {
return Result{}, err
}
if existingSize > 0 {
req.Header.Set("Range", "bytes="+strconv.FormatInt(existingSize, 10)+"-")
}
@@ -104,6 +136,36 @@ func (h HTTP) Download(ctx context.Context, task client.DownloadTask, progress P
return Result{Path: path, Name: name, Size: counter.downloaded}, nil
}
func (h HTTP) outputPath(task client.DownloadTask) (string, string, error) {
parsed, err := httpURL(task.SourceURI)
if err != nil {
return "", "", err
}
name := outputName(task, filenameFromURL(parsed))
return filepath.Join(h.Dir, task.ID, name), name, nil
}
func completedHTTPCheckpoint(task client.DownloadTask) (int64, bool) {
if task.TotalBytes == nil || *task.TotalBytes <= 0 {
return 0, false
}
if task.DownloadedBytes != *task.TotalBytes {
return 0, false
}
return *task.TotalBytes, true
}
func httpURL(raw string) (*url.URL, error) {
parsed, err := url.Parse(raw)
if err != nil {
return nil, err
}
if parsed.Scheme != "http" && parsed.Scheme != "https" {
return nil, fmt.Errorf("unsupported http source url: %s", raw)
}
return parsed, nil
}
func existingFileSize(path string) (int64, error) {
info, err := os.Stat(path)
if err == nil {
+62 -20
View File
@@ -115,33 +115,19 @@ func (q QBittorrent) RestoreSeed(ctx context.Context, ref SeedRef) (*Seed, error
}, nil
}
func (q QBittorrent) Recover(ctx context.Context, task client.DownloadTask) (Result, bool, error) {
func (q QBittorrent) InspectTask(ctx context.Context, task client.DownloadTask) (TaskSnapshot, bool, error) {
if task.SourceType == "http" {
return HTTP{Dir: q.Dir}.Recover(ctx, task)
return HTTP{Dir: q.Dir}.InspectTask(ctx, task)
}
qbt, err := q.login(ctx)
if err != nil {
return Result{}, false, err
return TaskSnapshot{}, false, err
}
torrent, ok, err := q.findTask(ctx, qbt, task)
if err != nil {
return Result{}, false, err
if err != nil || !ok {
return TaskSnapshot{}, ok, err
}
if ok && (torrent.Progress >= 1 || (torrent.AmountLeft == 0 && torrent.TotalSize > 0)) {
result, err := resultFromQBittorrentFiles(ctx, qbt, task, filepath.Join(q.Dir, task.ID), torrent)
return result, err == nil, err
}
if ok {
return Result{}, false, fmt.Errorf(
"qbittorrent task %s is not a completed upload resume candidate: state=%s progress=%f amount_left=%d total=%d",
torrent.Hash,
torrent.State,
torrent.Progress,
torrent.AmountLeft,
torrent.TotalSize,
)
}
return Result{}, false, nil
return q.snapshotTask(ctx, qbt, task, torrent)
}
func (q QBittorrent) login(ctx context.Context) (*qbittorrent.Client, error) {
@@ -223,6 +209,62 @@ func (q QBittorrent) resultFromTorrent(
return result, nil
}
func (q QBittorrent) snapshotTask(
ctx context.Context,
qbt *qbittorrent.Client,
task client.DownloadTask,
torrent qbittorrent.Torrent,
) (TaskSnapshot, bool, error) {
total := torrent.TotalSize
if total <= 0 {
total = torrent.Size
}
var totalPtr *int64
if total > 0 {
totalPtr = &total
}
snapshot := TaskSnapshot{
State: qbittorrentTaskState(torrent),
Downloaded: torrent.Completed,
Total: totalPtr,
Bps: torrent.DlSpeed,
Detail: qbittorrentDetail(ctx, qbt, torrent),
}
if snapshot.State != TaskStateCompleted {
return snapshot, true, nil
}
result, err := resultFromQBittorrentFiles(ctx, qbt, task, filepath.Join(q.Dir, task.ID), torrent)
if err != nil {
return TaskSnapshot{}, false, err
}
if q.RetainSeed {
result.Seed = &Seed{
Engine: "qbittorrent",
ID: torrent.Hash,
InfoHash: torrent.Hash,
Path: filepath.Join(q.Dir, task.ID),
Snapshot: q.seedSnapshot(torrent.Hash),
Cleanup: q.cleanupSeed(torrent.Hash, filepath.Join(q.Dir, task.ID)),
}
}
snapshot.Result = &result
return snapshot, true, nil
}
func qbittorrentTaskState(torrent qbittorrent.Torrent) TaskState {
total := torrent.TotalSize
if total <= 0 {
total = torrent.Size
}
if torrent.Progress >= 1 || (torrent.AmountLeft == 0 && total > 0) {
return TaskStateCompleted
}
if isQBittorrentErrorState(torrent.State) {
return TaskStateFailed
}
return TaskStateDownloading
}
func (q QBittorrent) findTask(ctx context.Context, qbt *qbittorrent.Client, task client.DownloadTask) (qbittorrent.Torrent, bool, error) {
tag := qbittorrentTrackingTag(task.ID)
torrents, err := qbt.GetTorrentsCtx(ctx, qbittorrent.TorrentFilterOptions{Tag: tag})
+74 -31
View File
@@ -9,6 +9,7 @@ import (
"os/exec"
"path/filepath"
"runtime"
"runtime/debug"
"sort"
"strings"
"sync"
@@ -26,11 +27,11 @@ var errBillingPaused = errors.New("billing paused")
var errTaskPausing = errors.New("task pausing")
var errTaskCanceling = errors.New("task canceling")
type taskResumeStage int
type taskWorkStage int
const (
taskResumeDownload taskResumeStage = iota
taskResumeUpload
taskWorkStageDownload taskWorkStage = iota
taskWorkStageUploadExistingResult
)
type Worker struct {
@@ -185,32 +186,23 @@ func (w *Worker) tick(ctx context.Context) error {
func (w *Worker) process(ctx context.Context, task client.DownloadTask) {
defer w.finish(task.ID)
log := w.taskLogger(task)
defer w.recoverTaskPanic(ctx, task.ID, log)
log.Info("task started", "source_uri", task.SourceURI, "target_folder", task.TargetFolder)
currentDetail := task.Detail
stage := resumeStage(task)
if stage == taskResumeUpload {
result, recovered, err := w.engine.Recover(ctx, task)
if err != nil {
msg := taskErrorMessage(err)
log.Error("failed to recover completed download result", "error", err)
if _, updateErr := w.updateTask(ctx, task.ID, client.TaskPatch{Status: "failed", ErrorMessage: &msg}); updateErr != nil {
log.Error("failed to mark task failed", "error", updateErr)
}
return
}
if recovered {
log.Info("recovered completed download result", "path", result.Path, "name", result.Name, "size", result.Size)
w.uploadAndComplete(ctx, log, task, result, currentDetail)
return
}
msg := "completed download is not recoverable from downloader runtime"
log.Error("failed to recover completed download result", "error", msg)
if _, updateErr := w.updateTask(ctx, task.ID, client.TaskPatch{Status: "failed", ErrorMessage: &msg}); updateErr != nil {
log.Error("failed to mark task failed", "error", updateErr)
}
if nextTaskWorkStage(task) == taskWorkStageUploadExistingResult {
w.uploadExistingResult(ctx, log, task, currentDetail)
return
}
w.downloadThenUpload(ctx, log, task, currentDetail)
}
func (w *Worker) downloadThenUpload(
ctx context.Context,
log *slog.Logger,
task client.DownloadTask,
currentDetail *client.DownloadTaskDetail,
) {
if _, err := w.updateTask(ctx, task.ID, client.TaskPatch{Status: "downloading"}); err != nil {
log.Error("failed to mark task downloading", "error", err)
if w.resolveControlledTaskUpdate(ctx, task.ID, err, log) {
@@ -295,23 +287,74 @@ func (w *Worker) process(ctx context.Context, task client.DownloadTask) {
w.uploadAndComplete(ctx, log, task, result, currentDetail)
}
func resumeStage(task client.DownloadTask) taskResumeStage {
func (w *Worker) uploadExistingResult(
ctx context.Context,
log *slog.Logger,
task client.DownloadTask,
currentDetail *client.DownloadTaskDetail,
) {
snapshot, found, err := w.engine.InspectTask(ctx, task)
if err != nil {
msg := taskErrorMessage(err)
log.Error("failed to inspect downloader runtime task", "error", err)
if _, updateErr := w.updateTask(ctx, task.ID, client.TaskPatch{Status: "failed", ErrorMessage: &msg}); updateErr != nil {
log.Error("failed to mark task failed", "error", updateErr)
}
return
}
if !found {
msg := "download task is missing from downloader runtime"
log.Error("failed to inspect downloader runtime task", "error", msg)
if _, updateErr := w.updateTask(ctx, task.ID, client.TaskPatch{Status: "failed", ErrorMessage: &msg}); updateErr != nil {
log.Error("failed to mark task failed", "error", updateErr)
}
return
}
if snapshot.State != engine.TaskStateCompleted {
panic(fmt.Errorf(
"server task requires upload but runtime task is not completed: runtime_state=%s downloaded=%d total=%v",
snapshot.State,
snapshot.Downloaded,
optionalInt64(snapshot.Total),
))
}
if snapshot.Result == nil {
panic("runtime reported completed task without a local result")
}
log.Info("using completed runtime result", "path", snapshot.Result.Path, "name", snapshot.Result.Name, "size", snapshot.Result.Size)
w.uploadAndComplete(ctx, log, task, *snapshot.Result, currentDetail)
}
func (w *Worker) recoverTaskPanic(ctx context.Context, taskID string, log *slog.Logger) {
value := recover()
if value == nil {
return
}
err := fmt.Errorf("panic: %v", value)
msg := taskErrorMessage(err)
log.Error("task panicked", "panic", value, "stack", string(debug.Stack()))
if _, updateErr := w.updateTask(context.WithoutCancel(ctx), taskID, client.TaskPatch{Status: "failed", ErrorMessage: &msg}); updateErr != nil {
log.Error("failed to mark task failed after panic", "error", updateErr)
}
}
func nextTaskWorkStage(task client.DownloadTask) taskWorkStage {
if task.Status == "uploading" {
return taskResumeUpload
return taskWorkStageUploadExistingResult
}
if task.Status != "assigned" && task.Status != "downloading" && task.Status != "interrupted" {
return taskResumeDownload
return taskWorkStageDownload
}
if task.StorageUploadedBytes > 0 {
return taskResumeUpload
return taskWorkStageUploadExistingResult
}
if task.Detail != nil && (task.Detail.Phase == "uploading" || task.Detail.Phase == "completed") {
return taskResumeUpload
return taskWorkStageUploadExistingResult
}
if task.TotalBytes != nil && *task.TotalBytes > 0 && task.DownloadedBytes >= *task.TotalBytes {
return taskResumeUpload
return taskWorkStageUploadExistingResult
}
return taskResumeDownload
return taskWorkStageDownload
}
func (w *Worker) uploadAndComplete(
+180 -46
View File
@@ -9,6 +9,7 @@ import (
"net/http/httptest"
"os"
"path/filepath"
"strconv"
"strings"
"testing"
"time"
@@ -205,8 +206,11 @@ func TestWorkerLifecycleRetriesUploadWithoutRedownloading(t *testing.T) {
}
eng := &recordingEngine{
downloadResult: engine.Result{Path: payloadPath, Name: "payload.bin", Size: payloadSize},
recoverResult: engine.Result{Path: payloadPath, Name: "payload.bin", Size: payloadSize},
recovered: true,
taskSnapshot: engine.TaskSnapshot{
State: engine.TaskStateCompleted,
Result: &engine.Result{Path: payloadPath, Name: "payload.bin", Size: payloadSize},
},
taskFound: true,
}
first := NewWithAPI(config.Config{}, api)
@@ -244,8 +248,8 @@ func TestWorkerLifecycleRetriesUploadWithoutRedownloading(t *testing.T) {
if eng.downloadCalls != 1 {
t.Fatalf("expected retry to avoid a second download, got %d download calls", eng.downloadCalls)
}
if eng.recoverCalls != 1 {
t.Fatalf("expected retry to recover the completed download, got %d recover calls", eng.recoverCalls)
if eng.inspectCalls != 1 {
t.Fatalf("expected retry to inspect the runtime task, got %d inspect calls", eng.inspectCalls)
}
if uploadRequests != 2 {
t.Fatalf("expected both attempts to upload the local result, got %d upload requests", uploadRequests)
@@ -256,35 +260,80 @@ func TestWorkerLifecycleRetriesUploadWithoutRedownloading(t *testing.T) {
}
}
func TestUploadResumeRecoveryErrorFailsWithoutRedownloading(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{recoverErr: errors.New("runtime state is inconsistent")}
w := NewWithAPI(config.Config{}, api)
w.engine = eng
func TestWorkerLifecycleRetriesHTTPUploadFromCheckpointWithoutRedownloading(t *testing.T) {
payload := "downloaded payload"
downloadRequests := 0
downloadServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
downloadRequests++
w.Header().Set("Content-Length", strconv.Itoa(len(payload)))
_, _ = w.Write([]byte(payload))
}))
defer downloadServer.Close()
total := int64(100)
w.process(context.Background(), client.DownloadTask{
uploadRequests := 0
uploadServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
uploadRequests++
if _, err := io.Copy(io.Discard, r.Body); err != nil {
t.Fatal(err)
}
w.WriteHeader(http.StatusOK)
}))
defer uploadServer.Close()
api := &recordingAPI{
createObjectDraft: client.ObjectDraft{ID: "object-1", Name: "payload.bin", UploadURL: uploadServer.URL},
confirmErrs: []error{errors.New("unauthorized"), nil},
}
downloadDir := t.TempDir()
payloadSize := int64(len(payload))
first := NewWithAPI(config.Config{}, api)
first.engine = engine.HTTP{Dir: downloadDir}
first.process(context.Background(), client.DownloadTask{
ID: "task-1",
Status: "assigned",
SourceType: "magnet",
SourceURI: "magnet:?xt=urn:btih:abc123",
TotalBytes: &total,
Detail: &client.DownloadTaskDetail{Phase: "uploading"},
SourceType: "http",
SourceURI: downloadServer.URL + "/payload.bin",
Name: "payload.bin",
UploadToken: "upload-token",
})
if eng.downloadCalls != 0 {
t.Fatalf("expected upload resume recovery failure not to restart download, got %d download calls", eng.downloadCalls)
}
failed := lastPatchWithStatus(t, api.patches, "failed")
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "runtime state is inconsistent") {
t.Fatalf("expected recovery error to be reported, got %#v", failed.ErrorMessage)
if failed.DownloadedBytes == nil || *failed.DownloadedBytes != payloadSize {
t.Fatalf("expected failed task to persist downloaded bytes %d, got %#v", payloadSize, failed.DownloadedBytes)
}
if downloadRequests != 1 {
t.Fatalf("expected initial attempt to download once, got %d requests", downloadRequests)
}
second := NewWithAPI(config.Config{}, api)
second.engine = engine.HTTP{Dir: downloadDir}
second.process(context.Background(), client.DownloadTask{
ID: "task-1",
Status: "assigned",
SourceType: "http",
SourceURI: downloadServer.URL + "/payload.bin",
Name: "payload.bin",
DownloadedBytes: payloadSize,
TotalBytes: &payloadSize,
Detail: &client.DownloadTaskDetail{Phase: "uploading"},
UploadToken: "upload-token",
})
if downloadRequests != 1 {
t.Fatalf("expected retry not to request download source again, got %d requests", downloadRequests)
}
if uploadRequests != 2 {
t.Fatalf("expected both attempts to upload the local file, got %d upload requests", uploadRequests)
}
last := api.patches[len(api.patches)-1]
if last.Status != "completed" {
t.Fatalf("expected retry to complete task, got last patch %#v", last)
}
}
func TestUploadResumeMissingRuntimeFailsWithoutRedownloading(t *testing.T) {
func TestUploadExistingResultInspectErrorFailsWithoutRedownloading(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{recovered: false}
eng := &recordingEngine{inspectErr: errors.New("runtime state is inconsistent")}
w := NewWithAPI(config.Config{}, api)
w.engine = eng
@@ -300,14 +349,95 @@ func TestUploadResumeMissingRuntimeFailsWithoutRedownloading(t *testing.T) {
})
if eng.downloadCalls != 0 {
t.Fatalf("expected missing runtime during upload resume not to restart download, got %d download calls", eng.downloadCalls)
t.Fatalf("expected runtime inspection failure not to restart download, got %d download calls", eng.downloadCalls)
}
failed := lastPatchWithStatus(t, api.patches, "failed")
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "not recoverable") {
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "runtime state is inconsistent") {
t.Fatalf("expected runtime inspection error to be reported, got %#v", failed.ErrorMessage)
}
}
func TestUploadExistingResultInspectPanicFailsWithoutRedownloading(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{inspectPanic: "runtime invariant violated"}
w := NewWithAPI(config.Config{}, api)
w.engine = eng
total := int64(100)
w.process(context.Background(), client.DownloadTask{
ID: "task-1",
Status: "assigned",
SourceType: "magnet",
SourceURI: "magnet:?xt=urn:btih:abc123",
TotalBytes: &total,
Detail: &client.DownloadTaskDetail{Phase: "uploading"},
UploadToken: "upload-token",
})
if eng.downloadCalls != 0 {
t.Fatalf("expected runtime inspection panic not to restart download, got %d download calls", eng.downloadCalls)
}
failed := lastPatchWithStatus(t, api.patches, "failed")
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "panic: runtime invariant violated") {
t.Fatalf("expected panic to be reported, got %#v", failed.ErrorMessage)
}
}
func TestUploadExistingResultMissingRuntimeFailsWithoutRedownloading(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{taskFound: false}
w := NewWithAPI(config.Config{}, api)
w.engine = eng
total := int64(100)
w.process(context.Background(), client.DownloadTask{
ID: "task-1",
Status: "assigned",
SourceType: "magnet",
SourceURI: "magnet:?xt=urn:btih:abc123",
TotalBytes: &total,
Detail: &client.DownloadTaskDetail{Phase: "uploading"},
UploadToken: "upload-token",
})
if eng.downloadCalls != 0 {
t.Fatalf("expected missing runtime task not to restart download, got %d download calls", eng.downloadCalls)
}
failed := lastPatchWithStatus(t, api.patches, "failed")
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "missing from downloader runtime") {
t.Fatalf("expected missing runtime to be reported, got %#v", failed.ErrorMessage)
}
}
func TestUploadExistingResultIncompleteRuntimeTaskFailsWithoutRedownloading(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{
taskSnapshot: engine.TaskSnapshot{State: engine.TaskStateDownloading, Downloaded: 10},
taskFound: true,
}
w := NewWithAPI(config.Config{}, api)
w.engine = eng
total := int64(100)
w.process(context.Background(), client.DownloadTask{
ID: "task-1",
Status: "assigned",
SourceType: "magnet",
SourceURI: "magnet:?xt=urn:btih:abc123",
TotalBytes: &total,
Detail: &client.DownloadTaskDetail{Phase: "uploading"},
UploadToken: "upload-token",
})
if eng.downloadCalls != 0 {
t.Fatalf("expected incomplete runtime task not to restart download, got %d download calls", eng.downloadCalls)
}
failed := lastPatchWithStatus(t, api.patches, "failed")
if failed.ErrorMessage == nil || !strings.Contains(*failed.ErrorMessage, "server task requires upload but runtime task is not completed") {
t.Fatalf("expected invariant failure to be reported, got %#v", failed.ErrorMessage)
}
}
func TestDownloadShutdownMarksTaskInterrupted(t *testing.T) {
api := &recordingAPI{}
eng := &recordingEngine{downloadErr: context.Canceled}
@@ -417,63 +547,63 @@ func TestTaskErrorMessageTruncatesToSchemaLimit(t *testing.T) {
}
}
func TestResumeStage(t *testing.T) {
func TestNextTaskWorkStage(t *testing.T) {
total := int64(100)
cases := []struct {
name string
task client.DownloadTask
want taskResumeStage
want taskWorkStage
}{
{
name: "uploading status",
task: client.DownloadTask{Status: "uploading"},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "assigned with upload bytes",
task: client.DownloadTask{Status: "assigned", StorageUploadedBytes: 1},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "assigned with uploading phase",
task: client.DownloadTask{Status: "assigned", Detail: &client.DownloadTaskDetail{Phase: "uploading"}},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "assigned with completed phase",
task: client.DownloadTask{Status: "assigned", Detail: &client.DownloadTaskDetail{Phase: "completed"}},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "assigned with completed download bytes",
task: client.DownloadTask{Status: "assigned", DownloadedBytes: 100, TotalBytes: &total},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "assigned partial download",
task: client.DownloadTask{Status: "assigned", DownloadedBytes: 99, TotalBytes: &total},
want: taskResumeDownload,
want: taskWorkStageDownload,
},
{
name: "downloading completed bytes",
task: client.DownloadTask{Status: "downloading", DownloadedBytes: 100, TotalBytes: &total},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
{
name: "downloading partial download",
task: client.DownloadTask{Status: "downloading", DownloadedBytes: 99, TotalBytes: &total},
want: taskResumeDownload,
want: taskWorkStageDownload,
},
{
name: "interrupted with uploading phase",
task: client.DownloadTask{Status: "interrupted", Detail: &client.DownloadTaskDetail{Phase: "uploading"}},
want: taskResumeUpload,
want: taskWorkStageUploadExistingResult,
},
}
for _, tc := range cases {
t.Run(tc.name, func(t *testing.T) {
if got := resumeStage(tc.task); got != tc.want {
if got := nextTaskWorkStage(tc.task); got != tc.want {
t.Fatalf("expected %v, got %v", tc.want, got)
}
})
@@ -731,12 +861,13 @@ func findPatchWithStatus(patches []client.TaskPatch, status string) (client.Task
type recordingEngine struct {
downloadResult engine.Result
downloadErr error
recoverResult engine.Result
recoverErr error
recovered bool
taskSnapshot engine.TaskSnapshot
inspectErr error
inspectPanic any
taskFound bool
restoreSeed *engine.Seed
downloadCalls int
recoverCalls int
inspectCalls int
restoreCalls int
}
@@ -752,12 +883,15 @@ func (e *recordingEngine) Check(context.Context) error {
return nil
}
func (e *recordingEngine) Recover(context.Context, client.DownloadTask) (engine.Result, bool, error) {
e.recoverCalls++
if e.recoverErr != nil {
return engine.Result{}, false, e.recoverErr
func (e *recordingEngine) InspectTask(context.Context, client.DownloadTask) (engine.TaskSnapshot, bool, error) {
e.inspectCalls++
if e.inspectPanic != nil {
panic(e.inspectPanic)
}
return e.recoverResult, e.recovered, nil
if e.inspectErr != nil {
return engine.TaskSnapshot{}, false, e.inspectErr
}
return e.taskSnapshot, e.taskFound, nil
}
func (e *recordingEngine) RestoreSeed(context.Context, engine.SeedRef) (*engine.Seed, error) {