Fix OOM of thumbnail generation

drivers/local/driver.go

@@ -30,6 +30,10 @@ type Local struct {
 	model.Storage
 	Addition
 	mkdirPerm int32
+
+	// zero means no limit
+	thumbConcurrency int
+	thumbTokenBucket TokenBucket
 }
 
 func (d *Local) Config() driver.Config {
@@ -62,6 +66,18 @@ func (d *Local) Init(ctx context.Context) error {
 			return err
 		}
 	}
+	if d.ThumbConcurrency != "" {
+		v, err := strconv.ParseUint(d.ThumbConcurrency, 10, 32)
+		if err != nil {
+			return err
+		}
+		d.thumbConcurrency = int(v)
+	}
+	if d.thumbConcurrency == 0 {
+		d.thumbTokenBucket = NewNopTokenBucket()
+	} else {
+		d.thumbTokenBucket = NewStaticTokenBucket(d.thumbConcurrency)
+	}
 	return nil
 }
 
@@ -126,7 +142,6 @@ func (d *Local) FileInfoToObj(f fs.FileInfo, reqPath string, fullPath string) mo
 		},
 	}
 	return &file
-
 }
 func (d *Local) GetMeta(ctx context.Context, path string) (model.Obj, error) {
 	f, err := os.Stat(path)
@@ -178,7 +193,13 @@ func (d *Local) Link(ctx context.Context, file model.Obj, args model.LinkArgs) (
 	fullPath := file.GetPath()
 	var link model.Link
 	if args.Type == "thumb" && utils.Ext(file.GetName()) != "svg" {
-		buf, thumbPath, err := d.getThumb(file)
+		var buf *bytes.Buffer
+		var thumbPath *string
+		err := d.thumbTokenBucket.Do(ctx, func() error {
+			var err error
+			buf, thumbPath, err = d.getThumb(file)
+			return err
+		})
 		if err != nil {
 			return nil, err
 		}

drivers/local/meta.go

@@ -9,6 +9,7 @@ type Addition struct {
 	driver.RootPath
 	Thumbnail        bool   `json:"thumbnail" required:"true" help:"enable thumbnail"`
 	ThumbCacheFolder string `json:"thumb_cache_folder"`
+	ThumbConcurrency string `json:"thumb_concurrency" default:"16" required:"false" help:"Number of concurrent thumbnail generation goroutines. This controls how many thumbnails can be generated in parallel."`
 	ShowHidden       bool   `json:"show_hidden" default:"true" required:"false" help:"show hidden directories and files"`
 	MkdirPerm        string `json:"mkdir_perm" default:"777"`
 	RecycleBinPath   string `json:"recycle_bin_path" default:"delete permanently" help:"path to recycle bin, delete permanently if empty or keep 'delete permanently'"`

drivers/local/token_bucket.go

@@ -0,0 +1,61 @@
+package local
+
+import "context"
+
+type TokenBucket interface {
+	Take() <-chan struct{}
+	Put()
+	Do(context.Context, func() error) error
+}
+
+// StaticTokenBucket is a bucket with a fixed number of tokens,
+// where the retrieval and return of tokens are manually controlled.
+// In the initial state, the bucket is full.
+type StaticTokenBucket struct {
+	bucket chan struct{}
+}
+
+func NewStaticTokenBucket(size int) StaticTokenBucket {
+	bucket := make(chan struct{}, size)
+	for range size {
+		bucket <- struct{}{}
+	}
+	return StaticTokenBucket{bucket: bucket}
+}
+
+func (b StaticTokenBucket) Take() <-chan struct{} {
+	return b.bucket
+}
+
+func (b StaticTokenBucket) Put() {
+	b.bucket <- struct{}{}
+}
+
+func (b StaticTokenBucket) Do(ctx context.Context, f func() error) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case <-b.bucket:
+		defer b.Put()
+	}
+	return f()
+}
+
+// NopTokenBucket all function calls to this bucket will success immediately
+type NopTokenBucket struct {
+	nop chan struct{}
+}
+
+func NewNopTokenBucket() NopTokenBucket {
+	nop := make(chan struct{})
+	close(nop)
+	return NopTokenBucket{nop}
+}
+
+func (b NopTokenBucket) Take() <-chan struct{} {
+	return b.nop
+}
+
+func (b NopTokenBucket) Put() {}
+
+func (b NopTokenBucket) Do(_ context.Context, f func() error) error { return f() }

drivers/pikpak/driver.go

@@ -90,36 +90,43 @@ func (d *PikPak) Init(ctx context.Context) (err error) {
 
 	// 如果已经有RefreshToken，直接获取AccessToken
 	if d.Addition.RefreshToken != "" {
-		if d.RefreshTokenMethod == "oauth2" {
-			// 使用 oauth2 刷新令牌
-			// 初始化 oauth2Token
-			d.initializeOAuth2Token(ctx, oauth2Config, d.Addition.RefreshToken)
-			if err := d.refreshTokenByOAuth2(); err != nil {
-				return err
-			}
-		} else {
-			if err := d.refreshToken(d.Addition.RefreshToken); err != nil {
-				return err
-			}
-		}
-
+		// 使用 oauth2 刷新令牌
+		// 初始化 oauth2Token
+		d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) {
+			return oauth2Config.TokenSource(ctx, &oauth2.Token{
+				RefreshToken: d.Addition.RefreshToken,
+			}).Token()
+		}))
 	} else {
 		// 如果没有填写RefreshToken，尝试登录 获取 refreshToken
 		if err := d.login(); err != nil {
 			return err
 		}
-		if d.RefreshTokenMethod == "oauth2" {
-			d.initializeOAuth2Token(ctx, oauth2Config, d.RefreshToken)
-		}
+		d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) {
+			return oauth2Config.TokenSource(ctx, &oauth2.Token{
+				RefreshToken: d.RefreshToken,
+			}).Token()
+		}))
+	}
 
+	token, err := d.oauth2Token.Token()
+	if err != nil {
+		return err
 	}
+	d.RefreshToken = token.RefreshToken
+	d.AccessToken = token.AccessToken
 
 	// 获取CaptchaToken
-	err = d.RefreshCaptchaTokenAtLogin(GetAction(http.MethodGet, "https://api-drive.mypikpak.com/drive/v1/files"), d.Common.GetUserID())
+	err = d.RefreshCaptchaTokenAtLogin(GetAction(http.MethodGet, "https://api-drive.mypikpak.com/drive/v1/files"), d.Username)
 	if err != nil {
 		return err
 	}
 
+	// 获取用户ID
+	userID := token.Extra("sub").(string)
+	if userID != "" {
+		d.Common.SetUserID(userID)
+	}
 	// 更新UserAgent
 	if d.Platform == "android" {
 		d.Common.UserAgent = BuildCustomUserAgent(utils.GetMD5EncodeStr(d.Username+d.Password), AndroidClientID, AndroidPackageName, AndroidSdkVersion, AndroidClientVersion, AndroidPackageName, d.Common.UserID)

drivers/pikpak/meta.go

@@ -7,14 +7,13 @@ import (
 
 type Addition struct {
 	driver.RootID
-	Username           string `json:"username" required:"true"`
-	Password           string `json:"password" required:"true"`
-	Platform           string `json:"platform" required:"true" type:"select" options:"android,web"`
-	RefreshToken       string `json:"refresh_token" required:"true" default:""`
-	RefreshTokenMethod string `json:"refresh_token_method" required:"true" type:"select" options:"oauth2,http"`
-	CaptchaToken       string `json:"captcha_token" default:""`
-	DeviceID           string `json:"device_id"  required:"false" default:""`
-	DisableMediaLink   bool   `json:"disable_media_link" default:"true"`
+	Username         string `json:"username" required:"true"`
+	Password         string `json:"password" required:"true"`
+	Platform         string `json:"platform" required:"true" type:"select" options:"android,web"`
+	RefreshToken     string `json:"refresh_token" required:"true" default:""`
+	CaptchaToken     string `json:"captcha_token" default:""`
+	DeviceID         string `json:"device_id"  required:"false" default:""`
+	DisableMediaLink bool   `json:"disable_media_link" default:"true"`
 }
 
 var config = driver.Config{

drivers/pikpak/util.go

@@ -2,7 +2,6 @@ package pikpak
 
 import (
 	"bytes"
-	"context"
 	"crypto/md5"
 	"crypto/sha1"
 	"encoding/hex"
@@ -14,7 +13,6 @@ import (
 	"github.com/aliyun/aliyun-oss-go-sdk/oss"
 	jsoniter "github.com/json-iterator/go"
 	"github.com/pkg/errors"
-	"golang.org/x/oauth2"
 	"io"
 	"net/http"
 	"path/filepath"
@@ -114,63 +112,39 @@ func (d *PikPak) login() error {
 	return nil
 }
 
-func (d *PikPak) refreshToken(refreshToken string) error {
-	url := "https://user.mypikpak.com/v1/auth/token"
-	var e ErrResp
-	res, err := base.RestyClient.SetRetryCount(1).R().SetError(&e).
-		SetHeader("user-agent", "").SetBody(base.Json{
-		"client_id":     d.ClientID,
-		"client_secret": d.ClientSecret,
-		"grant_type":    "refresh_token",
-		"refresh_token": refreshToken,
-	}).SetQueryParam("client_id", d.ClientID).Post(url)
-	if err != nil {
-		d.Status = err.Error()
-		op.MustSaveDriverStorage(d)
-		return err
-	}
-	if e.ErrorCode != 0 {
-		if e.ErrorCode == 4126 {
-			// refresh_token invalid, re-login
-			return d.login()
-		}
-		d.Status = e.Error()
-		op.MustSaveDriverStorage(d)
-		return errors.New(e.Error())
-	}
-	data := res.Body()
-	d.Status = "work"
-	d.RefreshToken = jsoniter.Get(data, "refresh_token").ToString()
-	d.AccessToken = jsoniter.Get(data, "access_token").ToString()
-	d.Common.SetUserID(jsoniter.Get(data, "sub").ToString())
-	d.Addition.RefreshToken = d.RefreshToken
-	op.MustSaveDriverStorage(d)
-	return nil
-}
-
-func (d *PikPak) initializeOAuth2Token(ctx context.Context, oauth2Config *oauth2.Config, refreshToken string) {
-	d.oauth2Token = oauth2.ReuseTokenSource(nil, utils.TokenSource(func() (*oauth2.Token, error) {
-		return oauth2Config.TokenSource(ctx, &oauth2.Token{
-			RefreshToken: refreshToken,
-		}).Token()
-	}))
-}
-
-func (d *PikPak) refreshTokenByOAuth2() error {
-	token, err := d.oauth2Token.Token()
-	if err != nil {
-		return err
-	}
-	d.Status = "work"
-	d.RefreshToken = token.RefreshToken
-	d.AccessToken = token.AccessToken
-	// 获取用户ID
-	userID := token.Extra("sub").(string)
-	d.Common.SetUserID(userID)
-	d.Addition.RefreshToken = d.RefreshToken
-	op.MustSaveDriverStorage(d)
-	return nil
-}
+//func (d *PikPak) refreshToken() error {
+//	url := "https://user.mypikpak.com/v1/auth/token"
+//	var e ErrResp
+//	res, err := base.RestyClient.SetRetryCount(1).R().SetError(&e).
+//		SetHeader("user-agent", "").SetBody(base.Json{
+//		"client_id":     ClientID,
+//		"client_secret": ClientSecret,
+//		"grant_type":    "refresh_token",
+//		"refresh_token": d.RefreshToken,
+//	}).SetQueryParam("client_id", ClientID).Post(url)
+//	if err != nil {
+//		d.Status = err.Error()
+//		op.MustSaveDriverStorage(d)
+//		return err
+//	}
+//	if e.ErrorCode != 0 {
+//		if e.ErrorCode == 4126 {
+//			// refresh_token invalid, re-login
+//			return d.login()
+//		}
+//		d.Status = e.Error()
+//		op.MustSaveDriverStorage(d)
+//		return errors.New(e.Error())
+//	}
+//	data := res.Body()
+//	d.Status = "work"
+//	d.RefreshToken = jsoniter.Get(data, "refresh_token").ToString()
+//	d.AccessToken = jsoniter.Get(data, "access_token").ToString()
+//	d.Common.SetUserID(jsoniter.Get(data, "sub").ToString())
+//	d.Addition.RefreshToken = d.RefreshToken
+//	op.MustSaveDriverStorage(d)
+//	return nil
+//}
 
 func (d *PikPak) request(url string, method string, callback base.ReqCallback, resp interface{}) ([]byte, error) {
 	req := base.RestyClient.R()
@@ -207,19 +181,22 @@ func (d *PikPak) request(url string, method string, callback base.ReqCallback, r
 		return res.Body(), nil
 	case 4122, 4121, 16:
 		// access_token 过期
-		if d.RefreshTokenMethod == "oauth2" {
-			if err1 := d.refreshTokenByOAuth2(); err1 != nil {
-				return nil, err1
-			}
-		} else {
-			if err1 := d.refreshToken(d.RefreshToken); err1 != nil {
-				return nil, err1
-			}
+
+		//if err1 := d.refreshToken(); err1 != nil {
+		//	return nil, err1
+		//}
+		t, err := d.oauth2Token.Token()
+		if err != nil {
+			return nil, err
 		}
+		d.AccessToken = t.AccessToken
+		d.RefreshToken = t.RefreshToken
+		d.Addition.RefreshToken = t.RefreshToken
+		op.MustSaveDriverStorage(d)
 
 		return d.request(url, method, callback, resp)
 	case 9: // 验证码token过期
-		if err = d.RefreshCaptchaTokenAtLogin(GetAction(method, url), d.GetUserID()); err != nil {
+		if err = d.RefreshCaptchaTokenAtLogin(GetAction(method, url), d.Common.UserID); err != nil {
 			return nil, err
 		}
 		return d.request(url, method, callback, resp)
@@ -360,10 +337,6 @@ func (c *Common) GetDeviceID() string {
 	return c.DeviceID
 }
 
-func (c *Common) GetUserID() string {
-	return c.UserID
-}
-
 // RefreshCaptchaTokenAtLogin 刷新验证码token(登录后)
 func (d *PikPak) RefreshCaptchaTokenAtLogin(action, userID string) error {
 	metas := map[string]string{