feat(local): allow to migrate static token buckets

drivers/local/driver.go

@@ -76,7 +76,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
 }

drivers/local/token_bucket.go

@@ -23,6 +23,33 @@ 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{}{}
+			}
+
+			go func() {
+				for range migrateSize {
+					<-oldStaticBucket.bucket
+					bucket <- struct{}{}
+				}
+			}()
+			return StaticTokenBucket{bucket: bucket}
+		}
+	}
+	return NewStaticTokenBucket(size)
+}
+
 func (b StaticTokenBucket) Take() <-chan struct{} {
 	return b.bucket
 }