naza/pkg/ratelimit/leakybucket.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"
	"sync"
	"time"
)

var ErrResourceNotAvailable = errors.New("naza.ratelimit: resource not available")

// 漏桶
type LeakyBucket struct {
	intervalMs int64

	mu       sync.Mutex
	lastTick int64
}

// @param intervalMs 多长时间以上，允许获取到一个资源，单位毫秒
func NewLeakyBucket(intervalMs int) *LeakyBucket {
	return &LeakyBucket{
		intervalMs: int64(intervalMs),
		// 注意，第一次获取资源，需要与创建对象时的时间点做比较
		lastTick: time.Now().UnixNano() / 1e6,
	}
}

// 尝试获取资源，获取成功返回nil，获取失败返回ErrResourceNotAvailable
// 如果获取失败，上层可自由选择多久后重试或丢弃本次任务
func (lb *LeakyBucket) TryAquire() error {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	nowMs := time.Now().UnixNano() / 1e6

	// 距离上次获取成功时间超过了间隔阈值，返回成功
	if nowMs-lb.lastTick > lb.intervalMs {
		lb.lastTick = nowMs
		return nil
	}

	return ErrResourceNotAvailable
}

// 阻塞直到获取到资源
func (lb *LeakyBucket) WaitUntilAquire() {
	lb.mu.Lock()
	nowMs := time.Now().UnixNano() / 1e6

	diff := nowMs - lb.lastTick
	if diff > lb.intervalMs {
		lb.lastTick = nowMs
		lb.mu.Unlock()
		return
	}

	// 没有达到间隔，我们更新lastTick再出锁，使得其他想获取资源的协程以新的lastTick作为判断条件
	lb.lastTick += lb.intervalMs
	lb.mu.Unlock()

	// 我们不需要等整个interval间隔，因为可能已经过去了一段时间了，
	// 注意，diff是根据更新前的lastTick计算得到的
	time.Sleep(time.Duration(lb.intervalMs-diff) * time.Millisecond)
	return
}

// 最快可获取到资源距离当前的时长， 但是不保证获取时一定能抢到
// 返回0，说明可以获取，返回非0，则是对应的时长，单位毫秒
func (lb *LeakyBucket) MaybeAvailableIntervalMs() int64 {
	lb.mu.Lock()
	defer lb.mu.Unlock()
	nowMs := time.Now().UnixNano() / 1e6

	if nowMs-lb.lastTick > lb.intervalMs {
		return 0
	}

	return lb.lastTick + lb.intervalMs - nowMs
}