Goトークンバケットベースのフローリミッターの実装

package limiter

import (
 "fmt"
 "testing"
 "time"

 "golang.org/x/time/rate"
)

func TestLimter(t *testing.T) {
 limiter := rate.NewLimiter(rate.Every(time.Millisecond*31), 2)
 //Sleep(time.Second)
 for i := 0; i < 10; i++ {
  var ok bool
  if limiter.Allow() {
   ok = true
  }
  time.Sleep(time.Millisecond * 20)
  fmt.Println(ok, limiter.Burst())
 }
}

func NewLimiter(r Limit, b int) *Limiter {
 return &Limiter{
  limit: r,
  burst: b,
 }
}

// The methods AllowN, ReserveN, and WaitN consume n tokens.
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
Time }

type Reservation struct {
  ok bool // whether the condition is met to assign tokens
  lim *Limiter // Limiter to send the token
  tokens int // number of tokens
  timeToAct time.Time // the time to meet the token issuance
  limit Limit // speed of token issuance
}

// Allow is shorthand for AllowN(time.Now(), 1).
func (lim *Limiter) Allow() bool {
 return lim.AllowN(time.Now(), 1)
}

// AllowN reports whether n events may happen at time now.
// Use this method if you intend to drop / skip events that exceed the rate limit.
// Otherwise use Reserve or Wait.
func (lim *Limiter) AllowN(now time.Time, n int) bool {
 return lim.reserveN(now, n, 0).ok
}

// maxFutureReserve specifies the maximum reservation wait duration allowed.
// reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN.
func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
 lim.mu.Lock()
 // first determine if the frequency is infinite, if infinite, it means no limit for now
 if lim.limit == Inf {
  lim.mu.Unlock()
  return Reservation{
   ok: true,
   lim: lim,
   tokens: n,
   timeToAct: now,
  }
 }
   // The number of tokens that can be obtained as of now time, the last time the token was taken is last
 now, last, tokens := lim.advance(now)

 // Calculate the remaining number of tokens resulting from the request.
 // Update the number of tokens and remove the ones that need to be taken away
 tokens -= float64(n)

 // Calculate the wait duration
 // If the number of tokens is negative, it means it needs to wait, calculate the wait duration WaitDuration
 var waitDuration time.
 if tokens < 0 {
  waitDuration = lim.limit.durationFromTokens(-tokens)
 }

 // Decide result
 // Calculate if the allocation requirements are met
 // 1. the size to be allocated does not exceed the bucket capacity
 // 2. wait no longer than the set wait time
 ok := n <= lim.burst && waitDuration <= maxFutureReserve

 // Prepare reservation
 // finally construct a Resvervation object
 r := Reservation{
  ok: ok,
  lim: lim,
  limit: lim.limit,
 limit, }
 if ok {
  r.tokens = n
  r.timeToAct = now.Add(waitDuration)
 }

 // Update state
   // need to update the value of the current limit 
 if ok {
  lim.last = now
  lim.tokens = tokens
  lim.lastEvent = r.timeToAct
 } else {
  lim.last = last
 }

 lim.mu.Unlock()
 return r
}