naza/pkg/ratelimit/tokenbucket.go

// Copyright 2020, Chef.  All rights reserved.
// https://github.com/q191201771/naza
//
// Use of this source code is governed by a MIT-style license
// that can be found in the License file.
//
// Author: Chef (191201771@qq.com)

package ratelimit

import (
	"errors"
	"fmt"
	"sync"
	"time"

	"github.com/q191201771/naza/pkg/nazaatomic"
)

var ErrTokenNotEnough = errors.New("naza.ratelimit: token not enough")

// 令牌桶
type TokenBucket struct {
	capacity                  int
	prodTokenInterval         time.Duration
	prodTokenNumEveryInterval int

	disposeFlag nazaatomic.Bool

	mu        sync.Mutex
	available int
	cond      *sync.Cond
}

// @param capacity: 桶容量大小
// @param prodTokenIntervalMSec: 生产令牌的时间间隔，单位毫秒
// @param prodTokenNumEveryInterval: 每次生产多少个令牌
func NewTokenBucket(capacity int, prodTokenIntervalMSec int, prodTokenNumEveryInterval int) *TokenBucket {
	tb := &TokenBucket{
		capacity:                  capacity,
		prodTokenInterval:         time.Duration(time.Duration(prodTokenIntervalMSec) * time.Millisecond),
		prodTokenNumEveryInterval: prodTokenNumEveryInterval,
	}
	tb.cond = sync.NewCond(&tb.mu)
	tb.asyncProdToken()
	return tb
}

func (tb *TokenBucket) TryAquire() error {
	return tb.TryAquireWithNum(1)
}

func (tb *TokenBucket) WaitUntilAquire() {
	tb.WaitUntilAquireWithNum(1)
}

// 尝试获取相应数量的令牌，获取成功返回nil，获取失败返回ErrTokenNotEnough
// 如果获取失败，上层可自由选择多久后重试或丢弃本次任务
func (tb *TokenBucket) TryAquireWithNum(num int) error {
	tb.checkAquireNum(num)

	tb.mu.Lock()
	defer tb.mu.Unlock()
	if tb.available >= num {
		tb.available -= num
		return nil
	}

	return ErrTokenNotEnough
}

// 阻塞直到获取到相应数量的令牌
func (tb *TokenBucket) WaitUntilAquireWithNum(num int) {
	tb.checkAquireNum(num)

	for {
		tb.mu.Lock()
		if tb.available >= num {
			tb.available -= num
			tb.mu.Unlock()
			return
		}

		// 等待下次令牌生产时被唤醒
		// wait的内部会将自身添加到事件监听队列中然后释放锁，当接收到事件时，内部会重新获取锁然后返回
		tb.cond.Wait()
		tb.mu.Unlock()
	}
}

// 销毁令牌桶
func (tb *TokenBucket) Dispose() {
	tb.disposeFlag.Store(true)
}

func (tb *TokenBucket) asyncProdToken() {
	go func() {
		t := time.NewTicker(tb.prodTokenInterval)
		defer t.Stop()
		for {
			if tb.disposeFlag.Load() {
				break
			}
			select {
			case <-t.C:
				tb.mu.Lock()
				tb.available += tb.prodTokenNumEveryInterval
				if tb.available > tb.capacity {
					tb.available = tb.capacity
				}
				// It is allowed but not required for the caller to hold c.L
				// during the call.
				tb.cond.Broadcast()
				tb.mu.Unlock()
			}
		}
	}()
}

func (tb *TokenBucket) checkAquireNum(num int) {
	if num > tb.capacity {
		panic(fmt.Sprintf("aquire num should not bigger than capacity. num=%d, capacity=%d", num, tb.capacity))
	}
}