限流算法, 以 Golang 方式

type counter struct {  Maximal int  Period  time.Duration  count   int  tick    int64 // in nanosecond}
func (s *counter) Take(count int) bool {  if time.Now().UnixNano() > s.tick {    // if timeout, reset counter regally at first    s.count = 0    s.tick = time.Now().Add(s.Period).UnixNano()  }
  s.count += count            // it's acceptable in HPC scene  return s.count <= s.Maximal // it's acceptable in HPC scene}

// Makes it easy to test time based things.var now = time.Now
// LeakyBucket represents a bucket that leaks at a constant rate.type LeakyBucket struct {  // The identifying key, used for map lookups.  key string
  // How large the bucket is.  capacity int64
  // Amount the bucket leaks per second.  rate float64
  // The priority of the bucket in a min-heap priority queue, where p is the  // exact time the bucket will have leaked enough to be empty. Buckets that  // are empty or will be the soonest are at the top of the heap. This allows  // for quick pruning of empty buckets that scales very well. p is adjusted  // any time an amount is added to the Queue().  p time.Time
  // The index is maintained by the heap.Interface methods.  index int}
func NewLeakyBucket(rate float64, capacity int64) *LeakyBucket {  return &LeakyBucket{    rate:     rate,    capacity: capacity,    p:        now(),  }}
func (b *LeakyBucket) Add(amount int64) int64 {  count := b.Count()  if count >= b.capacity {    // The bucket is full.    return 0  }
  if !now().Before(b.p) {    // The bucket needs to be reset.    b.p = now()  }  remaining := b.capacity - count  if amount > remaining {    amount = remaining  }  t := time.Duration(float64(time.Second) * (float64(amount) / b.rate))  b.p = b.p.Add(t)
  return amount}

func New(maxCount int64, d time.Duration) *leakyBucket {  return (&leakyBucket{    int64(maxCount),    make(chan struct{}),    int64(d) / int64(maxCount),    time.Now().UnixNano(),    0,  }).start(d)}
type leakyBucket struct {  Maximal     int64  exitCh      chan struct{}  rate        int64  refreshTime int64 // in nanoseconds  count       int64}
func (s *leakyBucket) Count() int64            { return atomic.LoadInt64(&s.count) }func (s *leakyBucket) Available() int64        { return int64(s.count) }func (s *leakyBucket) Capacity() int64         { return int64(s.Maximal) }
func (s *leakyBucket) Close() {  close(s.exitCh)}
func (s *leakyBucket) start(d time.Duration) *leakyBucket {  if s.rate < 1000 {    log.Errorf("the rate cannot be less than 1000us, it's %v", s.rate)    return nil  }
  // fmt.Printf("rate: %v\n", time.Duration(s.rate))
  // go s.looper(d)  return s}
func (s *leakyBucket) looper(d time.Duration) {  // nothing to do}
func (s *leakyBucket) max(a, b int64) int64 {  if a < b {    return b  }  return a}
func (s *leakyBucket) take(count int) (requestAt time.Time, ok bool) {  requestAt = time.Now()
  s.count = s.max(0, s.count-(requestAt.UnixNano()-s.refreshTime)/s.rate*int64(count))  s.refreshTime = requestAt.UnixNano()
  if s.count < s.Maximal {    s.count += int64(count)    ok = true  }
  return}
func (s *leakyBucket) Take(count int) (ok bool) {  _, ok = s.take(count)  return}
func (s *leakyBucket) TakeBlocked(count int) (requestAt time.Time) {  var ok bool  requestAt, ok = s.take(count)  for !ok {    time.Sleep(time.Duration(s.rate - (1000 - 1)))    _, ok = s.take(count)  }  time.Sleep(time.Duration(s.rate-int64(time.Now().Sub(requestAt))) - time.Millisecond)  return}

import "github.com/hedzr/rate/leakybucket"
func TestLeakyBucketLimiter(b *testing.T) {  var counter int  l := leakybucket.New(100, time.Second, false) // one req per 10ms  defer l.Close()  time.Sleep(300 * time.Millisecond)  prev := time.Now()  for i := 0; i < 20; i++ {    ok := l.TakeBlocked(1)    now := time.Now()    counter++    fmt.Println(i, now.Sub(prev))    prev = now    time.Sleep(1 * time.Millisecond)  }  b.Logf("%v requests allowed.", counter)}

$ go test -v -race -test.run='^TestLeakyBucketLimiter$' ./leakybucket==== RUN   TestLeakyBucketLimiter0 9.468477ms1 10.47512ms2 12.589327ms3 10.364861ms4 11.820381ms5 10.353834ms6 12.504309ms7 10.167151ms8 11.623466ms9 10.487361ms10 10.72498ms11 12.435839ms12 11.73813ms13 10.956558ms14 11.893156ms15 11.964504ms16 11.856545ms17 11.094344ms18 10.978074ms19 10.632583ms    lb_test.go:36: 20 requests allowed.--- PASS: TestLeakyBucketLimiter (0.53s)PASSok      github.com/hedzr/rate/leakybucket       0.957s

func (r *Limiter) getRL(ctx gin.Context) rateapi.Limiter {  if r.limiter == nil {    r.limiter = leakybucket.New(100, time.Second, false)  }  return r.limiter}
func (r *Limiter) SimpleMiddleware() gin.HandlerFunc {  return func(ctx *gin.Context) {    limiter := r.getRL(ctx)    if err != nil {      ctx.AbortWithError(429, err)    } else if limiter != nil {      if limiter.Take(1) {        ctx.Writer.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", limiter.Available()))        ctx.Writer.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", limiter.Capacity()))        ctx.Next()      } else {        err = errors.New("Too many requests")        ctx.Writer.Header().Set("X-RateLimit-Remaining", fmt.Sprintf("%d", limiter.Available()))        ctx.Writer.Header().Set("X-RateLimit-Limit", fmt.Sprintf("%d", limiter.Capacity()))        ctx.AbortWithError(429, err)      }    } else {      ctx.Next()    }  }}

import "github.com/hedzr/rate/middleware"
func engine() *gin.Engine {  r := gin.Default()  r.Use(middleware.NewLimiterForGin(time.Second, 100, func(ctx *gin.Context) (string, error){    key := ctx.Request.Header.Get("X-API-KEY")    if key != "" {      return key, nil    }    ctx.JSON(403, gin.H{"code": 2901, "message": "API key is missing"})    return "", errors.New("API key is missing")  }))  return r}

func New(maxCount int64, d time.Duration) *tokenBucket {  return (&tokenBucket{    int32(maxCount),    d,    int64(d) / int64(maxCount),    make(chan struct{}),    int32(maxCount),  }).start(d)}
type tokenBucket struct {  Maximal int32  period  time.Duration  rate    int64  exitCh  chan struct{}  count   int32}
func (s *tokenBucket) Count() int32            { return atomic.LoadInt32(&s.count) }func (s *tokenBucket) Available() int64        { return int64(s.count) }func (s *tokenBucket) Capacity() int64         { return int64(s.Maximal) }
func (s *tokenBucket) Close() {  close(s.exitCh)}
func (s *tokenBucket) start(d time.Duration) *tokenBucket {  if s.rate < 1000 {    log.Errorf("the rate cannot be less than 1000us, it's %v", s.rate)    return nil  }
  go s.looper(d)  return s}
func (s *tokenBucket) looper(d time.Duration) {  ticker := time.NewTicker(d / time.Duration(s.Maximal))  // fmt.Printf("token building spped is: 1req/%v\n", d/time.Duration(s.Maximal))  defer func() {    ticker.Stop()  }()  for {    select {    case <-s.exitCh:      return    case <-ticker.C:      vn := atomic.AddInt32(&s.count, 1)      if vn < s.Maximal {        continue      }
      vn %= s.Maximal      if vn > 0 {        atomic.StoreInt32(&s.count, s.Maximal)      }    }  }}
func (s *tokenBucket) take(count int) bool {  if vn := atomic.AddInt32(&s.count, -1*int32(count)); vn >= 0 {    return true  }  atomic.StoreInt32(&s.count, 0)  return false}
func (s *tokenBucket) Take(count int) bool {  ok := s.take(count)  return ok}
func (s *tokenBucket) TakeBlocked(count int) (requestAt time.Time) {  requestAt = time.Now().UTC()  ok := s.take(count)  for !ok {    time.Sleep(time.Duration(s.rate - (1000 - 1)))    ok = s.take(count)  }  // time.Sleep(time.Duration(s.rate-int64(time.Now().Sub(requestAt))) - time.Millisecond)  return requestAt}

package main
import (  "fmt"  "github.com/hedzr/rate"  "time")
func main() {  // rate.LeakyBucket, rate.TokenBucket, ...  l := rate.New(rate.LeakyBucket, 100, time.Second)  for i := 0; i < 120; i++ {    ok := l.Take(1)    if !ok {      fmt.Printf("#%d Take() returns not ok, counter: %v\n", i, rate.CountOf(l))      time.Sleep(50 * time.Millisecond)    }  }}

import (  "github.com/hedzr/rate"  "github.com/hedzr/rate/middleware")
func engine() *gin.Engine {  config := &middleware.Config{    Name:          "A",    Description:   "A",    Algorithm:     stirng(rate.TokenBucket),    Interval:      time.Second,    MaxRequests:   1000,    HeaderKeyName: "X-API-TOKEN",    ExceptionKeys: nil,  }  r := gin.Default()  rg := r.Group("/api")  rg.Use(middleware.ForGin(config))  // ...    config2 := ...  r.Post("/login", middleware.ForGin(config2), loginHandler)  return r}