Skip to content

Commit

Permalink
fix(local): thumbnails oom (#7124 close #7082)
Browse files Browse the repository at this point in the history
* add my_build.sh

* Fix OOM of thumbnail generation of LoaclDrive by using a task queue to control thread count

* remove my_build.sh

* chore(local): allow ThumbConcurrency set to zero

* revert(local): changes to thumbnail generating functions

* feat(local): implement static token bucket

* feat(local): use static token bucket to limit thumbnails generating concurrent

---------

Co-authored-by: KKJas <[email protected]>
  • Loading branch information
Mmx233 and Muione authored Sep 3, 2024
1 parent 34ada81 commit 4874c9e
Show file tree
Hide file tree
Showing 3 changed files with 85 additions and 2 deletions.
25 changes: 23 additions & 2 deletions drivers/local/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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
}

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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
}
Expand Down
1 change: 1 addition & 0 deletions drivers/local/meta.go
Original file line number Diff line number Diff line change
Expand Up @@ -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'"`
Expand Down
61 changes: 61 additions & 0 deletions drivers/local/token_bucket.go
Original file line number Diff line number Diff line change
@@ -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() }

0 comments on commit 4874c9e

Please sign in to comment.