feat(aliyundrive_open): limit rate for `List` and `Link` (close #4290)

drivers/aliyundrive_open/driver.go

@@ -21,6 +21,9 @@ type AliyundriveOpen struct {
 	base string
 
 	DriveId string
+
+	limitList func(ctx context.Context, dir model.Obj) ([]model.Obj, error)
+	limitLink func(ctx context.Context, file model.Obj) (*model.Link, error)
 }
 
 func (d *AliyundriveOpen) Config() driver.Config {
@@ -37,6 +40,8 @@ func (d *AliyundriveOpen) Init(ctx context.Context) error {
 		return err
 	}
 	d.DriveId = utils.Json.Get(res, "default_drive_id").ToString()
+	d.limitList = utils.LimitRateCtx(d.list, time.Second/4)
+	d.limitLink = utils.LimitRateCtx(d.link, time.Second)
 	return nil
 }
 
@@ -44,7 +49,7 @@ func (d *AliyundriveOpen) Drop(ctx context.Context) error {
 	return nil
 }
 
-func (d *AliyundriveOpen) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
+func (d *AliyundriveOpen) list(ctx context.Context, dir model.Obj) ([]model.Obj, error) {
 	files, err := d.getFiles(dir.GetID())
 	if err != nil {
 		return nil, err
@@ -54,7 +59,11 @@ func (d *AliyundriveOpen) List(ctx context.Context, dir model.Obj, args model.Li
 	})
 }
 
-func (d *AliyundriveOpen) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
+func (d *AliyundriveOpen) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
+	return d.limitList(ctx, dir)
+}
+
+func (d *AliyundriveOpen) link(ctx context.Context, file model.Obj) (*model.Link, error) {
 	res, err := d.request("/adrive/v1.0/openFile/getDownloadUrl", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{
 			"drive_id":   d.DriveId,
@@ -73,6 +82,10 @@ func (d *AliyundriveOpen) Link(ctx context.Context, file model.Obj, args model.L
 	}, nil
 }
 
+func (d *AliyundriveOpen) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
+	return d.limitLink(ctx, file)
+}
+
 func (d *AliyundriveOpen) MakeDir(ctx context.Context, parentDir model.Obj, dirName string) error {
 	_, err := d.request("/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
 		req.SetBody(base.Json{

pkg/utils/fn_limiter.go

@@ -0,0 +1,114 @@
+package utils
+
+import (
+	"context"
+	"reflect"
+	"time"
+)
+
+func LimitRateReflect(f interface{}, interval time.Duration) func(...interface{}) []interface{} {
+	// Use closures to save the time of the last function call
+	var lastCall time.Time
+
+	fValue := reflect.ValueOf(f)
+	fType := fValue.Type()
+
+	if fType.Kind() != reflect.Func {
+		panic("f must be a function")
+	}
+
+	//if fType.NumOut() == 0 {
+	//	panic("f must have at least one output parameter")
+	//}
+
+	outCount := fType.NumOut()
+	outTypes := make([]reflect.Type, outCount)
+
+	for i := 0; i < outCount; i++ {
+		outTypes[i] = fType.Out(i)
+	}
+
+	// Returns a new function, which is used to limit the function to be called only once at a specified time interval
+	return func(args ...interface{}) []interface{} {
+		// Calculate the time interval since the last function call
+		elapsed := time.Since(lastCall)
+		// If the interval is less than the specified time, wait for the remaining time
+		if elapsed < interval {
+			time.Sleep(interval - elapsed)
+		}
+		// Update the time of the last function call
+		lastCall = time.Now()
+
+		inCount := fType.NumIn()
+		in := make([]reflect.Value, inCount)
+
+		if len(args) != inCount {
+			panic("wrong number of arguments")
+		}
+
+		for i := 0; i < inCount; i++ {
+			in[i] = reflect.ValueOf(args[i])
+		}
+
+		out := fValue.Call(in)
+
+		if len(out) != outCount {
+			panic("function returned wrong number of values")
+		}
+
+		result := make([]interface{}, outCount)
+
+		for i := 0; i < outCount; i++ {
+			result[i] = out[i].Interface()
+		}
+
+		return result
+	}
+}
+
+type Fn[T any, R any] func(T) (R, error)
+type FnCtx[T any, R any] func(context.Context, T) (R, error)
+
+func LimitRate[T any, R any](f Fn[T, R], interval time.Duration) Fn[T, R] {
+	// Use closures to save the time of the last function call
+	var lastCall time.Time
+	// Returns a new function, which is used to limit the function to be called only once at a specified time interval
+	return func(t T) (R, error) {
+		// Calculate the time interval since the last function call
+		elapsed := time.Since(lastCall)
+		// If the interval is less than the specified time, wait for the remaining time
+		if elapsed < interval {
+			time.Sleep(interval - elapsed)
+		}
+		// Update the time of the last function call
+		lastCall = time.Now()
+		// Execute the function that needs to be limited
+		return f(t)
+	}
+}
+
+func LimitRateCtx[T any, R any](f FnCtx[T, R], interval time.Duration) FnCtx[T, R] {
+	// Use closures to save the time of the last function call
+	var lastCall time.Time
+	// Returns a new function, which is used to limit the function to be called only once at a specified time interval
+	return func(ctx context.Context, t T) (R, error) {
+		// Calculate the time interval since the last function call
+		elapsed := time.Since(lastCall)
+		// If the interval is less than the specified time, wait for the remaining time
+		if elapsed < interval {
+			t := time.NewTimer(interval - elapsed)
+			select {
+			case <-ctx.Done():
+				t.Stop()
+				var zero R
+				return zero, ctx.Err()
+			case <-t.C:
+
+			}
+		}
+		// Update the time of the last function call
+		lastCall = time.Now()
+		// Execute the function that needs to be limited
+		return f(ctx, t)
+	}
+}

pkg/utils/fn_limiter_test.go

@@ -0,0 +1,59 @@
+package utils_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/alist-org/alist/v3/pkg/utils"
+)
+
+func myFunction(a int) (int, error) {
+	// do something
+	return a + 1, nil
+}
+
+func TestLimitRate(t *testing.T) {
+	myLimitedFunction := utils.LimitRate(myFunction, time.Second)
+	result, _ := myLimitedFunction(1)
+	t.Log(result) // Output: 2
+	result, _ = myLimitedFunction(2)
+	t.Log(result) // Output: 3
+}
+
+type Test struct {
+	limitFn func(string) (string, error)
+}
+
+func (t *Test) myFunction(a string) (string, error) {
+	// do something
+	return a + " world", nil
+}
+
+func TestLimitRateStruct(t *testing.T) {
+	test := &Test{}
+	test.limitFn = utils.LimitRate(test.myFunction, time.Second)
+	result, _ := test.limitFn("hello")
+	t.Log(result) // Output: hello world
+	result, _ = test.limitFn("hi")
+	t.Log(result) // Output: hi world
+}
+
+func myFunctionCtx(ctx context.Context, a int) (int, error) {
+	// do something
+	return a + 1, nil
+}
+func TestLimitRateCtx(t *testing.T) {
+	myLimitedFunction := utils.LimitRateCtx(myFunctionCtx, time.Second)
+	result, _ := myLimitedFunction(context.Background(), 1)
+	t.Log(result) // Output: 2
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		time.Sleep(500 * time.Millisecond)
+		cancel()
+	}()
+	result, err := myLimitedFunction(ctx, 2)
+	t.Log(result, err) // Output: 0 context canceled
+	result, _ = myLimitedFunction(context.Background(), 3)
+	t.Log(result) // Output: 4
+}