Golang限流库与漏桶和令牌桶的使用介绍

// New returns a Limiter that will limit to the given RPS.
func New(rate int, opts ...Option) Limiter {
	return newAtomicBased(rate, opts...)
}
// newAtomicBased returns a new atomic based limiter.
func newAtomicBased(rate int, opts ...Option) *atomicLimiter {
	// TODO consider moving config building to the implementation
	// independent code.
	config := buildConfig(opts)
	perRequest := config.per / time.Duration(rate)
	l := &atomicLimiter{
		perRequest: perRequest,
		maxSlack:   -1 * time.Duration(config.slack) * perRequest,
		clock:      config.clock,
	}
	initialState := state{
		last:     time.Time{},
		sleepFor: 0,
	}
	atomic.StorePointer(&l.state, unsafe.Pointer(&initialState))
	return l
}

// Clock is the minimum necessary interface to instantiate a rate limiter with
// a clock or mock clock, compatible with clocks created using
// github.com/andres-erbsen/clock.
type Clock interface {
   Now() time.Time
   Sleep(time.Duration)
}

type state struct {
   last     time.Time
   sleepFor time.Duration
}

// Take blocks to ensure that the time spent between multiple
// Take calls is on average time.Second/rate.
func (t *atomicLimiter) Take() time.Time {
   var (
      newState state
      taken    bool
      interval time.Duration
   )
   for !taken {
      now := t.clock.Now()
      previousStatePointer := atomic.LoadPointer(&t.state)
      oldState := (*state)(previousStatePointer)
      newState = state{
         last:     now,
         sleepFor: oldState.sleepFor,
      }
      // If this is our first request, then we allow it.
      if oldState.last.IsZero() {
         taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
         continue
      }
      // 计算是否需要进行等待取水操作
      newState.sleepFor += t.perRequest(每两滴水之间的间隔时间) - now.Sub(oldState.last)(当前时间与上次取水时间的间隔)
       // 如果等待取水时间特别小，就需要松紧度进行维护
      if newState.sleepFor < t.maxSlack {
         newState.sleepFor = t.maxSlack
      }
       // 如果等待时间大于0，就进行更新
      if newState.sleepFor > 0 {
         newState.last = newState.last.Add(newState.sleepFor)
         interval, newState.sleepFor = newState.sleepFor, 0
      }
      taken = atomic.CompareAndSwapPointer(&t.state, previousStatePointer, unsafe.Pointer(&newState))
   }
   t.clock.Sleep(interval)
   // 最后返回需要等待的时间
    return newState.last
}

func (c *clock) Sleep(d time.Duration) {<!--{C}%3C!%2D%2D%20%2D%2D%3E--> time.Sleep(d) }

func ratelimit1() func(ctx *gin.Context) {
	r1 := rate1.New(100)
	return func(ctx *gin.Context) {
		now := time.Now()
		//  Take 返回的是一个 time.Duration的时间
		if r1.Take().Sub(now) > 0 {
			// 返回的时间比当前的时间还大，说明需要进行等待
			// 如果需要等待, 就 time.Sleep(r1.Take().Sub(now())) 然后放行
			// 如果不需要等待请求时间，就直接进行Abort 然后返回
			response(ctx, http.StatusRequestTimeout, "rate1 limit...")
			fmt.Println("rate1 limit...")
			ctx.Abort()
			return
		}
		// 放行
		ctx.Next()
	}
}

func response(c *gin.Context, code int, info any) {
   c.JSON(code, info)
}
func pingHandler(c *gin.Context) {
   response(c, 200, "ping ok~")
}

func TestRun(t *testing.T) {
   r := gin.Default()
   r.GET("/ping1", ratelimit1(), pingHandler)
   r.GET("/ping2", ratelimit2(), helloHandler)
   _ = r.Run(":4399")
}

// NewBucket returns a new token bucket that fills at the
// rate of one token every fillInterval, up to the given
// maximum capacity. Both arguments must be
// positive. The bucket is initially full.
func NewBucket(fillInterval time.Duration, capacity int64) *Bucket {
   return NewBucketWithClock(fillInterval, capacity, nil)
}
// NewBucketWithClock is identical to NewBucket but injects a testable clock
// interface.
func NewBucketWithClock(fillInterval time.Duration, capacity int64, clock Clock) *Bucket {
   return NewBucketWithQuantumAndClock(fillInterval, capacity, 1, clock)
}

/ NewBucketWithQuantumAndClock is like NewBucketWithQuantum, but
// also has a clock argument that allows clients to fake the passing
// of time. If clock is nil, the system clock will be used.
func NewBucketWithQuantumAndClock(fillInterval time.Duration, capacity, quantum int64, clock Clock) *Bucket {
   if clock == nil {
      clock = realClock{}
   }
    // 填充速率
   if fillInterval <= 0 {
      panic("token bucket fill interval is not > 0")
   }
    // 最大令牌容量
   if capacity <= 0 {
      panic("token bucket capacity is not > 0")
   }
    // 单次令牌生成量
   if quantum <= 0 {
      panic("token bucket quantum is not > 0")
   }
   return &Bucket{
      clock:           clock,
      startTime:       clock.Now(),
      latestTick:      0,
      fillInterval:    fillInterval,
      capacity:        capacity,
      quantum:         quantum,
      availableTokens: capacity,
   }
}

// TakeAvailable takes up to count immediately available tokens from the
// bucket. It returns the number of tokens removed, or zero if there are
// no available tokens. It does not block.
func (tb *Bucket) TakeAvailable(count int64) int64 {
   tb.mu.Lock()
   defer tb.mu.Unlock()
   return tb.takeAvailable(tb.clock.Now(), count)
}

// takeAvailable is the internal version of TakeAvailable - it takes the
// current time as an argument to enable easy testing.
func (tb *Bucket) takeAvailable(now time.Time, count int64) int64 {
   // 如果需要取出的令牌数小于等于零，那么就返回0个令牌
    if count <= 0 {
      return 0
   }
    // 根据时间对当前桶中令牌数进行计算
   tb.adjustavailableTokens(tb.currentTick(now))
    // 计算之后的令牌总数小于等于0，说明当前令牌不足取出，那么就直接返回0个令牌
   if tb.availableTokens <= 0 {
      return 0
   }
    // 如果当前存储的令牌数量多于请求数量，那么就返回取出令牌数
   if count > tb.availableTokens {
      count = tb.availableTokens
   }
    // 调整令牌数
   tb.availableTokens -= count
   return count
}

// adjustavailableTokens adjusts the current number of tokens
// available in the bucket at the given time, which must
// be in the future (positive) with respect to tb.latestTick.
func (tb *Bucket) adjustavailableTokens(tick int64) {
   lastTick := tb.latestTick
   tb.latestTick = tick
    // 如果当前令牌数大于最大等于容量，直接返回最大容量
   if tb.availableTokens >= tb.capacity {
      return
   }
    // 当前令牌数 += (当前时间 - 上次取出令牌数的时间) * quannum(每次生成令牌量)
   tb.availableTokens += (tick - lastTick) * tb.quantum
    // 如果当前令牌数大于最大等于容量, 将当前令牌数 = 最大容量 然后返回 当前令牌数
   if tb.availableTokens > tb.capacity {
      tb.availableTokens = tb.capacity
   }
   return
}

func ratelimit2() func(ctx *gin.Context) {
	// 生成速率 最大容量
	r2 := rate2.NewBucket(time.Second, 200)
	return func(ctx *gin.Context) {
		//r2.Take() // 允许欠账，令牌不够也可以接收请求
		if r2.TakeAvailable(1) == 1 {
			// 如果想要取出1个令牌并且能够取出，就放行
			ctx.Next()
			return
		}
		response(ctx, http.StatusRequestTimeout, "rate2 limit...")
		ctx.Abort()
		return
	}
}