fix(downloader): reset runtime on task restart

This commit is contained in:
saltbo
2026-06-06 00:54:26 -04:00
parent d8daba2542
commit 65eb5d9738
21 changed files with 4289 additions and 5 deletions
+25
View File
@@ -387,6 +387,10 @@
"canceled"
]
},
"attempt": {
"type": "integer",
"minimum": 1
},
"assignment": {
"type": "object",
"nullable": true,
@@ -796,6 +800,7 @@
},
"required": [
"state",
"attempt",
"assignment",
"progress",
"billing",
@@ -1023,6 +1028,10 @@
"canceled"
]
},
"attempt": {
"type": "integer",
"minimum": 1
},
"assignment": {
"type": "object",
"nullable": true,
@@ -1432,6 +1441,7 @@
},
"required": [
"state",
"attempt",
"assignment",
"progress",
"billing",
@@ -1769,6 +1779,10 @@
"canceled"
]
},
"attempt": {
"type": "integer",
"minimum": 1
},
"assignment": {
"type": "object",
"nullable": true,
@@ -2178,6 +2192,7 @@
},
"required": [
"state",
"attempt",
"assignment",
"progress",
"billing",
@@ -2685,6 +2700,10 @@
"canceled"
]
},
"attempt": {
"type": "integer",
"minimum": 1
},
"assignment": {
"type": "object",
"nullable": true,
@@ -3094,6 +3113,7 @@
},
"required": [
"state",
"attempt",
"assignment",
"progress",
"billing",
@@ -3336,6 +3356,10 @@
"canceled"
]
},
"attempt": {
"type": "integer",
"minimum": 1
},
"assignment": {
"type": "object",
"nullable": true,
@@ -3745,6 +3769,7 @@
},
"required": [
"state",
"attempt",
"assignment",
"progress",
"billing",
+5
View File
@@ -59,6 +59,10 @@ func (t DownloadTask) State() string {
return t.Status.State
}
func (t DownloadTask) Attempt() int {
return t.Status.Attempt
}
func (t DownloadTask) Runtime() *DownloadTaskRuntime {
return t.Status.Runtime
}
@@ -93,6 +97,7 @@ type DownloadTaskLabels struct {
type DownloadTaskStatus struct {
State string `json:"state"`
Attempt int `json:"attempt"`
Assignment *DownloadTaskAssignment `json:"assignment"`
Progress DownloadTaskProgress `json:"progress"`
Runtime *DownloadTaskRuntime `json:"runtime"`
+48
View File
@@ -124,6 +124,35 @@ func (a Aria2) Check(ctx context.Context) error {
return nil
}
func (a Aria2) ResetTask(ctx context.Context, task client.DownloadTask) error {
aria, err := a.client(ctx)
if err != nil {
return err
}
defer aria.Close()
taskDir := filepath.Clean(filepath.Join(a.Dir, task.ID))
statuses, err := a.taskStatuses(ctx, &aria)
if err != nil {
return err
}
var resetErrs []error
for _, status := range statuses {
if !aria2StatusBelongsToTask(status, taskDir, aria2TaskGID(task.ID)) {
continue
}
if err := aria.ForceRemove(status.GID); err != nil {
resetErrs = append(resetErrs, fmt.Errorf("force remove aria2 gid %s: %w", status.GID, err))
}
if err := aria.RemoveDownloadResult(status.GID); err != nil {
resetErrs = append(resetErrs, fmt.Errorf("remove aria2 result %s: %w", status.GID, err))
}
}
if err := os.RemoveAll(taskDir); err != nil {
resetErrs = append(resetErrs, fmt.Errorf("remove task dir %s: %w", taskDir, err))
}
return errors.Join(resetErrs...)
}
func (a Aria2) SaveSession(ctx context.Context) error {
client, err := a.client(ctx)
if err != nil {
@@ -586,6 +615,25 @@ func aria2StatusMatchesTask(status arigo.Status, taskDir string, gid string, inf
return false
}
func aria2StatusBelongsToTask(status arigo.Status, taskDir string, gid string) bool {
if status.GID == gid || status.Following == gid || status.BelongsTo == gid {
return true
}
if filepath.Clean(status.Dir) == taskDir {
return true
}
for _, file := range status.Files {
if file.Path == "" {
continue
}
abs, _ := downloadedPath(taskDir, file.Path)
if strings.HasPrefix(filepath.Clean(abs), taskDir+string(filepath.Separator)) {
return true
}
}
return false
}
func aria2StatusTaskDir(status arigo.Status, fallback string) string {
if status.Dir == "" {
return fallback
+4
View File
@@ -71,6 +71,10 @@ type Engine interface {
Download(ctx context.Context, task client.DownloadTask, progress Progress) (Result, error)
}
type TaskResetter interface {
ResetTask(ctx context.Context, task client.DownloadTask) error
}
type SeedRestorer interface {
RestoreSeed(ctx context.Context, ref SeedRef) (*Seed, error)
}
+7
View File
@@ -42,6 +42,13 @@ func (h HTTP) Check(ctx context.Context) error {
return os.Remove(path)
}
func (h HTTP) ResetTask(ctx context.Context, task client.DownloadTask) error {
if task.SourceType() != "http" {
return nil
}
return os.RemoveAll(filepath.Join(h.Dir, task.ID))
}
func (h HTTP) InspectTask(ctx context.Context, task client.DownloadTask) (TaskSnapshot, bool, error) {
if task.SourceType() != "http" {
return TaskSnapshot{}, false, nil
+50
View File
@@ -77,6 +77,47 @@ func (q QBittorrent) Check(ctx context.Context) error {
return nil
}
func (q QBittorrent) ResetTask(ctx context.Context, task client.DownloadTask) error {
if task.SourceType() == "http" {
return HTTP{Dir: q.Dir}.ResetTask(ctx, task)
}
qbt, err := q.login(ctx)
if err != nil {
return err
}
taskDir := filepath.Clean(filepath.Join(q.Dir, task.ID))
tag := qbittorrentTrackingTag(task.ID)
tagged, err := qbt.GetTorrentsCtx(ctx, qbittorrent.TorrentFilterOptions{Tag: tag})
if err != nil {
return err
}
all, err := qbt.GetTorrentsCtx(ctx, qbittorrent.TorrentFilterOptions{})
if err != nil {
return err
}
seen := map[string]struct{}{}
hashes := make([]string, 0, len(tagged))
for _, torrent := range append(tagged, all...) {
if torrent.Hash == "" {
continue
}
if _, ok := seen[torrent.Hash]; ok {
continue
}
if filepath.Clean(torrent.SavePath) != taskDir && !torrentHasTag(torrent.Tags, tag) {
continue
}
seen[torrent.Hash] = struct{}{}
hashes = append(hashes, torrent.Hash)
}
if len(hashes) > 0 {
if err := qbt.DeleteTorrentsCtx(ctx, hashes, false); err != nil {
return err
}
}
return os.RemoveAll(taskDir)
}
func (q QBittorrent) RestoreSeed(ctx context.Context, ref SeedRef) (*Seed, error) {
qbt, err := q.login(ctx)
if err != nil {
@@ -311,6 +352,15 @@ func qbittorrentTrackingTag(taskID string) string {
return "ztid=" + taskID
}
func torrentHasTag(tags string, want string) bool {
for _, tag := range strings.Split(tags, ",") {
if strings.TrimSpace(tag) == want {
return true
}
}
return false
}
func (q QBittorrent) cleanupSeed(hash string, localPath string) func(context.Context) error {
return func(ctx context.Context) error {
var errs []error
@@ -1782,6 +1782,7 @@ type PostApiDownloadTasksIdActions200JSONResponseBody0 struct {
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -4045,6 +4046,7 @@ type GetApiDownloadTasksResponse struct {
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -4197,6 +4199,7 @@ type PostApiDownloadTasksResponse struct {
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -4383,6 +4386,7 @@ type GetApiDownloadTasksIdResponse struct {
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -4531,6 +4535,7 @@ type PatchApiDownloadTasksIdResponse struct {
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -5526,6 +5531,7 @@ func ParseGetApiDownloadTasksResponse(rsp *http.Response) (*GetApiDownloadTasksR
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -5680,6 +5686,7 @@ func ParsePostApiDownloadTasksResponse(rsp *http.Response) (*PostApiDownloadTask
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -5876,6 +5883,7 @@ func ParseGetApiDownloadTasksIdResponse(rsp *http.Response) (*GetApiDownloadTask
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -6026,6 +6034,7 @@ func ParsePatchApiDownloadTasksIdResponse(rsp *http.Response) (*PatchApiDownload
DownloaderId string `json:"downloaderId"`
UploadToken *string `json:"uploadToken,omitempty"`
} `json:"assignment"`
Attempt int `json:"attempt"`
Billing struct {
AuthorizedBytes int64 `json:"authorizedBytes"`
ChargedBytes int64 `json:"chargedBytes"`
@@ -0,0 +1,56 @@
package worker
import (
"encoding/json"
"os"
"path/filepath"
)
type attemptLedger struct {
Attempts map[string]int `json:"attempts"`
}
func attemptLedgerPath(stateDir string) string {
return filepath.Join(stateDir, "attempts.json")
}
func loadAttemptLedger(stateDir string) (attemptLedger, error) {
path := attemptLedgerPath(stateDir)
data, err := os.ReadFile(path)
if err != nil {
if os.IsNotExist(err) {
return attemptLedger{Attempts: map[string]int{}}, nil
}
return attemptLedger{}, err
}
if len(data) == 0 {
return attemptLedger{Attempts: map[string]int{}}, nil
}
var ledger attemptLedger
if err := json.Unmarshal(data, &ledger); err != nil {
return attemptLedger{}, err
}
if ledger.Attempts == nil {
ledger.Attempts = map[string]int{}
}
return ledger, nil
}
func saveAttemptLedger(stateDir string, ledger attemptLedger) error {
if err := os.MkdirAll(stateDir, 0o755); err != nil {
return err
}
if ledger.Attempts == nil {
ledger.Attempts = map[string]int{}
}
data, err := json.MarshalIndent(ledger, "", " ")
if err != nil {
return err
}
path := attemptLedgerPath(stateDir)
tmp := path + ".tmp"
if err := os.WriteFile(tmp, append(data, '\n'), 0o644); err != nil {
return err
}
return os.Rename(tmp, path)
}
+12
View File
@@ -306,6 +306,18 @@ func (w *Worker) cleanupRetainedSeed(ctx context.Context, seed retainedSeed, rea
}
}
func (w *Worker) cleanupRetainedSeedForTask(ctx context.Context, taskID string, reason string) {
for _, seed := range w.retainedSeedSnapshot() {
if seed.taskID == taskID {
w.cleanupRetainedSeed(ctx, seed, reason)
break
}
}
if err := w.removeSeedLedger(taskID); err != nil {
w.logger.Warn("failed to remove retained seed ledger entry", "task_id", taskID, "error", err)
}
}
func (w *Worker) upsertSeedLedger(seed retainedSeed, size int64) error {
if w.cfg.StateDir == "" {
return nil
+80 -4
View File
@@ -41,6 +41,7 @@ type Worker struct {
logger *slog.Logger
running map[string]context.CancelCauseFunc
retainedSeeds []retainedSeed
attempts map[string]int
started []*exec.Cmd
wg sync.WaitGroup
mu sync.Mutex
@@ -70,10 +71,11 @@ func New(cfg config.Config) (*Worker, error) {
func NewWithAPI(cfg config.Config, api apiClient) *Worker {
return &Worker{
cfg: cfg,
api: api,
logger: slog.Default(),
running: map[string]context.CancelCauseFunc{},
cfg: cfg,
api: api,
logger: slog.Default(),
running: map[string]context.CancelCauseFunc{},
attempts: map[string]int{},
}
}
@@ -189,6 +191,14 @@ func (w *Worker) process(ctx context.Context, task client.DownloadTask) {
log := w.taskLogger(task)
defer w.recoverTaskPanic(ctx, task.ID, log)
log.Info("task started", "source_uri", task.SourceURI(), "target_folder", task.TargetFolder())
if err := w.resetTaskForAttempt(ctx, task, log); err != nil {
msg := taskErrorMessage(err)
log.Error("failed to reset task for restart", "attempt", task.Attempt(), "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
}
currentDetail := task.Runtime()
if nextTaskWorkStage(task) == taskWorkStageUploadExistingResult {
w.uploadExistingResult(ctx, log, task, currentDetail)
@@ -359,6 +369,72 @@ func nextTaskWorkStage(task client.DownloadTask) taskWorkStage {
return taskWorkStageDownload
}
func (w *Worker) resetTaskForAttempt(ctx context.Context, task client.DownloadTask, log *slog.Logger) error {
attempt := task.Attempt()
if attempt <= 0 {
return fmt.Errorf("download task has invalid attempt %d", attempt)
}
if w.cfg.StateDir == "" {
seen := w.memoryAttempt(task.ID)
if seen == attempt {
return nil
}
if attempt > 1 {
if err := w.resetRuntimeTask(ctx, task, log); err != nil {
return err
}
}
w.setMemoryAttempt(task.ID, attempt)
return nil
}
ledger, err := loadAttemptLedger(w.cfg.StateDir)
if err != nil {
return fmt.Errorf("load attempt ledger: %w", err)
}
seen := ledger.Attempts[task.ID]
if seen == attempt {
return nil
}
if attempt > 1 {
if err := w.resetRuntimeTask(ctx, task, log); err != nil {
return err
}
}
ledger.Attempts[task.ID] = attempt
if err := saveAttemptLedger(w.cfg.StateDir, ledger); err != nil {
return fmt.Errorf("save attempt ledger: %w", err)
}
return nil
}
func (w *Worker) memoryAttempt(taskID string) int {
w.mu.Lock()
defer w.mu.Unlock()
return w.attempts[taskID]
}
func (w *Worker) setMemoryAttempt(taskID string, attempt int) {
w.mu.Lock()
defer w.mu.Unlock()
if w.attempts == nil {
w.attempts = map[string]int{}
}
w.attempts[taskID] = attempt
}
func (w *Worker) resetRuntimeTask(ctx context.Context, task client.DownloadTask, log *slog.Logger) error {
resetter, ok := w.engine.(engine.TaskResetter)
if !ok {
return fmt.Errorf("engine %s does not support task reset", w.engine.Name())
}
w.cleanupRetainedSeedForTask(ctx, task.ID, "restart")
log.Info("resetting downloader runtime task", "attempt", task.Attempt())
if err := resetter.ResetTask(ctx, task); err != nil {
return fmt.Errorf("reset downloader runtime task: %w", err)
}
return nil
}
func (w *Worker) uploadAndComplete(
ctx context.Context,
log *slog.Logger,
+70 -1
View File
@@ -566,6 +566,68 @@ func TestNextTaskWorkStage(t *testing.T) {
}
}
func TestResetTaskForRestartAttemptResetsRuntimeAndRecordsAttempt(t *testing.T) {
stateDir := t.TempDir()
seedPath := t.TempDir()
if err := saveSeedLedger(stateDir, seedLedger{Seeds: []seedLedgerEntry{{
TaskID: "task-1",
Engine: "aria2",
SeedID: "gid",
InfoHash: "abc123",
Path: seedPath,
RetainedAt: time.Now().Add(-time.Minute),
ExpiresAt: time.Now().Add(time.Hour),
}}}); err != nil {
t.Fatal(err)
}
task := clientTaskWithStatus("task-1", "assigned")
task.Status.Attempt = 2
eng := &recordingEngine{}
w := NewWithAPI(config.Config{StateDir: stateDir}, nil)
w.engine = eng
if err := w.resetTaskForAttempt(context.Background(), task, w.logger); err != nil {
t.Fatal(err)
}
if eng.resetCalls != 1 {
t.Fatalf("expected reset once, got %d", eng.resetCalls)
}
attempts, err := loadAttemptLedger(stateDir)
if err != nil {
t.Fatal(err)
}
if attempts.Attempts["task-1"] != 2 {
t.Fatalf("expected attempt 2 to be recorded, got %#v", attempts.Attempts)
}
seedLedger, err := loadSeedLedger(stateDir)
if err != nil {
t.Fatal(err)
}
if len(seedLedger.Seeds) != 0 {
t.Fatalf("expected restart to remove retained seed ledger, got %#v", seedLedger.Seeds)
}
}
func TestResetTaskForRestartAttemptSkipsAlreadyRecordedAttempt(t *testing.T) {
stateDir := t.TempDir()
if err := saveAttemptLedger(stateDir, attemptLedger{Attempts: map[string]int{"task-1": 2}}); err != nil {
t.Fatal(err)
}
task := clientTaskWithStatus("task-1", "assigned")
task.Status.Attempt = 2
eng := &recordingEngine{}
w := NewWithAPI(config.Config{StateDir: stateDir}, nil)
w.engine = eng
if err := w.resetTaskForAttempt(context.Background(), task, w.logger); err != nil {
t.Fatal(err)
}
if eng.resetCalls != 0 {
t.Fatalf("expected no reset for recorded attempt, got %d", eng.resetCalls)
}
}
func TestRetainSeedKeepsDownloadedResult(t *testing.T) {
dir := t.TempDir()
stateDir := t.TempDir()
@@ -844,7 +906,7 @@ func clientTask(id string) client.DownloadTask {
Source: client.DownloadTaskSource{Type: "magnet"},
Labels: client.DownloadTaskLabels{Tags: []string{}},
},
Status: client.DownloadTaskStatus{},
Status: client.DownloadTaskStatus{Attempt: 1},
}
}
@@ -898,12 +960,14 @@ func findPatchWithStatus(patches []client.TaskPatch, status string) (client.Task
type recordingEngine struct {
downloadResult engine.Result
downloadErr error
resetErr error
taskSnapshot engine.TaskSnapshot
inspectErr error
inspectPanic any
taskFound bool
restoreSeed *engine.Seed
downloadCalls int
resetCalls int
inspectCalls int
restoreCalls int
}
@@ -936,6 +1000,11 @@ func (e *recordingEngine) RestoreSeed(context.Context, engine.SeedRef) (*engine.
return e.restoreSeed, nil
}
func (e *recordingEngine) ResetTask(context.Context, client.DownloadTask) error {
e.resetCalls++
return e.resetErr
}
func (e *recordingEngine) Download(context.Context, client.DownloadTask, engine.Progress) (engine.Result, error) {
e.downloadCalls++
return e.downloadResult, e.downloadErr
@@ -0,0 +1 @@
ALTER TABLE `download_tasks` ADD `attempt` integer DEFAULT 1 NOT NULL;
File diff suppressed because it is too large Load Diff
+7
View File
@@ -302,6 +302,13 @@
"when": 1780719304552,
"tag": "0043_normalize-download-task-runtime",
"breakpoints": true
},
{
"idx": 44,
"version": "6",
"when": 1780721031283,
"tag": "0044_add-download-task-attempt",
"breakpoints": true
}
]
}
+1
View File
@@ -340,6 +340,7 @@ export const downloadTasks = sqliteTable(
tags: text('tags').notNull().default('[]'),
assignedDownloaderId: text('assigned_downloader_id'),
status: text('status').notNull(),
attempt: integer('attempt').notNull().default(1),
billingAuthorizedBytes: integer('billing_authorized_bytes').notNull().default(0),
billingChargedBytes: integer('billing_charged_bytes').notNull().default(0),
billingChargedCredits: integer('billing_charged_credits').notNull().default(0),
@@ -843,6 +843,7 @@ describe('Download tasks API integration', () => {
await expect(restartRes.json()).resolves.toMatchObject({
status: {
state: 'assigned',
attempt: 2,
assignment: { downloaderId: createdDownloader.downloader.id },
progress: {
download: { bytes: 0, totalBytes: null },
+1
View File
@@ -533,6 +533,7 @@ export async function performDownloadTaskAction(
.set({
status: 'queued',
assignedDownloaderId: null,
attempt: task.attempt + 1,
billingAuthorizedBytes: 0,
billingChargedBytes: 0,
billingChargedCredits: 0,
+1
View File
@@ -51,6 +51,7 @@ export function toDownloadTask(row: DownloadTaskRow): DownloadTask {
},
status: {
state: row.status as DownloadTask['status']['state'],
attempt: row.attempt,
assignment: row.assignedDownloaderId
? { downloaderId: row.assignedDownloaderId, assignedAt: row.assignedAt?.toISOString() ?? null }
: null,
+1
View File
@@ -429,6 +429,7 @@ const APP_SCHEMA_SQL = `
tags TEXT NOT NULL DEFAULT '[]',
assigned_downloader_id TEXT,
status TEXT NOT NULL,
attempt INTEGER NOT NULL DEFAULT 1,
billing_authorized_bytes INTEGER NOT NULL DEFAULT 0,
billing_charged_bytes INTEGER NOT NULL DEFAULT 0,
billing_charged_credits INTEGER NOT NULL DEFAULT 0,
+1
View File
@@ -120,6 +120,7 @@ export const downloadTaskSchema = z.object({
}),
status: z.object({
state: downloadTaskStatusSchema,
attempt: z.number().int().min(1),
assignment: z
.object({
downloaderId: z.string(),
+1
View File
@@ -275,6 +275,7 @@ export interface DownloadTaskSpec {
export interface DownloadTaskExecutionStatus {
state: DownloadTaskStatus
attempt: number
assignment: DownloadTaskAssignment | null
progress: DownloadTaskProgress
billing: DownloadTaskBilling