Go如何实现HTTP请求限流示例

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

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

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

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

好了我们先看Demo代码的实现：

limit.go

package main

import (
  "net/http"

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

var limiter = rate.NewLimiter(2, 5)

func limit(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    if limiter.Allow() == false {
      http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
      return
    }

    next.ServeHTTP(w, r)
  })
}

main.go

package main

import (
  "net/http"
)

func main() {
  mux := http.NewServeMux()
  mux.HandleFunc("/", okHandler)

  // Wrap the servemux with the limit middleware.
  http.ListenAndServe(":4000", limit(mux))
}

func okHandler(w http.ResponseWriter, r *http.Request) {
  w.Write([]byte("OK"))
}

我们看看 rate.NewLimiter的源码：

// Copyright 2015 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package rate provides a rate limiter.
package rate

import (
 "fmt"
 "math"
 "sync"
 "time"

 "golang.org/x/net/context"
)

// Limit defines the maximum frequency of some events.
// Limit is represented as number of events per second.
// A zero Limit allows no events.
type Limit float64

// Inf is the infinite rate limit; it allows all events (even if burst is zero).
const Inf = Limit(math.MaxFloat64)

// Every converts a minimum time interval between events to a Limit.
func Every(interval time.Duration) Limit {
 if interval   burst {
  tokens = burst
 }
 // update state
 r.lim.last = now
 r.lim.tokens = tokens
 if r.timeToAct == r.lim.lastEvent {
  prevEvent := r.timeToAct.Add(r.limit.durationFromTokens(float64(-r.tokens)))
  if !prevEvent.Before(now) {
   r.lim.lastEvent = prevEvent
  }
 }

 return
}

// Reserve is shorthand for ReserveN(time.Now(), 1).
func (lim *Limiter) Reserve() *Reservation {
 return lim.ReserveN(time.Now(), 1)
}

// ReserveN returns a Reservation that indicates how long the caller must wait before n events happen.
// The Limiter takes this Reservation into account when allowing future events.
// ReserveN returns false if n exceeds the Limiter's burst size.
// Usage example:
//  r, ok := lim.ReserveN(time.Now(), 1)
//  if !ok {
//   // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
//  }
//  time.Sleep(r.Delay())
//  Act()
// Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
// If you need to respect a deadline or cancel the delay, use Wait instead.
// To drop or skip events exceeding rate limit, use Allow instead.
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation {
 r := lim.reserveN(now, n, InfDuration)
 return &r
}

// Wait is shorthand for WaitN(ctx, 1).
func (lim *Limiter) Wait(ctx context.Context) (err error) {
 return lim.WaitN(ctx, 1)
}

// WaitN blocks until lim permits n events to happen.
// It returns an error if n exceeds the Limiter's burst size, the Context is
// canceled, or the expected wait time exceeds the Context's Deadline.
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
 if n > lim.burst {
  return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
 }
 // Check if ctx is already cancelled
 select {
 case   maxElapsed {
  elapsed = maxElapsed
 }

 // Calculate the new number of tokens, due to time that passed.
 delta := lim.limit.tokensFromDuration(elapsed)
 tokens := lim.tokens + delta
 if burst := float64(lim.burst); tokens > burst {
  tokens = burst
 }

 return now, last, tokens
}

// durationFromTokens is a unit conversion function from the number of tokens to the duration
// of time it takes to accumulate them at a rate of limit tokens per second.
func (limit Limit) durationFromTokens(tokens float64) time.Duration {
 seconds := tokens / float64(limit)
 return time.Nanosecond * time.Duration(1e9*seconds)
}

// tokensFromDuration is a unit conversion function from a time duration to the number of tokens
// which could be accumulated during that duration at a rate of limit tokens per second.
func (limit Limit) tokensFromDuration(d time.Duration) float64 {
 return d.Seconds() * float64(limit)
}

算法描述：

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

实现用户粒度的限流

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

package main
import (
  "net/http"
  "sync"
  "time"

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

// Create a custom visitor struct which holds the rate limiter for each
// visitor and the last time that the visitor was seen.
type visitor struct {
  limiter *rate.Limiter
  lastSeen time.Time
}

// Change the the map to hold values of the type visitor.
var visitors = make(map[string]*visitor)
var mtx sync.Mutex
// Run a background goroutine to remove old entries from the visitors map.
func init() {
  go cleanupVisitors()
}

func addVisitor(ip string) *rate.Limiter {
  limiter := rate.NewLimiter(2, 5)
  mtx.Lock()
  // Include the current time when creating a new visitor.
  visitors[ip] = &visitor{limiter, time.Now()}
  mtx.Unlock()
  return limiter
}

func getVisitor(ip string) *rate.Limiter {
  mtx.Lock()
  v, exists := visitors[ip]
  if !exists {
    mtx.Unlock()
    return addVisitor(ip)
  }
  // Update the last seen time for the visitor.
  v.lastSeen = time.Now()
  mtx.Unlock()
  return v.limiter
}

// Every minute check the map for visitors that haven't been seen for
// more than 3 minutes and delete the entries.
func cleanupVisitors() {
  for {
    time.Sleep(time.Minute)
    mtx.Lock()
    for ip, v := range visitors {
      if time.Now().Sub(v.lastSeen) > 3*time.Minute {
        delete(visitors, ip)
      }
    }
    mtx.Unlock()
  }
}

func limit(next http.Handler) http.Handler {
  return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    limiter := getVisitor(r.RemoteAddr)
    if limiter.Allow() == false {
      http.Error(w, http.StatusText(429), http.StatusTooManyRequests)
      return
    }
    next.ServeHTTP(w, r)
  })
}

当然这只是一个简单的实现方案，如果我们要在微服务的API-GateWay中去实现限流还是要考虑很多东西的。建议大家可以看看 https://github.com/didip/tollbooth 的源码。

以上就是本文的全部内容，希望对大家的学习有所帮助，也希望大家多多支持。

查看更多关于Go如何实现HTTP请求限流示例的详细内容...