zoukankan      html  css  js  c++  java
  • Golang实现请求限流的几种办法

    目录

    简单的并发控制

    使用计数器实现请求限流

    使用golang官方包实现httpserver频率限制

    使用Token Bucket(令牌桶算法)实现请求限流


    简单的并发控制

    利用 channel 的缓冲设定,我们就可以来实现并发的限制。我们只要在执行并发的同时,往一个带有缓冲的 channel 里写入点东西(随便写啥,内容不重要)。让并发的 goroutine在执行完成后把这个 channel 里的东西给读走。这样整个并发的数量就讲控制在这个 channel的缓冲区大小上。

    比如我们可以用一个 bool 类型的带缓冲 channel 作为并发限制的计数器。

    chLimit := make(chan bool, 1)

     然后在并发执行的地方,每创建一个新的 goroutine,都往 chLimit 里塞个东西。

    1.  
      for i, sleeptime := range input {
    2.  
      chs[i] = make(chan string, 1)
    3.  
      chLimit <- true
    4.  
      go limitFunc(chLimit, chs[i], i, sleeptime, timeout)
    5.  
      }

    这里通过 go 关键字并发执行的是新构造的函数。他在执行完后,会把 chLimit的缓冲区里给消费掉一个。

    1.  
      limitFunc := func(chLimit chan bool, ch chan string, task_id, sleeptime, timeout int) {
    2.  
      Run(task_id, sleeptime, timeout, ch)
    3.  
      <-chLimit
    4.  
      }

    这样一来,当创建的 goroutine 数量到达 chLimit 的缓冲区上限后。主 goroutine 就挂起阻塞了,直到这些 goroutine 执行完毕,消费掉了 chLimit 缓冲区中的数据,程序才会继续创建新的 goroutine 。我们并发数量限制的目的也就达到了。

    以下是完整代码:

    1.  
      package main
    2.  
       
    3.  
      import (
    4.  
      "fmt"
    5.  
      "time"
    6.  
      )
    7.  
       
    8.  
      func Run(task_id, sleeptime, timeout int, ch chan string) {
    9.  
      ch_run := make(chan string)
    10.  
      go run(task_id, sleeptime, ch_run)
    11.  
      select {
    12.  
      case re := <-ch_run:
    13.  
      ch <- re
    14.  
      case <-time.After(time.Duration(timeout) * time.Second):
    15.  
      re := fmt.Sprintf("task id %d , timeout", task_id)
    16.  
      ch <- re
    17.  
      }
    18.  
      }
    19.  
       
    20.  
      func run(task_id, sleeptime int, ch chan string) {
    21.  
       
    22.  
      time.Sleep(time.Duration(sleeptime) * time.Second)
    23.  
      ch <- fmt.Sprintf("task id %d , sleep %d second", task_id, sleeptime)
    24.  
      return
    25.  
      }
    26.  
       
    27.  
      func main() {
    28.  
      input := []int{3, 2, 1}
    29.  
      timeout := 2
    30.  
      chLimit := make(chan bool, 1)
    31.  
      chs := make([]chan string, len(input))
    32.  
      limitFunc := func(chLimit chan bool, ch chan string, task_id, sleeptime, timeout int) {
    33.  
      Run(task_id, sleeptime, timeout, ch)
    34.  
      <-chLimit
    35.  
      }
    36.  
      startTime := time.Now()
    37.  
      fmt.Println("Multirun start")
    38.  
      for i, sleeptime := range input {
    39.  
      chs[i] = make(chan string, 1)
    40.  
      chLimit <- true
    41.  
      go limitFunc(chLimit, chs[i], i, sleeptime, timeout)
    42.  
      }
    43.  
       
    44.  
      for _, ch := range chs {
    45.  
      fmt.Println(<-ch)
    46.  
      }
    47.  
      endTime := time.Now()
    48.  
      fmt.Printf("Multissh finished. Process time %s. Number of task is %d", endTime.Sub(startTime), len(input))
    49.  
      }

    运行结果:

    1.  
      Multirun start
    2.  
      task id 0 , timeout
    3.  
      task id 1 , timeout
    4.  
      task id 2 , sleep 1 second
    5.  
      Multissh finished. Process time 5s. Number of task is 3

    如果修改并发限制为2:

    chLimit := make(chan bool, 2)

    运行结果:

    1.  
       
    2.  
      Multirun start
    3.  
      task id 0 , timeout
    4.  
      task id 1 , timeout
    5.  
      task id 2 , sleep 1 second
    6.  
      Multissh finished. Process time 3s. Number of task is 3

    使用计数器实现请求限流

    限流的要求是在指定的时间间隔内,server 最多只能服务指定数量的请求。实现的原理是我们启动一个计数器,每次服务请求会把计数器加一,同时到达指定的时间间隔后会把计数器清零;这个计数器的实现代码如下所示:

    1.  
      type RequestLimitService struct {
    2.  
      Interval time.Duration
    3.  
      MaxCount int
    4.  
      Lock sync.Mutex
    5.  
      ReqCount int
    6.  
      }
    7.  
       
    8.  
      func NewRequestLimitService(interval time.Duration, maxCnt int) *RequestLimitService {
    9.  
      reqLimit := &RequestLimitService{
    10.  
      Interval: interval,
    11.  
      MaxCount: maxCnt,
    12.  
      }
    13.  
       
    14.  
      go func() {
    15.  
      ticker := time.NewTicker(interval)
    16.  
      for {
    17.  
      <-ticker.C
    18.  
      reqLimit.Lock.Lock()
    19.  
      fmt.Println("Reset Count...")
    20.  
      reqLimit.ReqCount = 0
    21.  
      reqLimit.Lock.Unlock()
    22.  
      }
    23.  
      }()
    24.  
       
    25.  
      return reqLimit
    26.  
      }
    27.  
       
    28.  
      func (reqLimit *RequestLimitService) Increase() {
    29.  
      reqLimit.Lock.Lock()
    30.  
      defer reqLimit.Lock.Unlock()
    31.  
       
    32.  
      reqLimit.ReqCount += 1
    33.  
      }
    34.  
       
    35.  
      func (reqLimit *RequestLimitService) IsAvailable() bool {
    36.  
      reqLimit.Lock.Lock()
    37.  
      defer reqLimit.Lock.Unlock()
    38.  
       
    39.  
      return reqLimit.ReqCount < reqLimit.MaxCount
    40.  
      }

    在服务请求的时候, 我们会对当前计数器和阈值进行比较,只有未超过阈值时才进行服务:

    1.  
      var RequestLimit = NewRequestLimitService(10 * time.Second, 5)
    2.  
       
    3.  
      func helloHandler(w http.ResponseWriter, r *http.Request) {
    4.  
      if RequestLimit.IsAvailable() {
    5.  
      RequestLimit.Increase()
    6.  
      fmt.Println(RequestLimit.ReqCount)
    7.  
      io.WriteString(w, "Hello world! ")
    8.  
      } else {
    9.  
      fmt.Println("Reach request limiting!")
    10.  
      io.WriteString(w, "Reach request limit! ")
    11.  
      }
    12.  
      }
    13.  
       
    14.  
      func main() {
    15.  
      fmt.Println("Server Started!")
    16.  
      http.HandleFunc("/", helloHandler)
    17.  
      http.ListenAndServe(":8000", nil)
    18.  
      }

     完整代码url:https://github.com/hiberabyss/JustDoIt/blob/master/RequestLimit/request_limit.go

    使用golang官方包实现httpserver频率限制

    使用golang来编写httpserver时,可以使用官方已经有实现好的包:

    1.  
      import(
    2.  
      "fmt"
    3.  
      "net"
    4.  
      "golang.org/x/net/netutil"
    5.  
      )
    6.  
       
    7.  
      func main() {
    8.  
      l, err := net.Listen("tcp", "127.0.0.1:0")
    9.  
      if err != nil {
    10.  
      fmt.Fatalf("Listen: %v", err)
    11.  
      }
    12.  
      defer l.Close()
    13.  
      l = LimitListener(l, max)
    14.  
       
    15.  
      http.Serve(l, http.HandlerFunc())
    16.  
       
    17.  
      //bla bla bla.................
    18.  
      }

    源码如下(url : https://github.com/golang/net/blob/master/netutil/listen.go),基本思路就是为连接数计数,通过make chan来建立一个最大连接数的channel, 每次accept就+1,close时候就-1. 当到达最大连接数时,就等待空闲连接出来之后再accept。

    1.  
      // Copyright 2013 The Go Authors. All rights reserved.
    2.  
      // Use of this source code is governed by a BSD-style
    3.  
      // license that can be found in the LICENSE file.
    4.  
       
    5.  
      // Package netutil provides network utility functions, complementing the more
    6.  
      // common ones in the net package.
    7.  
      package netutil // import "golang.org/x/net/netutil"
    8.  
       
    9.  
      import (
    10.  
      "net"
    11.  
      "sync"
    12.  
      )
    13.  
       
    14.  
      // LimitListener returns a Listener that accepts at most n simultaneous
    15.  
      // connections from the provided Listener.
    16.  
      func LimitListener(l net.Listener, n int) net.Listener {
    17.  
      return &limitListener{
    18.  
      Listener: l,
    19.  
      sem: make(chan struct{}, n),
    20.  
      done: make(chan struct{}),
    21.  
      }
    22.  
      }
    23.  
       
    24.  
      type limitListener struct {
    25.  
      net.Listener
    26.  
      sem chan struct{}
    27.  
      closeOnce sync.Once // ensures the done chan is only closed once
    28.  
      done chan struct{} // no values sent; closed when Close is called
    29.  
      }
    30.  
       
    31.  
      // acquire acquires the limiting semaphore. Returns true if successfully
    32.  
      // accquired, false if the listener is closed and the semaphore is not
    33.  
      // acquired.
    34.  
      func (l *limitListener) acquire() bool {
    35.  
      select {
    36.  
      case <-l.done:
    37.  
      return false
    38.  
      case l.sem <- struct{}{}:
    39.  
      return true
    40.  
      }
    41.  
      }
    42.  
      func (l *limitListener) release() { <-l.sem }
    43.  
       
    44.  
      func (l *limitListener) Accept() (net.Conn, error) {
    45.  
      //如果sem满了,就会阻塞在这
    46.  
      acquired := l.acquire()
    47.  
      // If the semaphore isn't acquired because the listener was closed, expect
    48.  
      // that this call to accept won't block, but immediately return an error.
    49.  
      c, err := l.Listener.Accept()
    50.  
      if err != nil {
    51.  
      if acquired {
    52.  
      l.release()
    53.  
      }
    54.  
      return nil, err
    55.  
      }
    56.  
      return &limitListenerConn{Conn: c, release: l.release}, nil
    57.  
      }
    58.  
       
    59.  
      func (l *limitListener) Close() error {
    60.  
      err := l.Listener.Close()
    61.  
      l.closeOnce.Do(func() { close(l.done) })
    62.  
      return err
    63.  
      }
    64.  
       
    65.  
      type limitListenerConn struct {
    66.  
      net.Conn
    67.  
      releaseOnce sync.Once
    68.  
      release func()
    69.  
      }
    70.  
       
    71.  
      func (l *limitListenerConn) Close() error {
    72.  
      err := l.Conn.Close()
    73.  
      //close时释放占用的sem
    74.  
      l.releaseOnce.Do(l.release)
    75.  
      return err
    76.  
      }

    使用Token Bucket(令牌桶算法)实现请求限流

    在开发高并发系统时有三把利器用来保护系统:缓存、降级和限流!为了保证在业务高峰期,线上系统也能保证一定的弹性和稳定性,最有效的方案就是进行服务降级了,而限流就是降级系统最常采用的方案之一。

    这里为大家推荐一个开源库https://github.com/didip/tollbooth,但是,如果您想要一些简单的、轻量级的或者只是想要学习的东西,实现自己的中间件来处理速率限制并不困难。今天我们就来聊聊如何实现自己的一个限流中间件

    首先我们需要安装一个提供了 Token bucket (令牌桶算法)的依赖包,上面提到的toolbooth 的实现也是基于它实现的:

    $ go get golang.org/x/time/rate

    先看Demo代码的实现:

    1.  
      package main
    2.  
       
    3.  
      import (
    4.  
      "net/http"
    5.  
      "golang.org/x/time/rate"
    6.  
      )
    7.  
       
    8.  
      var limiter = rate.NewLimiter(2, 5)
    9.  
      func limit(next http.Handler) http.Handler {
    10.  
      return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    11.  
      if limiter.Allow() == false {
    12.  
      http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
    13.  
      return
    14.  
      }
    15.  
      next.ServeHTTP(w, r)
    16.  
      })
    17.  
      }
    18.  
       
    19.  
      func main() {
    20.  
      mux := http.NewServeMux()
    21.  
      mux.HandleFunc("/", okHandler)
    22.  
      // Wrap the servemux with the limit middleware.
    23.  
      http.ListenAndServe(":4000", limit(mux))
    24.  
      }
    25.  
       
    26.  
      func okHandler(w http.ResponseWriter, r *http.Request) {
    27.  
      w.Write([]byte("OK"))
    28.  
      }

    然后看看 rate.NewLimiter的源码:

    算法描述:用户配置的平均发送速率为r,则每隔1/r秒一个令牌被加入到桶中(每秒会有r个令牌放入桶中),桶中最多可以存放b个令牌。如果令牌到达时令牌桶已经满了,那么这个令牌会被丢弃;

    1.  
      // Copyright 2015 The Go Authors. All rights reserved.
    2.  
      // Use of this source code is governed by a BSD-style
    3.  
      // license that can be found in the LICENSE file.
    4.  
      // Package rate provides a rate limiter.
    5.  
      package rate
    6.  
       
    7.  
      import (
    8.  
      "fmt"
    9.  
      "math"
    10.  
      "sync"
    11.  
      "time"
    12.  
       
    13.  
      "golang.org/x/net/context"
    14.  
      )
    15.  
       
    16.  
      // Limit defines the maximum frequency of some events.
    17.  
      // Limit is represented as number of events per second.
    18.  
      // A zero Limit allows no events.
    19.  
      type Limit float64
    20.  
       
    21.  
      // Inf is the infinite rate limit; it allows all events (even if burst is zero).
    22.  
      const Inf = Limit(math.MaxFloat64)
    23.  
       
    24.  
      // Every converts a minimum time interval between events to a Limit.
    25.  
      func Every(interval time.Duration) Limit {
    26.  
      if interval <= 0 {
    27.  
      return Inf
    28.  
      }
    29.  
      return 1 / Limit(interval.Seconds())
    30.  
      }
    31.  
       
    32.  
      // A Limiter controls how frequently events are allowed to happen.
    33.  
      // It implements a "token bucket" of size b, initially full and refilled
    34.  
      // at rate r tokens per second.
    35.  
      // Informally, in any large enough time interval, the Limiter limits the
    36.  
      // rate to r tokens per second, with a maximum burst size of b events.
    37.  
      // As a special case, if r == Inf (the infinite rate), b is ignored.
    38.  
      // See https://en.wikipedia.org/wiki/Token_bucket for more about token buckets.
    39.  
      //
    40.  
      // The zero value is a valid Limiter, but it will reject all events.
    41.  
      // Use NewLimiter to create non-zero Limiters.
    42.  
      //
    43.  
      // Limiter has three main methods, Allow, Reserve, and Wait.
    44.  
      // Most callers should use Wait.
    45.  
      //
    46.  
      // Each of the three methods consumes a single token.
    47.  
      // They differ in their behavior when no token is available.
    48.  
      // If no token is available, Allow returns false.
    49.  
      // If no token is available, Reserve returns a reservation for a future token
    50.  
      // and the amount of time the caller must wait before using it.
    51.  
      // If no token is available, Wait blocks until one can be obtained
    52.  
      // or its associated context.Context is canceled.
    53.  
      //
    54.  
      // The methods AllowN, ReserveN, and WaitN consume n tokens.
    55.  
      type Limiter struct {
    56.  
      //maximum token, token num per second
    57.  
      limit Limit
    58.  
      //burst field, max token num
    59.  
      burst int
    60.  
      mu sync.Mutex
    61.  
      //tokens num, change
    62.  
      tokens float64
    63.  
      // last is the last time the limiter's tokens field was updated
    64.  
      last time.Time
    65.  
      // lastEvent is the latest time of a rate-limited event (past or future)
    66.  
      lastEvent time.Time
    67.  
      }
    68.  
       
    69.  
      // Limit returns the maximum overall event rate.
    70.  
      func (lim *Limiter) Limit() Limit {
    71.  
      lim.mu.Lock()
    72.  
      defer lim.mu.Unlock()
    73.  
      return lim.limit
    74.  
      }
    75.  
       
    76.  
      // Burst returns the maximum burst size. Burst is the maximum number of tokens
    77.  
      // that can be consumed in a single call to Allow, Reserve, or Wait, so higher
    78.  
      // Burst values allow more events to happen at once.
    79.  
      // A zero Burst allows no events, unless limit == Inf.
    80.  
      func (lim *Limiter) Burst() int {
    81.  
      return lim.burst
    82.  
      }
    83.  
       
    84.  
      // NewLimiter returns a new Limiter that allows events up to rate r and permits
    85.  
      // bursts of at most b tokens.
    86.  
      func NewLimiter(r Limit, b int) *Limiter {
    87.  
      return &Limiter{
    88.  
      limit: r,
    89.  
      burst: b,
    90.  
      }
    91.  
      }
    92.  
       
    93.  
      // Allow is shorthand for AllowN(time.Now(), 1).
    94.  
      func (lim *Limiter) Allow() bool {
    95.  
      return lim.AllowN(time.Now(), 1)
    96.  
      }
    97.  
       
    98.  
      // AllowN reports whether n events may happen at time now.
    99.  
      // Use this method if you intend to drop / skip events that exceed the rate limit.
    100.  
      // Otherwise use Reserve or Wait.
    101.  
      func (lim *Limiter) AllowN(now time.Time, n int) bool {
    102.  
      return lim.reserveN(now, n, 0).ok
    103.  
      }
    104.  
       
    105.  
      // A Reservation holds information about events that are permitted by a Limiter to happen after a delay.
    106.  
      // A Reservation may be canceled, which may enable the Limiter to permit additional events.
    107.  
      type Reservation struct {
    108.  
      ok bool
    109.  
      lim *Limiter
    110.  
      tokens int
    111.  
      //This is the time to action
    112.  
      timeToAct time.Time
    113.  
      // This is the Limit at reservation time, it can change later.
    114.  
      limit Limit
    115.  
      }
    116.  
       
    117.  
      // OK returns whether the limiter can provide the requested number of tokens
    118.  
      // within the maximum wait time. If OK is false, Delay returns InfDuration, and
    119.  
      // Cancel does nothing.
    120.  
      func (r *Reservation) OK() bool {
    121.  
      return r.ok
    122.  
      }
    123.  
       
    124.  
      // Delay is shorthand for DelayFrom(time.Now()).
    125.  
      func (r *Reservation) Delay() time.Duration {
    126.  
      return r.DelayFrom(time.Now())
    127.  
      }
    128.  
       
    129.  
      // InfDuration is the duration returned by Delay when a Reservation is not OK.
    130.  
      const InfDuration = time.Duration(1<<63 - 1)
    131.  
       
    132.  
      // DelayFrom returns the duration for which the reservation holder must wait
    133.  
      // before taking the reserved action. Zero duration means act immediately.
    134.  
      // InfDuration means the limiter cannot grant the tokens requested in this
    135.  
      // Reservation within the maximum wait time.
    136.  
      func (r *Reservation) DelayFrom(now time.Time) time.Duration {
    137.  
      if !r.ok {
    138.  
      return InfDuration
    139.  
      }
    140.  
      delay := r.timeToAct.Sub(now)
    141.  
      if delay < 0 {
    142.  
      return 0
    143.  
      }
    144.  
      return delay
    145.  
      }
    146.  
       
    147.  
      // Cancel is shorthand for CancelAt(time.Now()).
    148.  
      func (r *Reservation) Cancel() {
    149.  
      r.CancelAt(time.Now())
    150.  
      return
    151.  
      }
    152.  
       
    153.  
      // CancelAt indicates that the reservation holder will not perform the reserved action
    154.  
      // and reverses the effects of this Reservation on the rate limit as much as possible,
    155.  
      // considering that other reservations may have already been made.
    156.  
      func (r *Reservation) CancelAt(now time.Time) {
    157.  
      if !r.ok {
    158.  
      return
    159.  
      }
    160.  
      r.lim.mu.Lock()
    161.  
      defer r.lim.mu.Unlock()
    162.  
      if r.lim.limit == Inf || r.tokens == 0 || r.timeToAct.Before(now) {
    163.  
      return
    164.  
      }
    165.  
      // calculate tokens to restore
    166.  
      // The duration between lim.lastEvent and r.timeToAct tells us how many tokens were reserved
    167.  
      // after r was obtained. These tokens should not be restored.
    168.  
      restoreTokens := float64(r.tokens) - r.limit.tokensFromDuration(r.lim.lastEvent.Sub(r.timeToAct))
    169.  
      if restoreTokens <= 0 {
    170.  
      return
    171.  
      }
    172.  
      // advance time to now
    173.  
      now, _, tokens := r.lim.advance(now)
    174.  
      // calculate new number of tokens
    175.  
      tokens += restoreTokens
    176.  
      if burst := float64(r.lim.burst); tokens > burst {
    177.  
      tokens = burst
    178.  
      }
    179.  
      // update state
    180.  
      r.lim.last = now
    181.  
      r.lim.tokens = tokens
    182.  
      if r.timeToAct == r.lim.lastEvent {
    183.  
      prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
    184.  
      if !prevEvent.Before(now) {
    185.  
      r.lim.lastEvent = prevEvent
    186.  
      }
    187.  
      }
    188.  
      return
    189.  
      }
    190.  
       
    191.  
      // Reserve is shorthand for ReserveN(time.Now(), 1).
    192.  
      func (lim *Limiter) Reserve() *Reservation {
    193.  
      return lim.ReserveN(time.Now(), 1)
    194.  
      }
    195.  
       
    196.  
      // ReserveN returns a Reservation that indicates how long the caller must wait before n events happen.
    197.  
      // The Limiter takes this Reservation into account when allowing future events.
    198.  
      // ReserveN returns false if n exceeds the Limiter's burst size.
    199.  
      // Usage example:
    200.  
      // r, ok := lim.ReserveN(time.Now(), 1)
    201.  
      // if !ok {
    202.  
      // // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
    203.  
      // }
    204.  
      // time.Sleep(r.Delay())
    205.  
      // Act()
    206.  
      // Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
    207.  
      // If you need to respect a deadline or cancel the delay, use Wait instead.
    208.  
      // To drop or skip events exceeding rate limit, use Allow instead.
    209.  
      func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation {
    210.  
      r := lim.reserveN(now, n, InfDuration)
    211.  
      return &r
    212.  
      }
    213.  
       
    214.  
      // Wait is shorthand for WaitN(ctx, 1).
    215.  
      func (lim *Limiter) Wait(ctx context.Context) (err error) {
    216.  
      return lim.WaitN(ctx, 1)
    217.  
      }
    218.  
       
    219.  
      // WaitN blocks until lim permits n events to happen.
    220.  
      // It returns an error if n exceeds the Limiter's burst size, the Context is
    221.  
      // canceled, or the expected wait time exceeds the Context's Deadline.
    222.  
      func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
    223.  
      if n > lim.burst {
    224.  
      return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
    225.  
      }
    226.  
      // Check if ctx is already cancelled
    227.  
      select {
    228.  
      case <-ctx.Done():
    229.  
      return ctx.Err()
    230.  
      default:
    231.  
      }
    232.  
      // Determine wait limit
    233.  
      now := time.Now()
    234.  
      waitLimit := InfDuration
    235.  
      if deadline, ok := ctx.Deadline(); ok {
    236.  
      waitLimit = deadline.Sub(now)
    237.  
      }
    238.  
      // Reserve
    239.  
      r := lim.reserveN(now, n, waitLimit)
    240.  
      if !r.ok {
    241.  
      return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
    242.  
      }
    243.  
      // Wait
    244.  
      t := time.NewTimer(r.DelayFrom(now))
    245.  
      defer t.Stop()
    246.  
      select {
    247.  
      case <-t.C:
    248.  
      // We can proceed.
    249.  
      return nil
    250.  
      case <-ctx.Done():
    251.  
      // Context was canceled before we could proceed. Cancel the
    252.  
      // reservation, which may permit other events to proceed sooner.
    253.  
      r.Cancel()
    254.  
      return ctx.Err()
    255.  
      }
    256.  
      }
    257.  
       
    258.  
      // SetLimit is shorthand for SetLimitAt(time.Now(), newLimit).
    259.  
      func (lim *Limiter) SetLimit(newLimit Limit) {
    260.  
      lim.SetLimitAt(time.Now(), newLimit)
    261.  
      }
    262.  
       
    263.  
      // SetLimitAt sets a new Limit for the limiter. The new Limit, and Burst, may be violated
    264.  
      // or underutilized by those which reserved (using Reserve or Wait) but did not yet act
    265.  
      // before SetLimitAt was called.
    266.  
      func (lim *Limiter) SetLimitAt(now time.Time, newLimit Limit) {
    267.  
      lim.mu.Lock()
    268.  
      defer lim.mu.Unlock()
    269.  
      now, _, tokens := lim.advance(now)
    270.  
      lim.last = now
    271.  
      lim.tokens = tokens
    272.  
      lim.limit = newLimit
    273.  
      }
    274.  
       
    275.  
      // reserveN is a helper method for AllowN, ReserveN, and WaitN.
    276.  
      // maxFutureReserve specifies the maximum reservation wait duration allowed.
    277.  
      // reserveN returns Reservation, not *Reservation, to avoid allocation in AllowN and WaitN.
    278.  
      func (lim *Limiter) reserveN(now time.Time, n int, maxFutureReserve time.Duration) Reservation {
    279.  
      lim.mu.Lock()
    280.  
      defer lim.mu.Unlock()
    281.  
      if lim.limit == Inf {
    282.  
      return Reservation{
    283.  
      ok: true,
    284.  
      lim: lim,
    285.  
      tokens: n,
    286.  
      timeToAct: now,
    287.  
      }
    288.  
      }
    289.  
      now, last, tokens := lim.advance(now)
    290.  
      // Calculate the remaining number of tokens resulting from the request.
    291.  
      tokens -= float64(n)
    292.  
      // Calculate the wait duration
    293.  
      var waitDuration time.Duration
    294.  
      if tokens < 0 {
    295.  
      waitDuration = lim.limit.durationFromTokens(-tokens)
    296.  
      }
    297.  
      // Decide result
    298.  
      ok := n <= lim.burst && waitDuration <= maxFutureReserve
    299.  
      // Prepare reservation
    300.  
      r := Reservation{
    301.  
      ok: ok,
    302.  
      lim: lim,
    303.  
      limit: lim.limit,
    304.  
      }
    305.  
      if ok {
    306.  
      r.tokens = n
    307.  
      r.timeToAct = now.Add(waitDuration)
    308.  
      }
    309.  
      // Update state
    310.  
      if ok {
    311.  
      lim.last = now
    312.  
      lim.tokens = tokens
    313.  
      lim.lastEvent = r.timeToAct
    314.  
      } else {
    315.  
      lim.last = last
    316.  
      }
    317.  
      return r
    318.  
      }
    319.  
       
    320.  
      // advance calculates and returns an updated state for lim resulting from the passage of time.
    321.  
      // lim is not changed.
    322.  
      func (lim *Limiter) advance(now time.Time) (newNow time.Time, newLast time.Time, newTokens float64) {
    323.  
      last := lim.last
    324.  
      if now.Before(last) {
    325.  
      last = now
    326.  
      }
    327.  
      // Avoid making delta overflow below when last is very old.
    328.  
      maxElapsed := lim.limit.durationFromTokens(float64(lim.burst) - lim.tokens)
    329.  
      elapsed := now.Sub(last)
    330.  
      if elapsed > maxElapsed {
    331.  
      elapsed = maxElapsed
    332.  
      }
    333.  
      // Calculate the new number of tokens, due to time that passed.
    334.  
      delta := lim.limit.tokensFromDuration(elapsed)
    335.  
      tokens := lim.tokens + delta
    336.  
      if burst := float64(lim.burst); tokens > burst {
    337.  
      tokens = burst
    338.  
      }
    339.  
      return now, last, tokens
    340.  
      }
    341.  
       
    342.  
      // durationFromTokens is a unit conversion function from the number of tokens to the duration
    343.  
      // of time it takes to accumulate them at a rate of limit tokens per second.
    344.  
      func (limit Limit) durationFromTokens(tokens float64) time.Duration {
    345.  
      seconds := tokens / float64(limit)
    346.  
      return time.Nanosecond * time.Duration(1e9*seconds)
    347.  
      }
    348.  
       
    349.  
      // tokensFromDuration is a unit conversion function from a time duration to the number of tokens
    350.  
      // which could be accumulated during that duration at a rate of limit tokens per second.
    351.  
      func (limit Limit) tokensFromDuration(d time.Duration) float64 {
    352.  
      return d.Seconds() * float64(limit)
    353.  
      }

    虽然在某些情况下使用单个全局速率限制器非常有用,但另一种常见情况是基于IP地址或API密钥等标识符为每个用户实施速率限制器。我们将使用IP地址作为标识符。简单实现代码如下:

    1.  
      package main
    2.  
      import (
    3.  
      "net/http"
    4.  
      "sync"
    5.  
      "time"
    6.  
      "golang.org/x/time/rate"
    7.  
      )
    8.  
      // Create a custom visitor struct which holds the rate limiter for each
    9.  
      // visitor and the last time that the visitor was seen.
    10.  
      type visitor struct {
    11.  
      limiter *rate.Limiter
    12.  
      lastSeen time.Time
    13.  
      }
    14.  
      // Change the the map to hold values of the type visitor.
    15.  
      var visitors = make(map[string]*visitor)
    16.  
      var mtx sync.Mutex
    17.  
      // Run a background goroutine to remove old entries from the visitors map.
    18.  
      func init() {
    19.  
      go cleanupVisitors()
    20.  
      }
    21.  
      func addVisitor(ip string) *rate.Limiter {
    22.  
      limiter := rate.NewLimiter(2, 5)
    23.  
      mtx.Lock()
    24.  
      // Include the current time when creating a new visitor.
    25.  
      visitors[ip] = &visitor{limiter, time.Now()}
    26.  
      mtx.Unlock()
    27.  
      return limiter
    28.  
      }
    29.  
      func getVisitor(ip string) *rate.Limiter {
    30.  
      mtx.Lock()
    31.  
      v, exists := visitors[ip]
    32.  
      if !exists {
    33.  
      mtx.Unlock()
    34.  
      return addVisitor(ip)
    35.  
      }
    36.  
      // Update the last seen time for the visitor.
    37.  
      v.lastSeen = time.Now()
    38.  
      mtx.Unlock()
    39.  
      return v.limiter
    40.  
      }
    41.  
      // Every minute check the map for visitors that haven't been seen for
    42.  
      // more than 3 minutes and delete the entries.
    43.  
      func cleanupVisitors() {
    44.  
      for {
    45.  
      time.Sleep(time.Minute)
    46.  
      mtx.Lock()
    47.  
      for ip, v := range visitors {
    48.  
      if time.Now().Sub(v.lastSeen) > 3*time.Minute {
    49.  
      delete(visitors, ip)
    50.  
      }
    51.  
      }
    52.  
      mtx.Unlock()
    53.  
      }
    54.  
      }
    55.  
      func limit(next http.Handler) http.Handler {
    56.  
      return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    57.  
      limiter := getVisitor(r.RemoteAddr)
    58.  
      if limiter.Allow() == false {
    59.  
      http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
    60.  
      return
    61.  
      }
    62.  
      next.ServeHTTP(w, r)
    63.  
      })
    64.  
      }

    转载: https://blog.csdn.net/micl200110041/article/details/82013032

  • 相关阅读:
    AGC算法
    Cordic算法
    git Remote: HTTP Basic: Access denied Git failed with a fatal error.
    mysql 定义用户变量
    Docker 报错处理
    IIS,Docker 部署.Net Core
    SpringBoot向后台传参的若干种方式
    修改Mysql 数据库以及表字符集
    安装Docker
    获取北京时间
  • 原文地址:https://www.cnblogs.com/lovezbs/p/13910575.html
Copyright © 2011-2022 走看看