From 8f4f7d12913fbfa745877a5d87c3e6857f923507 Mon Sep 17 00:00:00 2001
From: varg1714 <51254954+varg1714@users.noreply.github.com>
Date: Tue, 11 Nov 2025 21:59:10 +0800
Subject: [PATCH] feat(doubao): Add rate limiting (#1618)

---
 drivers/doubao/driver.go | 41 ++++++++++++++++++++++++++++++++++++++++
 drivers/doubao/meta.go   | 13 +++++++++----
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/drivers/doubao/driver.go b/drivers/doubao/driver.go
index d2ba04ea..6d4a4fc6 100644
--- a/drivers/doubao/driver.go
+++ b/drivers/doubao/driver.go
@@ -15,6 +15,7 @@ import (
 	"github.com/OpenListTeam/OpenList/v4/pkg/utils"
 	"github.com/go-resty/resty/v2"
 	"github.com/google/uuid"
+	"golang.org/x/time/rate"
 )
 
 type Doubao struct {
@@ -23,6 +24,7 @@ type Doubao struct {
 	*UploadToken
 	UserId       string
 	uploadThread int
+	limiter      *rate.Limiter
 }
 
 func (d *Doubao) Config() driver.Config {
@@ -61,6 +63,17 @@ func (d *Doubao) Init(ctx context.Context) error {
 		d.UploadToken = uploadToken
 	}
 
+	if d.LimitRate > 0 {
+		d.limiter = rate.NewLimiter(rate.Limit(d.LimitRate), 1)
+	}
+
+	return nil
+}
+
+func (d *Doubao) WaitLimit(ctx context.Context) error {
+	if d.limiter != nil {
+		return d.limiter.Wait(ctx)
+	}
 	return nil
 }
 
@@ -69,6 +82,10 @@ func (d *Doubao) Drop(ctx context.Context) error {
 }
 
 func (d *Doubao) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
+	if err := d.WaitLimit(ctx); err != nil {
+		return nil, err
+	}
+
 	var files []model.Obj
 	fileList, err := d.getFiles(dir.GetID(), "")
 	if err != nil {
@@ -95,6 +112,10 @@ func (d *Doubao) List(ctx context.Context, dir model.Obj, args model.ListArgs) (
 }
 
 func (d *Doubao) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
+	if err := d.WaitLimit(ctx); err != nil {
+		return nil, err
+	}
+
 	var downloadUrl string
 
 	if u, ok := file.(*Object); ok {
@@ -160,6 +181,10 @@ func (d *Doubao) Link(ctx context.Context, file model.Obj, args model.LinkArgs)
 }
 
 func (d *Doubao) MakeDir(ctx context.Context, parentDir model.Obj, dirName string) error {
+	if err := d.WaitLimit(ctx); err != nil {
+		return err
+	}
+
 	var r UploadNodeResp
 	_, err := d.request("/samantha/aispace/upload_node", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{
@@ -177,6 +202,10 @@ func (d *Doubao) MakeDir(ctx context.Context, parentDir model.Obj, dirName strin
 }
 
 func (d *Doubao) Move(ctx context.Context, srcObj, dstDir model.Obj) error {
+	if err := d.WaitLimit(ctx); err != nil {
+		return err
+	}
+
 	var r UploadNodeResp
 	_, err := d.request("/samantha/aispace/move_node", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{
@@ -191,6 +220,10 @@ func (d *Doubao) Move(ctx context.Context, srcObj, dstDir model.Obj) error {
 }
 
 func (d *Doubao) Rename(ctx context.Context, srcObj model.Obj, newName string) error {
+	if err := d.WaitLimit(ctx); err != nil {
+		return err
+	}
+
 	var r BaseResp
 	_, err := d.request("/samantha/aispace/rename_node", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{
@@ -207,6 +240,10 @@ func (d *Doubao) Copy(ctx context.Context, srcObj, dstDir model.Obj) (model.Obj,
 }
 
 func (d *Doubao) Remove(ctx context.Context, obj model.Obj) error {
+	if err := d.WaitLimit(ctx); err != nil {
+		return err
+	}
+
 	var r BaseResp
 	_, err := d.request("/samantha/aispace/delete_node", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{"node_list": []base.Json{{"id": obj.GetID()}}})
@@ -215,6 +252,10 @@ func (d *Doubao) Remove(ctx context.Context, obj model.Obj) error {
 }
 
 func (d *Doubao) Put(ctx context.Context, dstDir model.Obj, file model.FileStreamer, up driver.UpdateProgress) (model.Obj, error) {
+	if err := d.WaitLimit(ctx); err != nil {
+		return nil, err
+	}
+
 	// 根据MIME类型确定数据类型
 	mimetype := file.GetMimetype()
 	dataType := FileDataType
diff --git a/drivers/doubao/meta.go b/drivers/doubao/meta.go
index d2e350e3..83dade91 100644
--- a/drivers/doubao/meta.go
+++ b/drivers/doubao/meta.go
@@ -10,9 +10,10 @@ type Addition struct {
 	// driver.RootPath
 	driver.RootID
 	// define other
-	Cookie       string `json:"cookie" type:"text"`
-	UploadThread string `json:"upload_thread" default:"3"`
-	DownloadApi  string `json:"download_api" type:"select" options:"get_file_url,get_download_info" default:"get_file_url"`
+	Cookie       string  `json:"cookie" type:"text"`
+	UploadThread string  `json:"upload_thread" default:"3"`
+	DownloadApi  string  `json:"download_api" type:"select" options:"get_file_url,get_download_info" default:"get_file_url"`
+	LimitRate    float64 `json:"limit_rate" type:"float" default:"2" help:"limit all api request rate ([limit]r/1s)"`
 }
 
 var config = driver.Config{
@@ -23,6 +24,10 @@ var config = driver.Config{
 
 func init() {
 	op.RegisterDriver(func() driver.Driver {
-		return &Doubao{}
+		return &Doubao{
+			Addition: Addition{
+				LimitRate: 2,
+			},
+		}
 	})
 }