4874c9e43b
* 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 <75424880+Muione@users.noreply.github.com>
62 lines
1.3 KiB
Go
62 lines
1.3 KiB
Go
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() }
|