From 68c40d7f3f66857498e0eac7f7c5bdb7024487ad Mon Sep 17 00:00:00 2001 From: Levi Date: Fri, 17 May 2024 15:19:10 +0800 Subject: [PATCH] feat: add batch create and server info APIs (#540) --- pkg/download/downloader.go | 12 +++++++++++ pkg/download/downloader_test.go | 35 +++++++++++++++++++++++++++++++++ pkg/download/storage.go | 4 +++- pkg/rest/api.go | 28 ++++++++++++++++++++++++++ pkg/rest/model/action.go | 7 ------- pkg/rest/model/task.go | 5 +++++ pkg/rest/server.go | 2 ++ pkg/rest/server_test.go | 27 +++++++++++++++++++++++++ 8 files changed, 112 insertions(+), 8 deletions(-) delete mode 100644 pkg/rest/model/action.go diff --git a/pkg/download/downloader.go b/pkg/download/downloader.go index 59dfecdb..a1225c4e 100644 --- a/pkg/download/downloader.go +++ b/pkg/download/downloader.go @@ -301,6 +301,18 @@ func (d *Downloader) CreateDirect(req *base.Request, opts *base.Options) (taskId return d.doCreate(fetcher, opts) } +func (d *Downloader) CreateDirectBatch(reqs []*base.Request, opts *base.Options) (taskId []string, err error) { + taskIds := make([]string, 0) + for _, req := range reqs { + taskId, err := d.CreateDirect(req, opts) + if err != nil { + return nil, err + } + taskIds = append(taskIds, taskId) + } + return taskIds, nil +} + func (d *Downloader) Create(rrId string, opts *base.Options) (taskId string, err error) { d.fetcherMapLock.RLock() fetcher, ok := d.fetcherCache[rrId] diff --git a/pkg/download/downloader_test.go b/pkg/download/downloader_test.go index 3d9e3a29..8c4ca2ec 100644 --- a/pkg/download/downloader_test.go +++ b/pkg/download/downloader_test.go @@ -84,6 +84,41 @@ func TestDownloader_Create(t *testing.T) { } } +func TestDownloader_CreateDirectBatch(t *testing.T) { + listener := test.StartTestFileServer() + defer listener.Close() + + downloader := NewDownloader(nil) + if err := downloader.Setup(); err != nil { + t.Fatal(err) + } + defer downloader.Clear() + + reqs := make([]*base.Request, 0) + for i := 0; i < 6; i++ { + req := &base.Request{ + URL: "http://" + listener.Addr().String() + "/" + test.BuildName, + } + reqs = append(reqs, req) + } + + _, err := downloader.CreateDirectBatch(reqs, &base.Options{ + Path: test.Dir, + Name: test.DownloadName, + Extra: http.OptsExtra{ + Connections: 4, + }, + }) + if err != nil { + t.Fatal(err) + } + + tasks := downloader.GetTasks() + if len(tasks) != len(reqs) { + t.Errorf("CreateDirectBatch() got = %v, want %v", len(tasks), len(reqs)) + } +} + func TestDownloader_CreateWithProxy(t *testing.T) { // No proxy doTestDownloaderCreateWithProxy(t, false, func(proxyCfg *base.DownloaderProxyConfig) *base.DownloaderProxyConfig { diff --git a/pkg/download/storage.go b/pkg/download/storage.go index e615cdb6..f8def453 100644 --- a/pkg/download/storage.go +++ b/pkg/download/storage.go @@ -69,7 +69,9 @@ func (n *MemStorage) Setup(buckets []string) error { func (n *MemStorage) Put(bucket string, key string, v any) error { n.lock.Lock() defer n.lock.Unlock() - memData[bucket][key] = v + if bucketData, ok := memData[bucket]; ok { + bucketData[key] = v + } return nil } diff --git a/pkg/rest/api.go b/pkg/rest/api.go index 8a0ae75b..80b9bd6a 100644 --- a/pkg/rest/api.go +++ b/pkg/rest/api.go @@ -8,8 +8,20 @@ import ( "io" "net/http" "net/url" + "runtime" ) +func Info(w http.ResponseWriter, r *http.Request) { + info := map[string]any{ + "version": base.Version, + "runtime": runtime.Version(), + "os": runtime.GOOS, + "arch": runtime.GOARCH, + } + info["inDocker"] = base.InDocker == "true" + WriteJson(w, model.NewOkResult(info)) +} + func Resolve(w http.ResponseWriter, r *http.Request) { var req base.Request if ReadJson(r, w, &req) { @@ -45,6 +57,22 @@ func CreateTask(w http.ResponseWriter, r *http.Request) { } } +func CreateTaskBatch(w http.ResponseWriter, r *http.Request) { + var req model.CreateTaskBatch + if ReadJson(r, w, &req) { + if len(req.Reqs) == 0 { + WriteJson(w, model.NewErrorResult("param invalid: reqs", model.CodeInvalidParam)) + return + } + taskIds, err := Downloader.CreateDirectBatch(req.Reqs, req.Opt) + if err != nil { + WriteJson(w, model.NewErrorResult(err.Error())) + return + } + WriteJson(w, model.NewOkResult(taskIds)) + } +} + func PauseTask(w http.ResponseWriter, r *http.Request) { vars := mux.Vars(r) taskId := vars["id"] diff --git a/pkg/rest/model/action.go b/pkg/rest/model/action.go deleted file mode 100644 index ff29473e..00000000 --- a/pkg/rest/model/action.go +++ /dev/null @@ -1,7 +0,0 @@ -package model - -type Command struct { - Protocol string `json:"protocol"` - Action string `json:"action"` - Params any `json:"params"` -} diff --git a/pkg/rest/model/task.go b/pkg/rest/model/task.go index 4cd9e656..f93ccdb7 100644 --- a/pkg/rest/model/task.go +++ b/pkg/rest/model/task.go @@ -7,3 +7,8 @@ type CreateTask struct { Req *base.Request `json:"req"` Opt *base.Options `json:"opt"` } + +type CreateTaskBatch struct { + Reqs []*base.Request `json:"reqs"` + Opt *base.Options `json:"opt"` +} diff --git a/pkg/rest/server.go b/pkg/rest/server.go index 962c7b04..280c6821 100644 --- a/pkg/rest/server.go +++ b/pkg/rest/server.go @@ -88,8 +88,10 @@ func BuildServer(startCfg *model.StartConfig) (*http.Server, net.Listener, error } var r = mux.NewRouter() + r.Methods(http.MethodGet).Path("/api/v1/info").HandlerFunc(Info) r.Methods(http.MethodPost).Path("/api/v1/resolve").HandlerFunc(Resolve) 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/{id}/continue").HandlerFunc(ContinueTask) r.Methods(http.MethodPut).Path("/api/v1/tasks/pause").HandlerFunc(PauseAllTask) diff --git a/pkg/rest/server_test.go b/pkg/rest/server_test.go index ef6b8605..d591afad 100644 --- a/pkg/rest/server_test.go +++ b/pkg/rest/server_test.go @@ -63,6 +63,18 @@ var ( } ) +func TestInfo(t *testing.T) { + matchKeys := []string{"version", "runtime", "os", "arch", "inDocker"} + doTest(func() { + resp := httpRequestCheckOk[map[string]any](http.MethodGet, "/api/v1/info", nil) + for _, key := range matchKeys { + if _, ok := resp[key]; !ok { + t.Errorf("Info() missing key = %v", key) + } + } + }) +} + func TestResolve(t *testing.T) { doTest(func() { resp := httpRequestCheckOk[*download.ResolveResult](http.MethodPost, "/api/v1/resolve", taskReq) @@ -125,6 +137,21 @@ func TestCreateDirectTask(t *testing.T) { }) } +func TestCreateDirectTaskBatch(t *testing.T) { + doTest(func() { + reqs := make([]*base.Request, 0) + for i := 0; i < 6; i++ { + reqs = append(reqs, createReq.Req) + } + taskIds := httpRequestCheckOk[[]string](http.MethodPost, "/api/v1/tasks/batch", &model.CreateTaskBatch{ + Reqs: reqs, + }) + if len(taskIds) != len(reqs) { + t.Errorf("CreateDirectTaskBatch() got = %v, want %v", len(taskIds), len(reqs)) + } + }) +} + func TestCreateDirectTaskWithResoled(t *testing.T) { doTest(func() { var wg sync.WaitGroup