From bdb8ed177fc12d8a53c6fbb7f3bbe50bb98a1a0a Mon Sep 17 00:00:00 2001 From: Levi Date: Sat, 21 Sep 2024 10:19:52 +0800 Subject: [PATCH] feat: supports batch operation of tasks by filter (#753) --- pkg/download/downloader.go | 224 +++++++++++++++++++++----------- pkg/download/downloader_test.go | 154 +++++++++++++++++++++- pkg/download/extension.go | 8 +- pkg/download/model.go | 10 ++ pkg/rest/api.go | 104 ++++++++++----- pkg/rest/server.go | 4 +- pkg/rest/server_test.go | 4 +- 7 files changed, 394 insertions(+), 114 deletions(-) diff --git a/pkg/download/downloader.go b/pkg/download/downloader.go index 684ad66e..337bc099 100644 --- a/pkg/download/downloader.go +++ b/pkg/download/downloader.go @@ -12,6 +12,7 @@ import ( "github.com/rs/zerolog" "github.com/rs/zerolog/pkgerrors" "github.com/virtuald/go-paniclog" + "math" gohttp "net/http" "net/url" "os" @@ -371,48 +372,28 @@ func (d *Downloader) Create(rrId string, opts *base.Options) (taskId string, err return d.doCreate(fetcher, opts) } -func (d *Downloader) Pause(id string) (err error) { - task := d.GetTask(id) - if task == nil { +func (d *Downloader) Pause(filter *TaskFilter) (err error) { + if filter == nil || filter.IsEmpty() { + return d.pauseAll() + } + + filter.NotStatuses = []base.Status{base.DownloadStatusPause, base.DownloadStatusError, base.DownloadStatusDone} + pauseTasks := d.GetTasksByFilter(filter) + if len(pauseTasks) == 0 { return ErrTaskNotFound } - if err = d.doPause(task); err != nil { - return + for _, task := range pauseTasks { + if err = d.doPause(task); err != nil { + return + } } d.notifyRunning() + return } -func (d *Downloader) Continue(id string) (err error) { - task := d.GetTask(id) - if task == nil { - return ErrTaskNotFound - } - - func() { - d.lock.Lock() - defer d.lock.Unlock() - - remainRunningCount := d.remainRunningCount() - if remainRunningCount == 0 { - for _, t := range d.tasks { - if t.Status == base.DownloadStatusRunning { - if err = d.doPause(t); err != nil { - return - } - t.Status = base.DownloadStatusWait - d.waitTasks = append(d.waitTasks, t) - break - } - } - } - }() - - return d.doStart(task) -} - -func (d *Downloader) PauseAll() (err error) { +func (d *Downloader) pauseAll() (err error) { func() { d.lock.Lock() defer d.lock.Unlock() @@ -430,7 +411,65 @@ func (d *Downloader) PauseAll() (err error) { return } -func (d *Downloader) ContinueAll() (err error) { +// Continue specific tasks, if continue tasks will exceed maxRunning, it needs pause some running tasks before that +func (d *Downloader) Continue(filter *TaskFilter) (err error) { + if filter == nil || filter.IsEmpty() { + return d.continueAll() + } + + filter.NotStatuses = []base.Status{base.DownloadStatusRunning, base.DownloadStatusDone} + continueTasks := d.GetTasksByFilter(filter) + if len(continueTasks) == 0 { + return ErrTaskNotFound + } + + realContinueTasks := make([]*Task, 0) + func() { + d.lock.Lock() + defer d.lock.Unlock() + + continueCount := len(continueTasks) + remainRunningCount := d.remainRunningCount() + needRunningCount := int(math.Min(float64(d.cfg.MaxRunning), float64(continueCount))) + needPauseCount := needRunningCount - remainRunningCount + if needPauseCount > 0 { + pausedCount := 0 + for _, task := range d.tasks { + if task.Status == base.DownloadStatusRunning { + if err = d.doPause(task); err != nil { + return + } + task.Status = base.DownloadStatusWait + d.waitTasks = append(d.waitTasks, task) + pausedCount++ + } + if pausedCount == needPauseCount { + break + } + } + } + + for _, task := range continueTasks { + if len(realContinueTasks) < needRunningCount { + realContinueTasks = append(realContinueTasks, task) + } else { + task.Status = base.DownloadStatusWait + d.waitTasks = append(d.waitTasks, task) + } + } + }() + + for _, task := range realContinueTasks { + if err = d.doStart(task); err != nil { + return + } + } + + return +} + +// continueAll continue all tasks but does not affect tasks already running +func (d *Downloader) continueAll() (err error) { continuedTasks := make([]*Task, 0) func() { @@ -451,8 +490,7 @@ func (d *Downloader) ContinueAll() (err error) { }() for _, task := range continuedTasks { - tt := task - if err = d.doStart(tt); err != nil { + if err = d.doStart(task); err != nil { return } } @@ -460,40 +498,26 @@ func (d *Downloader) ContinueAll() (err error) { return } -func (d *Downloader) Delete(id string, force bool) (err error) { - task := d.GetTask(id) - if task == nil { - return ErrTaskNotFound +func (d *Downloader) ContinueBatch(filter *TaskFilter) (err error) { + if filter == nil || filter.IsEmpty() { + return d.continueAll() } - func() { - d.lock.Lock() - defer d.lock.Unlock() - - for i, t := range d.tasks { - if t.ID == id { - d.tasks = append(d.tasks[:i], d.tasks[i+1:]...) - break - } + continueTasks := d.GetTasksByFilter(filter) + for _, task := range continueTasks { + if err = d.doStart(task); err != nil { + return } - for i, t := range d.waitTasks { - if t.ID == id { - d.waitTasks = append(d.waitTasks[:i], d.waitTasks[i+1:]...) - break - } - } - }() - - err = d.doDelete(task, force) - if err != nil { - return } - d.notifyRunning() return } -func (d *Downloader) DeleteByStatues(statues []base.Status, force bool) (err error) { - deleteTasks := d.GetTasksByStatues(statues) +func (d *Downloader) Delete(filter *TaskFilter, force bool) (err error) { + if filter == nil || filter.IsEmpty() { + return d.deleteAll() + } + + deleteTasks := d.GetTasksByFilter(filter) if len(deleteTasks) == 0 { return } @@ -535,6 +559,23 @@ func (d *Downloader) DeleteByStatues(statues []base.Status, force bool) (err err return } +func (d *Downloader) deleteAll() (err error) { + func() { + d.lock.Lock() + defer d.lock.Unlock() + + d.tasks = make([]*Task, 0) + d.waitTasks = make([]*Task, 0) + }() + + for _, task := range d.tasks { + if err = d.doDelete(task, true); err != nil { + return + } + } + return +} + func (d *Downloader) Stats(id string) (sr any, err error) { task := d.GetTask(id) if task == nil { @@ -557,9 +598,6 @@ func (d *Downloader) Stats(id string) (sr any, err error) { func (d *Downloader) doDelete(task *Task, force bool) (err error) { err = func() error { - d.lock.Lock() - defer d.lock.Unlock() - if err := d.storage.Delete(bucketTask, task.ID); err != nil { return err } @@ -587,10 +625,10 @@ func (d *Downloader) doDelete(task *Task, force bool) (err error) { task = nil return nil }() + if err != nil { d.Logger.Error().Stack().Err(err).Msgf("delete task failed, task id: %s", task.ID) } - return } @@ -598,7 +636,7 @@ func (d *Downloader) Close() error { d.closed.Store(true) closeArr := []func() error{ - d.PauseAll, + d.pauseAll, } for _, fm := range d.cfg.FetchManagers { closeArr = append(closeArr, fm.Close) @@ -660,17 +698,51 @@ func (d *Downloader) GetTasks() []*Task { return d.tasks } -func (d *Downloader) GetTasksByStatues(statues []base.Status) []*Task { - if len(statues) == 0 { +// GetTasksByFilter get tasks by filter, if filter is nil, return all tasks +// return tasks and if match all tasks +func (d *Downloader) GetTasksByFilter(filter *TaskFilter) []*Task { + if filter == nil || filter.IsEmpty() { return d.tasks } + + idMatch := func(task *Task) bool { + if len(filter.IDs) == 0 { + return true + } + for _, id := range filter.IDs { + if task.ID == id { + return true + } + } + return false + } + statusMatch := func(task *Task) bool { + if len(filter.Statuses) == 0 { + return true + } + for _, status := range filter.Statuses { + if task.Status == status { + return true + } + } + return false + } + notStatusMatch := func(task *Task) bool { + if len(filter.NotStatuses) == 0 { + return true + } + for _, status := range filter.NotStatuses { + if task.Status == status { + return false + } + } + return true + } + tasks := make([]*Task, 0) for _, task := range d.tasks { - for _, status := range statues { - if task.Status == status { - tasks = append(tasks, task) - break - } + if idMatch(task) && statusMatch(task) && notStatusMatch(task) { + tasks = append(tasks, task) } } return tasks diff --git a/pkg/download/downloader_test.go b/pkg/download/downloader_test.go index 1234aac8..a7cfd10c 100644 --- a/pkg/download/downloader_test.go +++ b/pkg/download/downloader_test.go @@ -133,7 +133,7 @@ func TestDownloader_CreateDirectBatch(t *testing.T) { t.Fatal(err) } defer func() { - downloader.DeleteByStatues(nil, true) + downloader.Delete(nil, true) downloader.Clear() }() @@ -400,7 +400,7 @@ func TestDownloader_StoreAndRestore(t *testing.T) { t.Fatal(err) } time.Sleep(time.Millisecond * 1001) - err = downloader.Pause(id) + err = downloader.Pause(&TaskFilter{IDs: []string{id}}) if err != nil { t.Fatal(err) } @@ -424,7 +424,7 @@ func TestDownloader_StoreAndRestore(t *testing.T) { wg.Done() } }) - err = downloader.Continue(id) + err = downloader.Continue(&TaskFilter{IDs: []string{id}}) wg.Wait() if err != nil { t.Fatal(err) @@ -485,3 +485,151 @@ func TestDownloader_Protocol_Config(t *testing.T) { t.Errorf("GetConfig() got = %v, want %v", test.ToJson(storeCfg), test.ToJson(newStoreCfg)) } } + +func TestDownloader_GetTasksByFilter(t *testing.T) { + listener := test.StartTestFileServer() + defer listener.Close() + + downloader := NewDownloader(nil) + if err := downloader.Setup(); err != nil { + t.Fatal(err) + } + defer func() { + downloader.Delete(nil, true) + downloader.Clear() + }() + + reqs := make([]*base.Request, 0) + fileNames := make([]string, 0) + for i := 0; i < 10; i++ { + req := &base.Request{ + URL: "http://" + listener.Addr().String() + "/" + test.BuildName, + } + reqs = append(reqs, req) + if i == 0 { + fileNames = append(fileNames, test.DownloadName) + } else { + arr := strings.Split(test.DownloadName, ".") + fileNames = append(fileNames, arr[0]+" ("+strconv.Itoa(i)+")."+arr[1]) + } + } + + var wg sync.WaitGroup + wg.Add(len(reqs)) + downloader.Listener(func(event *Event) { + if event.Key == EventKeyDone { + wg.Done() + } + }) + + taskIds, err := downloader.CreateDirectBatch(reqs, &base.Options{ + Path: test.Dir, + Name: test.DownloadName, + Extra: http.OptsExtra{ + Connections: 4, + }, + }) + if err != nil { + t.Fatal(err) + } + + wg.Wait() + + t.Run("GetTasksByFilter nil", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(nil) + if len(tasks) != len(reqs) { + t.Errorf("GetTasksByFilter nil task got = %v, want %v", len(tasks), len(reqs)) + } + }) + + t.Run("GetTasksByFilter empty", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{}) + if len(tasks) != len(reqs) { + t.Errorf("GetTasksByFilter empty task got = %v, want %v", len(tasks), len(reqs)) + } + }) + + t.Run("GetTasksByFilter ids", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + IDs: taskIds, + }) + if len(tasks) != len(reqs) { + t.Errorf("GetTasksByFilter ids task got = %v, want %v", len(tasks), len(reqs)) + } + }) + + t.Run("GetTasksByFilter match ids", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + IDs: []string{taskIds[0]}, + }) + if len(tasks) != 1 { + t.Errorf("GetTasksByFilter ids task got = %v, want %v", len(tasks), 1) + } + }) + + t.Run("GetTasksByFilter not match ids", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + IDs: []string{"xxx"}, + }) + if len(tasks) != 0 { + t.Errorf("GetTasksByFilter ids task got = %v, want %v", len(tasks), 0) + } + }) + + t.Run("GetTasksByFilter status", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + Statuses: []base.Status{base.DownloadStatusDone}, + }) + if len(tasks) != len(reqs) { + t.Errorf("GetTasksByFilter status task got = %v, want %v", len(tasks), len(reqs)) + } + }) + + t.Run("GetTasksByFilter not match status", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + Statuses: []base.Status{base.DownloadStatusError}, + }) + if len(tasks) != 0 { + t.Errorf("GetTasksByFilter status task got = %v, want %v", len(tasks), 0) + } + }) + + t.Run("GetTasksByFilter match notStatus", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + NotStatuses: []base.Status{base.DownloadStatusRunning, base.DownloadStatusPause}, + }) + if len(tasks) != len(reqs) { + t.Errorf("GetTasksByFilter match notStatus task got = %v, want %v", len(tasks), len(reqs)) + } + }) + + t.Run("GetTasksByFilter not match notStatus", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + NotStatuses: []base.Status{base.DownloadStatusDone}, + }) + if len(tasks) != 0 { + t.Errorf("GetTasksByFilter not match notStatus task got = %v, want %v", len(tasks), 0) + } + }) + + t.Run("GetTasksByFilter match ids and status", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + IDs: []string{taskIds[0]}, + Statuses: []base.Status{base.DownloadStatusDone}, + }) + if len(tasks) != 1 { + t.Errorf("GetTasksByFilter match ids and status task got = %v, want %v", len(tasks), 1) + } + }) + + t.Run("GetTasksByFilter not match ids and status", func(t *testing.T) { + tasks := downloader.GetTasksByFilter(&TaskFilter{ + IDs: []string{taskIds[0]}, + Statuses: []base.Status{base.DownloadStatusError}, + }) + if len(tasks) != 0 { + t.Errorf("GetTasksByFilter not match ids and status task got = %v, want %v", len(tasks), 0) + } + }) + +} diff --git a/pkg/download/extension.go b/pkg/download/extension.go index 89bf1d21..5a64920f 100644 --- a/pkg/download/extension.go +++ b/pkg/download/extension.go @@ -720,11 +720,15 @@ func NewExtensionTask(download *Downloader, task *Task) *ExtensionTask { } func (t *ExtensionTask) Continue() error { - return t.download.Continue(t.ID) + return t.download.Continue(&TaskFilter{ + IDs: []string{t.ID}, + }) } func (t *ExtensionTask) Pause() error { - return t.download.Pause(t.ID) + return t.download.Pause(&TaskFilter{ + IDs: []string{t.ID}, + }) } func parseSettings(settings []*Setting) map[string]any { diff --git a/pkg/download/model.go b/pkg/download/model.go index 9b765994..085a20f0 100644 --- a/pkg/download/model.go +++ b/pkg/download/model.go @@ -112,6 +112,16 @@ func (t *Task) calcSpeed(speedArr []int64, downloaded int64, usedTime float64) i return int64(float64(total) / float64(len(speedArr)) / usedTime) } +type TaskFilter struct { + IDs []string + Statuses []base.Status + NotStatuses []base.Status +} + +func (f *TaskFilter) IsEmpty() bool { + return len(f.IDs) == 0 && len(f.Statuses) == 0 && len(f.NotStatuses) == 0 +} + type DownloaderConfig struct { Controller *controller.Controller FetchManagers []fetcher.FetcherManager diff --git a/pkg/rest/api.go b/pkg/rest/api.go index fde4c3cf..9e030b34 100644 --- a/pkg/rest/api.go +++ b/pkg/rest/api.go @@ -74,13 +74,27 @@ func CreateTaskBatch(w http.ResponseWriter, r *http.Request) { } func PauseTask(w http.ResponseWriter, r *http.Request) { - vars := mux.Vars(r) - taskId := vars["id"] - if taskId == "" { - WriteJson(w, model.NewErrorResult("param invalid: id", model.CodeInvalidParam)) + filter, errResult := parseIdFilter(r) + if errResult != nil { + WriteJson(w, errResult) return } - if err := Downloader.Pause(taskId); err != nil { + + if err := Downloader.Pause(filter); err != nil { + WriteJson(w, model.NewErrorResult(err.Error())) + return + } + WriteJson(w, model.NewNilResult()) +} + +func PauseTasks(w http.ResponseWriter, r *http.Request) { + filter, errResult := parseFilter(r) + if errResult != nil { + WriteJson(w, errResult) + return + } + + if err := Downloader.Pause(filter); err != nil { WriteJson(w, model.NewErrorResult(err.Error())) return } @@ -88,29 +102,27 @@ func PauseTask(w http.ResponseWriter, r *http.Request) { } func ContinueTask(w http.ResponseWriter, r *http.Request) { - vars := mux.Vars(r) - taskId := vars["id"] - if taskId == "" { - WriteJson(w, model.NewErrorResult("param invalid: id", model.CodeInvalidParam)) + filter, errResult := parseIdFilter(r) + if errResult != nil { + WriteJson(w, errResult) return } - if err := Downloader.Continue(taskId); err != nil { + + if err := Downloader.Continue(filter); err != nil { WriteJson(w, model.NewErrorResult(err.Error())) return } WriteJson(w, model.NewNilResult()) } -func PauseAllTask(w http.ResponseWriter, r *http.Request) { - if err := Downloader.PauseAll(); err != nil { - WriteJson(w, model.NewErrorResult(err.Error())) +func ContinueTasks(w http.ResponseWriter, r *http.Request) { + filter, errResult := parseFilter(r) + if errResult != nil { + WriteJson(w, errResult) return } - WriteJson(w, model.NewNilResult()) -} -func ContinueAllTask(w http.ResponseWriter, r *http.Request) { - if err := Downloader.ContinueAll(); err != nil { + if err := Downloader.Continue(filter); err != nil { WriteJson(w, model.NewErrorResult(err.Error())) return } @@ -118,14 +130,14 @@ func ContinueAllTask(w http.ResponseWriter, r *http.Request) { } func DeleteTask(w http.ResponseWriter, r *http.Request) { - vars := mux.Vars(r) - taskId := vars["id"] - force := r.FormValue("force") - if taskId == "" { - WriteJson(w, model.NewErrorResult("param invalid: id", model.CodeInvalidParam)) + filter, errResult := parseIdFilter(r) + if errResult != nil { + WriteJson(w, errResult) return } - if err := Downloader.Delete(taskId, force == "true"); err != nil { + force := r.FormValue("force") + + if err := Downloader.Delete(filter, force == "true"); err != nil { WriteJson(w, model.NewErrorResult(err.Error())) return } @@ -133,10 +145,14 @@ func DeleteTask(w http.ResponseWriter, r *http.Request) { } func DeleteTasks(w http.ResponseWriter, r *http.Request) { - r.ParseForm() - status := r.Form["status"] + filter, errResult := parseFilter(r) + if errResult != nil { + WriteJson(w, errResult) + return + } force := r.FormValue("force") - if err := Downloader.DeleteByStatues(convertStatues(status), force == "true"); err != nil { + + if err := Downloader.Delete(filter, force == "true"); err != nil { WriteJson(w, model.NewErrorResult(err.Error())) return } @@ -159,9 +175,13 @@ func GetTask(w http.ResponseWriter, r *http.Request) { } func GetTasks(w http.ResponseWriter, r *http.Request) { - r.ParseForm() - status := r.Form["status"] - tasks := Downloader.GetTasksByStatues(convertStatues(status)) + filter, errResult := parseFilter(r) + if errResult != nil { + WriteJson(w, errResult) + return + } + + tasks := Downloader.GetTasksByFilter(filter) WriteJson(w, model.NewOkResult(tasks)) } @@ -324,6 +344,32 @@ func GetStats(w http.ResponseWriter, r *http.Request) { WriteJson(w, model.NewOkResult(statsResult)) } +func parseIdFilter(r *http.Request) (*download.TaskFilter, any) { + vars := mux.Vars(r) + taskId := vars["id"] + if taskId == "" { + return nil, model.NewErrorResult("param invalid: id", model.CodeInvalidParam) + } + + filter := &download.TaskFilter{ + IDs: []string{taskId}, + } + return filter, nil +} + +func parseFilter(r *http.Request) (*download.TaskFilter, any) { + if err := r.ParseForm(); err != nil { + return nil, model.NewErrorResult(err.Error()) + } + + filter := &download.TaskFilter{ + IDs: r.Form["id"], + Statuses: convertStatues(r.Form["status"]), + NotStatuses: convertStatues(r.Form["notStatus"]), + } + return filter, nil +} + func convertStatues(statues []string) []base.Status { result := make([]base.Status, 0) for _, status := range statues { diff --git a/pkg/rest/server.go b/pkg/rest/server.go index 3bcfb8e0..9c0836d4 100644 --- a/pkg/rest/server.go +++ b/pkg/rest/server.go @@ -104,9 +104,9 @@ func BuildServer(startCfg *model.StartConfig) (*http.Server, net.Listener, error r.Methods(http.MethodPost).Path("/api/v1/tasks").HandlerFunc(CreateTask) r.Methods(http.MethodPost).Path("/api/v1/tasks/batch").HandlerFunc(CreateTaskBatch) r.Methods(http.MethodPut).Path("/api/v1/tasks/{id}/pause").HandlerFunc(PauseTask) + r.Methods(http.MethodPut).Path("/api/v1/tasks/pause").HandlerFunc(PauseTasks) r.Methods(http.MethodPut).Path("/api/v1/tasks/{id}/continue").HandlerFunc(ContinueTask) - r.Methods(http.MethodPut).Path("/api/v1/tasks/pause").HandlerFunc(PauseAllTask) - r.Methods(http.MethodPut).Path("/api/v1/tasks/continue").HandlerFunc(ContinueAllTask) + r.Methods(http.MethodPut).Path("/api/v1/tasks/continue").HandlerFunc(ContinueTasks) r.Methods(http.MethodDelete).Path("/api/v1/tasks/{id}").HandlerFunc(DeleteTask) r.Methods(http.MethodDelete).Path("/api/v1/tasks").HandlerFunc(DeleteTasks) r.Methods(http.MethodGet).Path("/api/v1/tasks/{id}").HandlerFunc(GetTask) diff --git a/pkg/rest/server_test.go b/pkg/rest/server_test.go index ff1b7f03..930d9817 100644 --- a/pkg/rest/server_test.go +++ b/pkg/rest/server_test.go @@ -582,8 +582,8 @@ func doTest(handler func()) { Downloader.Clear() }() defer func() { - Downloader.PauseAll() - Downloader.DeleteByStatues(nil, true) + Downloader.Pause(nil) + Downloader.Delete(nil, true) os.RemoveAll(cfg.StorageDir) }() taskReq.URL = "http://" + fileListener.Addr().String() + "/" + test.BuildName