add task cache to avoid too many requests about API

Muione · Jun 23, 2024 · dedce2f · dedce2f
1 parent 362d21e
commit dedce2f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 14 deletions.
diff --git a/internal/offline_download/pikpak/pikpak.go b/internal/offline_download/pikpak/pikpak.go
@@ -12,29 +12,33 @@ import (
 )
 
 type PikPak struct {
+	refreshTaskCache bool
 }
 
-func (p PikPak) Name() string {
+func (p *PikPak) Name() string {
 	return "pikpak"
 }
 
-func (p PikPak) Items() []model.SettingItem {
+func (p *PikPak) Items() []model.SettingItem {
 	return nil
 }
 
-func (p PikPak) Run(task *tool.DownloadTask) error {
+func (p *PikPak) Run(task *tool.DownloadTask) error {
 	return errs.NotSupport
 }
 
-func (p PikPak) Init() (string, error) {
+func (p *PikPak) Init() (string, error) {
+	p.refreshTaskCache = false
 	return "ok", nil
 }
 
-func (p PikPak) IsReady() bool {
+func (p *PikPak) IsReady() bool {
 	return true
 }
 
-func (p PikPak) AddURL(args *tool.AddUrlArgs) (string, error) {
+func (p *PikPak) AddURL(args *tool.AddUrlArgs) (string, error) {
+	// 添加新任务刷新缓存
+	p.refreshTaskCache = true
 	// args.TempDir 已经被修改为了 DstDirPath
 	storage, actualPath, err := op.GetStorageAndActualPath(args.TempDir)
 	if err != nil {
@@ -59,7 +63,7 @@ func (p PikPak) AddURL(args *tool.AddUrlArgs) (string, error) {
 	return t.ID, nil
 }
 
-func (p PikPak) Remove(task *tool.DownloadTask) error {
+func (p *PikPak) Remove(task *tool.DownloadTask) error {
 	storage, _, err := op.GetStorageAndActualPath(task.DstDirPath)
 	if err != nil {
 		return err
@@ -76,7 +80,7 @@ func (p PikPak) Remove(task *tool.DownloadTask) error {
 	return nil
 }
 
-func (p PikPak) Status(task *tool.DownloadTask) (*tool.Status, error) {
+func (p *PikPak) Status(task *tool.DownloadTask) (*tool.Status, error) {
 	storage, _, err := op.GetStorageAndActualPath(task.DstDirPath)
 	if err != nil {
 		return nil, err
@@ -85,7 +89,7 @@ func (p PikPak) Status(task *tool.DownloadTask) (*tool.Status, error) {
 	if !ok {
 		return nil, fmt.Errorf("unsupported storage driver for offline download, only Pikpak is supported")
 	}
-	tasks, err := GetTasks(pikpakDriver)
+	tasks, err := p.GetTasks(pikpakDriver)
 	if err != nil {
 		return nil, err
 	}

diff --git a/internal/offline_download/pikpak/util.go b/internal/offline_download/pikpak/util.go
@@ -1,28 +1,39 @@
 package pikpak
 
 import (
-	// "github.com/Xhofe/go-cache"
-
 	"context"
+	"time"
 
+	"github.com/Xhofe/go-cache"
 	"github.com/alist-org/alist/v3/drivers/pikpak"
 	"github.com/alist-org/alist/v3/internal/op"
-
 	"github.com/alist-org/alist/v3/pkg/singleflight"
 )
 
-// var taskCache = cache.NewMemCache(cache.WithShards[[]pikpak.OfflineTask](16))
+var taskCache = cache.NewMemCache(cache.WithShards[[]pikpak.OfflineTask](16))
 var taskG singleflight.Group[[]pikpak.OfflineTask]
 
-func GetTasks(pikpakDriver *pikpak.PikPak) ([]pikpak.OfflineTask, error) {
+func (p *PikPak) GetTasks(pikpakDriver *pikpak.PikPak) ([]pikpak.OfflineTask, error) {
 	key := op.Key(pikpakDriver, "/drive/v1/task")
+	if !p.refreshTaskCache {
+		if tasks, ok := taskCache.Get(key); ok {
+			return tasks, nil
+		}
+	}
+	p.refreshTaskCache = false
 	tasks, err, _ := taskG.Do(key, func() ([]pikpak.OfflineTask, error) {
 		ctx := context.Background()
 		phase := []string{"PHASE_TYPE_RUNNING", "PHASE_TYPE_ERROR", "PHASE_TYPE_PENDING", "PHASE_TYPE_COMPLETE"}
 		tasks, err := pikpakDriver.OfflineList(ctx, "", phase)
 		if err != nil {
 			return nil, err
 		}
+		// 添加缓存 15s
+		if len(tasks) > 0 {
+			taskCache.Set(key, tasks, cache.WithEx[[]pikpak.OfflineTask](time.Second*15))
+		} else {
+			taskCache.Del(key)
+		}
 		return tasks, nil
 	})
 	if err != nil {