令牌桶+漏桶算法

算法原理:

https://baike.baidu.com/item/漏桶算法
https://baike.baidu.com/item/令牌桶算法

代码实现参考:

令牌桶:
https://github.com/juju/ratelimit/blob/master/ratelimit.go
漏桶:
https://github.com/yangwenmai/ratelimit/blob/master/leakybucket/leaky_memory.go

令牌桶算法代码解析:

定义:

type Bucket struct {
	clock Clock

	// 第一次创建这个令牌桶的时间
	startTime time.Time

	// 最大容量
	capacity int64

	// 每次填充令牌的个数
	quantum int64

	// 填充的间隔
	fillInterval time.Duration

	// mu guards the fields below it.
	mu sync.Mutex

	// 当前可用的令牌数
	availableTokens int64

	// 最后一次填充的次数
	latestTick int64
}

获取令牌:

func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 {
	if count <= 0 {
		return 0
	}
	// 主动调整当前令牌数
	tb.adjustavailableTokens(tb.currentTick(now))
	if tb.availableTokens <= 0 {
		return 0
	}
	if count > tb.availableTokens {
		count = tb.availableTokens
	}
	tb.availableTokens -= count
	return count
}

调整令牌数:

// (当前时间-创建时间)/ 填充间隔 = 当前应该填充多少次
func (tb *Bucket) currentTick(now time.Time) int64 {
	return int64(now.Sub(tb.startTime) / tb.fillInterval)
}

func (tb *Bucket) adjustavailableTokens(tick int64) {
	if tb.availableTokens >= tb.capacity {
		return
	}
	// 当前应该填充次数-已经填充次数 = 需要调整填充次数
	// 增加的令牌数 = 需要调整填充次数 * 每次填充需要增加令牌的个数
	tb.availableTokens += (tick - tb.latestTick) * tb.quantum
	if tb.availableTokens > tb.capacity {
		tb.availableTokens = tb.capacity
	}
	// mark下最后一次调整的次数
	tb.latestTick = tick
	return
}

令牌桶分布式实现(基于redis):

  1. Bucket对象availableTokens和latestTick属性可以作为单独的key存储

漏桶算法代码解析

定义:

type bucket struct {
	// 容量
	capacity  uint
	// 剩余大小
	remaining uint
	// 下一次的重置容量时间
	reset     time.Time
	// 重置容量时间间隔
	rate      time.Duration
	mutex     sync.Mutex
}

往漏桶加请求:

// Add to the bucket.
func (b *bucket) Add(amount uint) (BucketState, error) {
	b.mutex.Lock()
	defer b.mutex.Unlock()
	if time.Now().After(b.reset) { // 需要重置
		b.reset = time.Now().Add(b.rate) // 更新时间
		b.remaining = b.capacity // 重置剩余容量
	}
	if amount > b.remaining { // 判断是否能过
		return BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, ErrorFull
	}
	b.remaining -= amount
	return BucketState{Capacity: b.capacity, Remaining: b.remaining, Reset: b.reset}, nil
}

你可能感兴趣的:(golang)