golang高并发系统限流策略漏桶和令牌桶算法源码剖析

// 定义全局限流器对象
var rateLimit ratelimit.Limiter
// 在 gin.HandlerFunc 加入限流逻辑
func leakyBucket() gin.HandlerFunc {
	prev := time.Now()
	return func(c *gin.Context) {
		now := rateLimit.Take()
		fmt.Println(now.Sub(prev)) // 为了打印时间间隔
		prev = now // 记录上一次的时间，没有这个打印的会有问题
	}
}
func main() {
	rateLimit = ratelimit.New(10)
	r := gin.Default()
	r.GET("/ping", leakyBucket(), func(c *gin.Context) {
		c.JSON(200, true)
	})
	r.Run() // listen and serve on 0.0.0.0:8080 (for windows "localhost:8080")
}

type limiter struct {
	sync.Mutex
	last       time.Time
	sleepFor   time.Duration
	perRequest time.Duration
	maxSlack   time.Duration
	clock      Clock
}
type Limiter interface {
	// Take should block to make sure that the RPS is met.
	Take() time.Time
}

l := &limiter{
		perRequest: time.Second / time.Duration(rate),
		maxSlack:   -10 * time.Second / time.Duration(rate),
	}

func (t *limiter) Take() time.Time {
	t.Lock()
	defer t.Unlock()
	now := t.clock.Now()
	if t.last.IsZero() {
		t.last = now
		return t.last
	}
	t.sleepFor += t.perRequest - now.Sub(t.last)
	if t.sleepFor < t.maxSlack {
		t.sleepFor = t.maxSlack
	}
	if t.sleepFor > 0 {
		t.clock.Sleep(t.sleepFor)
		t.last = now.Add(t.sleepFor)
		t.sleepFor = 0
	} else {
		t.last = now
	}
	return t.last
}

import (
	"net/http"
	"time"
	"github.com/gin-gonic/gin"
	"golang.org/x/time/rate"
)
var rateLimit *rate.Limiter
func tokenBucket() gin.HandlerFunc {
	return func(c *gin.Context) {
		if rateLimit.Allow() {
			c.String(http.StatusOK, "rate limit,Drop")
			c.Abort()
			return
		}
		c.Next()
	}
}
func main() {
	limit := rate.Every(100 * time.Millisecond)
	rateLimit = rate.NewLimiter(limit, 10)
	r := gin.Default()
	r.GET("/ping", tokenBucket(), func(c *gin.Context) {
		c.jsON(200, true)
	})
	r.Run() // listen and serve on 0.0.0.0:8080 (for windows "localhost:8080")
}

type Limit float64
// Inf is the infinite rate limit; it allows all events (even if burst is zero).
const Inf = Limit(math.MaxFloat64)
// Every converts a minimum time interval between events to a Limit.
func Every(interval time.Duration) Limit {
	if interval <= 0 {
		return Inf
	}
	return 1 / Limit(interval.Seconds())
}

type Limiter struct {
	mu     sync.Mutex
	limit  Limit
	burst  int
	tokens float64
	// last is the last time the limiter's tokens field was updated
	last time.Time
	// lastEvent is the latest time of a rate-limited event (past or future)
	lastEvent time.Time
}

type Reservation struct {
	ok        bool
	lim       *Limiter
	tokens    int
	timeToAct time.Time
	// This is the Limit at reservation time, it can change later.
	limit Limit
}

func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
	//last不能在当前时间now之后，否则计算出来的elapsed为负数，会导致令牌桶数量减少
  last := lim.last
	if now.Before(last) {
		last = now
	}
	//根据令牌桶的缺数计算出令牌桶未进行更新的最大时间
	maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
	elapsed := now.Sub(last) //令牌桶未进行更新的时间段
	if elapsed > maxElapsed {
		elapsed = maxElapsed
	}
	//根据未更新的时间(未向桶中加入令牌的时间段)计算出产生的令牌数
	delta := lim.limit.tokensFromDuration(elapsed)
	tokens := lim.tokens + delta //计算出可用的令牌数
	if burst := float64(lim.burst); tokens > burst {
		tokens = burst
	}
	return now, last, tokens
}

// @param n 要消费的token数量
// @param maxFutureReserve 愿意等待的最长时间
func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
	lim.mu.Lock()
	// 如果没有限制
	if lim.limit == Inf {
		lim.mu.Unlock()
		return Reservation{
			ok:        true, //桶中有足够的令牌
			lim:       lim,
			tokens:    n,
			timeToAct: now,
		}
	}
	//更新令牌桶的状态，tokens为目前可用的令牌数量
	now, last, tokens := lim.advance(now)
  // 计算取完之后桶还能剩能下多少token
	tokens -= float64(n)
	var waitDuration time.Duration
  // 如果token < 0, 说明目前的token不够，需要等待一段时间
	if tokens < 0 {
		waitDuration = lim.limit.durationFromTokens(-tokens)
	}
	ok := n <= lim.burst && waitDuration <= maxFutureReserve
	r := Reservation{
		ok:    ok,
		lim:   lim,
		limit: lim.limit,
	}
  // timeToAct表示当桶中满足token数目等于n的时间
	if ok {
		r.tokens = n
		r.timeToAct = now.Add(waitDuration)
	}
  // 更新桶里面的token数目
	// 更新last时间
	// lastEvent
	if ok {
		lim.last = now
		lim.tokens = tokens
		lim.lastEvent = r.timeToAct
	} else {
		lim.last = last
	}
	lim.mu.Unlock()
	return r
}

func (limit Limit) durationFromTokens(tokens float64) time.Duration {
	seconds := tokens / float64(limit)
	return time.Nanosecond * time.Duration(1e9*seconds)
}
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
	// Split the integer and fractional parts ourself to minimize rounding errors.
	// See golang.org/issues/34861.
	sec := float64(d/time.Second) * float64(limit)
	nsec := float64(d%time.Second) * float64(limit)
	return sec + nsec/1e9
}

func (r *Reservation) CancelAt(now time.Time) {
	if !r.ok {
		return
	}
	r.lim.mu.Lock()
	defer r.lim.mu.Unlock()
  
	if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {
		return
	}
	//计算出需要还原的令牌数量
	//这里的r.lim.lastEvent可能是本次Reservation的结束时间，也可能是后来的Reservation的结束时间，所以要把本次结束时间点(r.timeToAct)之后产生的令牌数减去
	restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
  // 当小于0，表示已经都预支完了，不能归还了
	if restoreTokens <= 0 {
		return
	}
	//从新计算令牌桶的状态
	now, _, tokens := r.lim.advance(now)
	//还原当前令牌桶的令牌数量，当前的令牌数tokens加上需要还原的令牌数restoreTokens
	tokens += restoreTokens
  //如果tokens大于桶的最大容量，则将tokens置为桶的最大容量
	if burst := float64(r.lim.burst); tokens > burst {
		tokens = burst
	}
	// update state
	r.lim.last = now //记录桶的更新时间
	r.lim.tokens = tokens //更新令牌数量
 // 如果都相等，说明跟没消费一样。直接还原成上次的状态吧
	if r.timeToAct == r.lim.lastEvent {
		prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
		if !prevEvent.Before(now) {
			r.lim.lastEvent = prevEvent
		}
	}
	return
}