限流-令牌桶实现(go版本)

703 阅读2分钟

漏斗算法思想是将所有请求先存到一个桶里。若此刻桶容量没满,表示当前请求是可以访问资源。若满了,则拒绝服务。同时桶会以固定速率取出桶里的请求来处理

具体实现方法可以将请求先暂存到一个队列中,若队列已满,则拒绝该请求。同时有一个周期性定时任务来消费队列里的数据

从实现可以看出,不管请求有多少或者瞬时流量有多大,请求的处理是固定速率的,所以令牌桶油流量整形的功能

相对计数器方法,令牌桶能有效避免抖动的问题,但当瞬时请求量很大时,后续的请求很有可能由于得不到及时处理而超时

import (
   "context"
   "errors"
   "github.com/google/uuid"
   "sync"
   "time"
)

type Request interface{}
type Command struct {
   id      string
   request Request
}

type Micros int64
type Handler func(request Request) (error, interface{})

type Flow interface {
   Do(ctx context.Context, request Request, timeout time.Duration) (error, interface{})
}
type CommandResult struct {
   err    error
   result interface{}
}

func NewCommandResult(err error, resp interface{}) *CommandResult {
   return &CommandResult{
      err:    err,
      result: resp,
   }
}

type TokenBucket struct {
   commandInterval Micros
   commands        chan *Command
   maxBucket       int
   ticker          *time.Ticker
   l               sync.Mutex
   handler         Handler
   dones           map[string]chan *CommandResult
}

func NewCommand(request Request) *Command {
   return &Command{
      id:      GetUUID(),
      request: request,
   }
}

func GetUUID() string {
   uid, err := uuid.NewUUID()
   if err != nil {
      panic(errors.New("uuid get fail"))
   }
   return uid.String()
}

func (token *TokenBucket) register(id string) chan *CommandResult {
   token.l.Lock()
   defer token.l.Unlock()
   done := token.dones[id]
   if done == nil {
      done = make(chan *CommandResult, 1)
      token.dones[id] = done
   }
   return done
}

func (token *TokenBucket) isRegister(id string) bool {
   token.l.Lock()
   defer token.l.Unlock()
   _, ok := token.dones[id]
   return ok
}

func (token *TokenBucket) trigger(id string, v *CommandResult) {
   token.l.Lock()
   ch := token.dones[id]
   delete(token.dones, id)
   token.l.Unlock()
   if ch != nil {
      ch <- v
      close(ch)
   }
}

func (token *TokenBucket) Do(ctx context.Context, request Request, timeout time.Duration) (error, interface{}) {
   command := NewCommand(request)
   done := token.register(command.id)
   select {
   case token.commands <- command:
   default:
      println("触发限流,time:" + time.Now().String() + "," + request.(string))
      return errors.New("触发限流"), nil
   }

   cctx, cancel := context.WithTimeout(ctx, timeout)
   defer cancel()

   select {
   case resp := <-done:
      return resp.err, resp.result
   case <-cctx.Done():
      println("请求超时,time:" + time.Now().String() + "," + request.(string))
      token.trigger(command.id, nil)
      return errors.New("请求超时,time:" + time.Now().String()), nil
   }
}

func (token *TokenBucket) start() {
   for range token.ticker.C {
      select {
      case command := <-token.commands:
         if token.isRegister(command.id) {
            err, resp := token.handler(command.request)
            token.trigger(command.id, NewCommandResult(err, resp))
         }
      default:
      }
   }
}

func NewTokenBucket(maxBuctet int /*桶里的最大请求数*/, commandInterval Micros /*请求处理的时间间隔,单位是微秒*/, handler Handler/*请求的处理方法*/) Flow {
   tokenBucket := &TokenBucket{
      commandInterval: commandInterval,
      commands:        make(chan *Command, maxBuctet),
      maxBucket:       maxBuctet,
      handler:         handler,
      dones:           make(map[string]chan *CommandResult),
      ticker:          time.NewTicker(time.Microsecond * time.Duration(commandInterval)),
   }
   go tokenBucket.start()
   return tokenBucket
}

具体使用方法:

token := NewTokenBucket(5, 100000, func(request Request) (e error, i interface{}) {
   println("handler:" + request.(string) + "," + time.Now().String())
   return nil, request
})
token.Do(context.Background(), "test"+strconv.Itoa(t), time.Millisecond*500)