Golang 简单实现并发限流

关于并发限流的实现有方法比较的多, 较多的是使用令牌桶算法。在Golang中可以利用 channel 的缓冲设定来实现并发限流的功能。只需要在收到请求后往 channel 里写入个数据（随便写点什么，内容不是很重要），然后在执行完成后把这个 channel 里的东西给取走，一个channel可以理解为一个先进先出的消息队列, 整个并发的数量可以根据channel队列的大小来判断。

以下主要代码说明:

type ConnLimiter struct {
    concurrentConn int //定义一个并发限制数
    bucket chan byte
}

//创建一个ConnLimiter
func NewConnLimiter(limit int) *ConnLimiter {
    return &ConnLimiter{
        concurrentConn: limit,
        bucket: make(chan byte, limit),
    }
}

//根据队列大小判断，是否达到并发限制
func (cl *ConnLimiter) GetConn() bool {
    if len(cl.bucket) >= cl.concurrentConn {
        log.Println("reached rate limit.")
        return false
    }
    cl.bucket <- 1
    return true
}

//请求结束后, 从chan中取走一个数据
func (cl *ConnLimiter) ReleaseConn() {
    <-cl.bucket
    log.Printf("release")
}

调用也比较简单, 判断可以放到middleware里。为了方便测试可以在代码中加入time.Sleep：

func rateLimitMiddleware(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if !connLimiter.GetConn() {
            w.WriteHeader(500)
            fmt.Fprintf(w, "too many requests!")
            return
        }
        defer connLimiter.ReleaseConn()
        next.ServeHTTP(w, r)
    })
}

示例代码: golang-ratelimit

Golang 简单实现并发限流

Golang Simple Ratelimit

CATALOG

FEATURED TAGS