perf(123pan): optimize rate limiting (#6859)

- eliminating fixed 200 ms delay in getFiles to prevent thread starvation
- allowing cancellation via context to mitigate potential DoS attacks by immediately cancelling excessive requests
pull/6878/head
seiuneko 2024-07-25 20:08:59 +08:00 committed by GitHub
parent d4e3355f56
commit 5fa70e4010
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 23 additions and 24 deletions

View File

@ -53,7 +53,7 @@ func (d *Pan123) Drop(ctx context.Context) error {
} }
func (d *Pan123) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) { func (d *Pan123) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
files, err := d.getFiles(dir.GetID(), dir.GetName()) files, err := d.getFiles(ctx, dir.GetID(), dir.GetName())
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -247,9 +247,6 @@ func (d *Pan123) Put(ctx context.Context, dstDir model.Obj, stream model.FileStr
} }
_, err = uploader.UploadWithContext(ctx, input) _, err = uploader.UploadWithContext(ctx, input)
} }
if err != nil {
return err
}
_, err = d.request(UploadComplete, http.MethodPost, func(req *resty.Request) { _, err = d.request(UploadComplete, http.MethodPost, func(req *resty.Request) {
req.SetBody(base.Json{ req.SetBody(base.Json{
"fileId": resp.Data.FileId, "fileId": resp.Data.FileId,
@ -258,11 +255,12 @@ func (d *Pan123) Put(ctx context.Context, dstDir model.Obj, stream model.FileStr
return err return err
} }
func (d *Pan123) APIRateLimit(api string) bool { func (d *Pan123) APIRateLimit(ctx context.Context, api string) error {
limiter, _ := d.apiRateLimit.LoadOrStore(api, value, _ := d.apiRateLimit.LoadOrStore(api,
rate.NewLimiter(rate.Every(time.Millisecond*700), 1)) rate.NewLimiter(rate.Every(700*time.Millisecond), 1))
ins := limiter.(*rate.Limiter) limiter := value.(*rate.Limiter)
return ins.Allow()
return limiter.Wait(ctx)
} }
var _ driver.Driver = (*Pan123)(nil) var _ driver.Driver = (*Pan123)(nil)

View File

@ -1,6 +1,7 @@
package _123 package _123
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"hash/crc32" "hash/crc32"
@ -14,7 +15,7 @@ import (
"github.com/alist-org/alist/v3/drivers/base" "github.com/alist-org/alist/v3/drivers/base"
"github.com/alist-org/alist/v3/pkg/utils" "github.com/alist-org/alist/v3/pkg/utils"
resty "github.com/go-resty/resty/v2" "github.com/go-resty/resty/v2"
jsoniter "github.com/json-iterator/go" jsoniter "github.com/json-iterator/go"
log "github.com/sirupsen/logrus" log "github.com/sirupsen/logrus"
) )
@ -233,15 +234,14 @@ func (d *Pan123) request(url string, method string, callback base.ReqCallback, r
return body, nil return body, nil
} }
func (d *Pan123) getFiles(parentId string, name string) ([]File, error) { func (d *Pan123) getFiles(ctx context.Context, parentId string, name string) ([]File, error) {
page := 1 page := 1
total := 0 total := 0
res := make([]File, 0) res := make([]File, 0)
// 2024-02-06 fix concurrency by 123pan // 2024-02-06 fix concurrency by 123pan
for { for {
if !d.APIRateLimit(FileList) { if err := d.APIRateLimit(ctx, FileList); err != nil {
time.Sleep(time.Millisecond * 200) return nil, err
continue
} }
var resp Files var resp Files
query := map[string]string{ query := map[string]string{

View File

@ -45,7 +45,7 @@ func (d *Pan123Share) Drop(ctx context.Context) error {
func (d *Pan123Share) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) { func (d *Pan123Share) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
// TODO return the files list, required // TODO return the files list, required
files, err := d.getFiles(dir.GetID()) files, err := d.getFiles(ctx, dir.GetID())
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -150,11 +150,12 @@ func (d *Pan123Share) Put(ctx context.Context, dstDir model.Obj, stream model.Fi
// return nil, errs.NotSupport // return nil, errs.NotSupport
//} //}
func (d *Pan123Share) APIRateLimit(api string) bool { func (d *Pan123Share) APIRateLimit(ctx context.Context, api string) error {
limiter, _ := d.apiRateLimit.LoadOrStore(api, value, _ := d.apiRateLimit.LoadOrStore(api,
rate.NewLimiter(rate.Every(time.Millisecond*700), 1)) rate.NewLimiter(rate.Every(700*time.Millisecond), 1))
ins := limiter.(*rate.Limiter) limiter := value.(*rate.Limiter)
return ins.Allow()
return limiter.Wait(ctx)
} }
var _ driver.Driver = (*Pan123Share)(nil) var _ driver.Driver = (*Pan123Share)(nil)

View File

@ -1,6 +1,7 @@
package _123Share package _123Share
import ( import (
"context"
"errors" "errors"
"fmt" "fmt"
"hash/crc32" "hash/crc32"
@ -80,13 +81,12 @@ func (d *Pan123Share) request(url string, method string, callback base.ReqCallba
return body, nil return body, nil
} }
func (d *Pan123Share) getFiles(parentId string) ([]File, error) { func (d *Pan123Share) getFiles(ctx context.Context, parentId string) ([]File, error) {
page := 1 page := 1
res := make([]File, 0) res := make([]File, 0)
for { for {
if !d.APIRateLimit(FileList) { if err := d.APIRateLimit(ctx, FileList); err != nil {
time.Sleep(time.Millisecond * 200) return nil, err
continue
} }
var resp Files var resp Files
query := map[string]string{ query := map[string]string{