-
-
Notifications
You must be signed in to change notification settings - Fork 5.7k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
3 changed files
with
188 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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 | ||
| } |