From c77ed5fcb00f4856899f34d0c3fc81fc91ce9ae1 Mon Sep 17 00:00:00 2001 From: Andy Hsu Date: Tue, 2 May 2023 22:06:03 +0800 Subject: [PATCH] feat(aliyundrive_open): limit rate for `List` and `Link` (close #4290) --- drivers/aliyundrive_open/driver.go | 17 ++++- pkg/utils/fn_limiter.go | 114 +++++++++++++++++++++++++++++ pkg/utils/fn_limiter_test.go | 59 +++++++++++++++ 3 files changed, 188 insertions(+), 2 deletions(-) create mode 100644 pkg/utils/fn_limiter.go create mode 100644 pkg/utils/fn_limiter_test.go diff --git a/drivers/aliyundrive_open/driver.go b/drivers/aliyundrive_open/driver.go index 735cda12a80..f85972b3752 100644 --- a/drivers/aliyundrive_open/driver.go +++ b/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{ diff --git a/pkg/utils/fn_limiter.go b/pkg/utils/fn_limiter.go new file mode 100644 index 00000000000..6f9d461fa69 --- /dev/null +++ b/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) + } +} diff --git a/pkg/utils/fn_limiter_test.go b/pkg/utils/fn_limiter_test.go new file mode 100644 index 00000000000..599cee95755 --- /dev/null +++ b/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 +}