Skip to content

Commit

Permalink
feat(local):thumbnail token bucket 平滑迁移 (AlistGo#7425))
Browse files Browse the repository at this point in the history
  • Loading branch information
long2005a1 committed Nov 9, 2024
1 parent ecb075e commit 924603b
Show file tree
Hide file tree
Showing 2 changed files with 37 additions and 3 deletions.
2 changes: 1 addition & 1 deletion drivers/local/driver.go
Original file line number Diff line number Diff line change
Expand Up @@ -77,7 +77,7 @@ func (d *Local) Init(ctx context.Context) error {
if d.thumbConcurrency == 0 {
d.thumbTokenBucket = NewNopTokenBucket()
} else {
d.thumbTokenBucket = NewStaticTokenBucket(d.thumbConcurrency)
d.thumbTokenBucket = NewStaticTokenBucketWithMigration(d.thumbTokenBucket, d.thumbConcurrency)
}
return nil
}
Expand Down
38 changes: 36 additions & 2 deletions drivers/local/token_bucket.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,38 @@ func NewStaticTokenBucket(size int) StaticTokenBucket {
return StaticTokenBucket{bucket: bucket}
}

func NewStaticTokenBucketWithMigration(oldBucket TokenBucket, size int) StaticTokenBucket {
if oldBucket != nil {
oldStaticBucket, ok := oldBucket.(StaticTokenBucket)
if ok {
oldSize := cap(oldStaticBucket.bucket)
migrateSize := oldSize
if size < migrateSize {
migrateSize = size
}

bucket := make(chan struct{}, size)
for range size - migrateSize {
bucket <- struct{}{}
}

if migrateSize != 0 {
go func() {
for range migrateSize {
<-oldStaticBucket.bucket
bucket <- struct{}{}
}
close(oldStaticBucket.bucket)
}()
}
return StaticTokenBucket{bucket: bucket}
}
}
return NewStaticTokenBucket(size)
}

// Take channel maybe closed when local driver is modified.
// don't call Put method after the channel is closed.
func (b StaticTokenBucket) Take() <-chan struct{} {
return b.bucket
}
Expand All @@ -35,8 +67,10 @@ func (b StaticTokenBucket) Do(ctx context.Context, f func() error) error {
select {
case <-ctx.Done():
return ctx.Err()
case <-b.bucket:
defer b.Put()
case _, ok := <-b.Take():
if ok {
defer b.Put()
}
}
return f()
}
Expand Down

0 comments on commit 924603b

Please sign in to comment.