Skip to content

Commit 37a6b26

Browse files
okatu-loliokatu-loli
committed
fix: handle aliyundrive open rate limits
1 parent 6bde813 commit 37a6b26

File tree

4 files changed

+260
-50
lines changed

4 files changed

+260
-50
lines changed

‎drivers/aliyundrive_open/driver.go‎

Lines changed: 19 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,10 @@ package aliyundrive_open
33
import (
44
"context"
55
"errors"
6-
"fmt"
76
"net/http"
87
"path/filepath"
98
"time"
109

11-
"github.com/Xhofe/rateg"
1210
"github.com/alist-org/alist/v3/drivers/base"
1311
"github.com/alist-org/alist/v3/internal/driver"
1412
"github.com/alist-org/alist/v3/internal/errs"
@@ -24,9 +22,8 @@ type AliyundriveOpen struct {
2422

2523
DriveId string
2624

27-
limitList func(ctx context.Context, data base.Json) (*Files, error)
28-
limitLink func(ctx context.Context, file model.Obj) (*model.Link, error)
29-
ref *AliyundriveOpen
25+
limiter *limiter
26+
ref *AliyundriveOpen
3027
}
3128

3229
func (d *AliyundriveOpen) Config() driver.Config {
@@ -38,25 +35,23 @@ func (d *AliyundriveOpen) GetAddition() driver.Additional {
3835
}
3936

4037
func (d *AliyundriveOpen) Init(ctx context.Context) error {
38+
d.limiter = getLimiterForUser(globalLimiterUserID)
4139
if d.LIVPDownloadFormat == "" {
4240
d.LIVPDownloadFormat = "jpeg"
4341
}
4442
if d.DriveType == "" {
4543
d.DriveType = "default"
4644
}
47-
res, err := d.request("/adrive/v1.0/user/getDriveInfo", http.MethodPost, nil)
45+
res, err := d.request(ctx, limiterOther, "/adrive/v1.0/user/getDriveInfo", http.MethodPost, nil)
4846
if err != nil {
47+
d.limiter.free()
48+
d.limiter = nil
4949
return err
5050
}
5151
d.DriveId = utils.Json.Get(res, d.DriveType+"_drive_id").ToString()
52-
d.limitList = rateg.LimitFnCtx(d.list, rateg.LimitFnOption{
53-
Limit: 4,
54-
Bucket: 1,
55-
})
56-
d.limitLink = rateg.LimitFnCtx(d.link, rateg.LimitFnOption{
57-
Limit: 1,
58-
Bucket: 1,
59-
})
52+
userID := utils.Json.Get(res, "user_id").ToString()
53+
d.limiter.free()
54+
d.limiter = getLimiterForUser(userID)
6055
return nil
6156
}
6257

@@ -70,6 +65,8 @@ func (d *AliyundriveOpen) InitReference(storage driver.Driver) error {
7065
}
7166

7267
func (d *AliyundriveOpen) Drop(ctx context.Context) error {
68+
d.limiter.free()
69+
d.limiter = nil
7370
d.ref = nil
7471
return nil
7572
}
@@ -87,9 +84,6 @@ func (d *AliyundriveOpen) GetRoot(ctx context.Context) (model.Obj, error) {
8784
}
8885

8986
func (d *AliyundriveOpen) List(ctx context.Context, dir model.Obj, args model.ListArgs) ([]model.Obj, error) {
90-
if d.limitList == nil {
91-
return nil, fmt.Errorf("driver not init")
92-
}
9387
files, err := d.getFiles(ctx, dir.GetID())
9488
if err != nil {
9589
return nil, err
@@ -108,7 +102,7 @@ func (d *AliyundriveOpen) List(ctx context.Context, dir model.Obj, args model.Li
108102
}
109103

110104
func (d *AliyundriveOpen) link(ctx context.Context, file model.Obj) (*model.Link, error) {
111-
res, err := d.request("/adrive/v1.0/openFile/getDownloadUrl", http.MethodPost, func(req *resty.Request) {
105+
res, err := d.request(ctx, limiterLink, "/adrive/v1.0/openFile/getDownloadUrl", http.MethodPost, func(req *resty.Request) {
112106
req.SetBody(base.Json{
113107
"drive_id": d.DriveId,
114108
"file_id": file.GetID(),
@@ -133,16 +127,13 @@ func (d *AliyundriveOpen) link(ctx context.Context, file model.Obj) (*model.Link
133127
}
134128

135129
func (d *AliyundriveOpen) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (*model.Link, error) {
136-
if d.limitLink == nil {
137-
return nil, fmt.Errorf("driver not init")
138-
}
139-
return d.limitLink(ctx, file)
130+
return d.link(ctx, file)
140131
}
141132

142133
func (d *AliyundriveOpen) MakeDir(ctx context.Context, parentDir model.Obj, dirName string) (model.Obj, error) {
143134
nowTime, _ := getNowTime()
144135
newDir := File{CreatedAt: nowTime, UpdatedAt: nowTime}
145-
_, err := d.request("/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
136+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
146137
req.SetBody(base.Json{
147138
"drive_id": d.DriveId,
148139
"parent_file_id": parentDir.GetID(),
@@ -168,7 +159,7 @@ func (d *AliyundriveOpen) MakeDir(ctx context.Context, parentDir model.Obj, dirN
168159

169160
func (d *AliyundriveOpen) Move(ctx context.Context, srcObj, dstDir model.Obj) (model.Obj, error) {
170161
var resp MoveOrCopyResp
171-
_, err := d.request("/adrive/v1.0/openFile/move", http.MethodPost, func(req *resty.Request) {
162+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/move", http.MethodPost, func(req *resty.Request) {
172163
req.SetBody(base.Json{
173164
"drive_id": d.DriveId,
174165
"file_id": srcObj.GetID(),
@@ -198,7 +189,7 @@ func (d *AliyundriveOpen) Move(ctx context.Context, srcObj, dstDir model.Obj) (m
198189

199190
func (d *AliyundriveOpen) Rename(ctx context.Context, srcObj model.Obj, newName string) (model.Obj, error) {
200191
var newFile File
201-
_, err := d.request("/adrive/v1.0/openFile/update", http.MethodPost, func(req *resty.Request) {
192+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/update", http.MethodPost, func(req *resty.Request) {
202193
req.SetBody(base.Json{
203194
"drive_id": d.DriveId,
204195
"file_id": srcObj.GetID(),
@@ -230,7 +221,7 @@ func (d *AliyundriveOpen) Rename(ctx context.Context, srcObj model.Obj, newName
230221

231222
func (d *AliyundriveOpen) Copy(ctx context.Context, srcObj, dstDir model.Obj) error {
232223
var resp MoveOrCopyResp
233-
_, err := d.request("/adrive/v1.0/openFile/copy", http.MethodPost, func(req *resty.Request) {
224+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/copy", http.MethodPost, func(req *resty.Request) {
234225
req.SetBody(base.Json{
235226
"drive_id": d.DriveId,
236227
"file_id": srcObj.GetID(),
@@ -256,7 +247,7 @@ func (d *AliyundriveOpen) Remove(ctx context.Context, obj model.Obj) error {
256247
if d.RemoveWay == "delete" {
257248
uri = "/adrive/v1.0/openFile/delete"
258249
}
259-
_, err := d.request(uri, http.MethodPost, func(req *resty.Request) {
250+
_, err := d.request(ctx, limiterOther, uri, http.MethodPost, func(req *resty.Request) {
260251
req.SetBody(base.Json{
261252
"drive_id": d.DriveId,
262253
"file_id": obj.GetID(),
@@ -295,7 +286,7 @@ func (d *AliyundriveOpen) Other(ctx context.Context, args model.OtherArgs) (inte
295286
default:
296287
return nil, errs.NotSupport
297288
}
298-
_, err := d.request(uri, http.MethodPost, func(req *resty.Request) {
289+
_, err := d.request(ctx, limiterOther, uri, http.MethodPost, func(req *resty.Request) {
299290
req.SetBody(data).SetResult(&resp)
300291
})
301292
if err != nil {

‎drivers/aliyundrive_open/limiter.go‎

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,97 @@
1+
package aliyundrive_open
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"sync"
7+
8+
"golang.org/x/time/rate"
9+
)
10+
11+
// Aliyun Open API rate limits are per user per app, so requests for the same
12+
// user should share one limiter across all storage instances.
13+
type limiterType int
14+
15+
const (
16+
limiterList limiterType = iota
17+
limiterLink
18+
limiterOther
19+
)
20+
21+
const (
22+
listRateLimit = 3.9
23+
linkRateLimit = 0.9
24+
otherRateLimit = 14.9
25+
globalLimiterUserID = ""
26+
)
27+
28+
type limiter struct {
29+
usedBy int
30+
list *rate.Limiter
31+
link *rate.Limiter
32+
other *rate.Limiter
33+
}
34+
35+
var (
36+
limiters = make(map[string]*limiter)
37+
limitersLock sync.Mutex
38+
)
39+
40+
func getLimiterForUser(userID string) *limiter {
41+
limitersLock.Lock()
42+
defer limitersLock.Unlock()
43+
defer func() {
44+
for id, lim := range limiters {
45+
if lim.usedBy <= 0 && id != globalLimiterUserID {
46+
delete(limiters, id)
47+
}
48+
}
49+
}()
50+
if lim, ok := limiters[userID]; ok {
51+
lim.usedBy++
52+
return lim
53+
}
54+
lim := &limiter{
55+
usedBy: 1,
56+
list: rate.NewLimiter(rate.Limit(listRateLimit), 1),
57+
link: rate.NewLimiter(rate.Limit(linkRateLimit), 1),
58+
other: rate.NewLimiter(rate.Limit(otherRateLimit), 1),
59+
}
60+
limiters[userID] = lim
61+
return lim
62+
}
63+
64+
func (l *limiter) wait(ctx context.Context, typ limiterType) error {
65+
if l == nil {
66+
return fmt.Errorf("driver not init")
67+
}
68+
switch typ {
69+
case limiterList:
70+
return l.list.Wait(ctx)
71+
case limiterLink:
72+
return l.link.Wait(ctx)
73+
case limiterOther:
74+
return l.other.Wait(ctx)
75+
default:
76+
return fmt.Errorf("unknown limiter type")
77+
}
78+
}
79+
80+
func (l *limiter) free() {
81+
if l == nil {
82+
return
83+
}
84+
limitersLock.Lock()
85+
defer limitersLock.Unlock()
86+
l.usedBy--
87+
}
88+
89+
func (d *AliyundriveOpen) wait(ctx context.Context, typ limiterType) error {
90+
if d == nil {
91+
return fmt.Errorf("driver not init")
92+
}
93+
if d.ref != nil {
94+
return d.ref.wait(ctx, typ)
95+
}
96+
return d.limiter.wait(ctx, typ)
97+
}

‎drivers/aliyundrive_open/upload.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -50,10 +50,10 @@ func calPartSize(fileSize int64) int64 {
5050
return partSize
5151
}
5252

53-
func (d *AliyundriveOpen) getUploadUrl(count int, fileId, uploadId string) ([]PartInfo, error) {
53+
func (d *AliyundriveOpen) getUploadUrl(ctx context.Context, count int, fileId, uploadId string) ([]PartInfo, error) {
5454
partInfoList := makePartInfos(count)
5555
var resp CreateResp
56-
_, err := d.request("/adrive/v1.0/openFile/getUploadUrl", http.MethodPost, func(req *resty.Request) {
56+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/getUploadUrl", http.MethodPost, func(req *resty.Request) {
5757
req.SetBody(base.Json{
5858
"drive_id": d.DriveId,
5959
"file_id": fileId,
@@ -84,10 +84,10 @@ func (d *AliyundriveOpen) uploadPart(ctx context.Context, r io.Reader, partInfo
8484
return nil
8585
}
8686

87-
func (d *AliyundriveOpen) completeUpload(fileId, uploadId string) (model.Obj, error) {
87+
func (d *AliyundriveOpen) completeUpload(ctx context.Context, fileId, uploadId string) (model.Obj, error) {
8888
// 3. complete
8989
var newFile File
90-
_, err := d.request("/adrive/v1.0/openFile/complete", http.MethodPost, func(req *resty.Request) {
90+
_, err := d.request(ctx, limiterOther, "/adrive/v1.0/openFile/complete", http.MethodPost, func(req *resty.Request) {
9191
req.SetBody(base.Json{
9292
"drive_id": d.DriveId,
9393
"file_id": fileId,
@@ -183,7 +183,7 @@ func (d *AliyundriveOpen) upload(ctx context.Context, dstDir model.Obj, stream m
183183
createData["pre_hash"] = hash
184184
}
185185
var createResp CreateResp
186-
_, err, e := d.requestReturnErrResp("/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
186+
_, err, e := d.requestReturnErrResp(ctx, limiterOther, "/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
187187
req.SetBody(createData).SetResult(&createResp)
188188
})
189189
if err != nil {
@@ -208,7 +208,7 @@ func (d *AliyundriveOpen) upload(ctx context.Context, dstDir model.Obj, stream m
208208
if err != nil {
209209
return nil, fmt.Errorf("cal proof code error: %s", err.Error())
210210
}
211-
_, err = d.request("/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
211+
_, err = d.request(ctx, limiterOther, "/adrive/v1.0/openFile/create", http.MethodPost, func(req *resty.Request) {
212212
req.SetBody(createData).SetResult(&createResp)
213213
})
214214
if err != nil {
@@ -229,7 +229,7 @@ func (d *AliyundriveOpen) upload(ctx context.Context, dstDir model.Obj, stream m
229229
}
230230
// refresh upload url if 50 minutes passed
231231
if time.Since(preTime) > 50*time.Minute {
232-
createResp.PartInfoList, err = d.getUploadUrl(count, createResp.FileId, createResp.UploadId)
232+
createResp.PartInfoList, err = d.getUploadUrl(ctx, count, createResp.FileId, createResp.UploadId)
233233
if err != nil {
234234
return nil, err
235235
}
@@ -266,5 +266,5 @@ func (d *AliyundriveOpen) upload(ctx context.Context, dstDir model.Obj, stream m
266266

267267
log.Debugf("[aliyundrive_open] create file success, resp: %+v", createResp)
268268
// 3. complete
269-
return d.completeUpload(createResp.FileId, createResp.UploadId)
269+
return d.completeUpload(ctx, createResp.FileId, createResp.UploadId)
270270
}

0 commit comments

Comments
 (0)