feat(local): implement static token bucket

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() }