Golang 限速器

限速器

之前看到这篇golang并发编程的两种限速方法,觉得 sleep 等待的方式不是特别好,唤醒线程的时间比较长。而且1s内的请求只能均匀的到来,如瞬间来 N 个, 那么只有一个能立刻返回,剩下的只能等待。

【修正】
根据下图的说明,无论是 sleep 还是 chan block, 发生的事情都是 G 和 P 分离,等待下次轮训,再接着执行,所以可能性能是几乎一样的。(这种说法可能不准确,没有深入)
图片描述

限速器的作用还是比较重要的,特别是协程使用较多的应用,如果不加限制,可能会OOM。

简单说下思路:默认时间间隔定为1秒,简化问题。如果 调用 Limit() bool 接口,如果一秒内运行次数到达某个值,那么就阻塞, 直到下一个1秒重新计数。

代码

package lib
import (
 "sync/atomic"
 "time"
)
type RateLimiter struct {
 limit uint64
 count uint64
 ticker *time.Ticker
 lockCh chan struct{}
}
func NewRateLimiter(limit uint64) *RateLimiter {
 ticker := time.NewTicker(time.Second)
 r := &RateLimiter{
 limit: limit,
 count: 0,
 ticker: ticker,
 lockCh: make(chan struct{}),
 }
 go func() {
 for range ticker.C {
 if r.count > r.limit {
 select {
 case <-r.lockCh:
 default:
 r.resetCount()
 }
 }
 
 if r.count > 0 {
 r.resetCount() 
 }
 }
 }()
 return r
}
func (r *RateLimiter) Limit() bool {
 r.addCount(1)
 if r.getCount() > r.limit {
 var s struct{}
 r.lockCh <- s
 }
 return true
}
func (r *RateLimiter) addCount(interval uint64) {
 atomic.AddUint64(&r.count, interval)
}
func (r *RateLimiter) getCount() uint64 {
 return atomic.LoadUint64(&r.count)
}
func (r *RateLimiter) resetCount() {
 atomic.StoreUint64(&r.count, 1)
}

测试

package lib
import (
 "log"
 "testing"
 "time"
)
func TestLimit(t *testing.T) {
 limiter := NewRateLimiter(3)
 start := time.Now()
 for i := 0; i < 30; i++ {
 if limiter.Limit() {
 log.Printf("i is %d \n", i)
 }
 }
 end := time.Now()
 d := end.Sub(start)
 log.Println("spends seconds: ", d.Seconds())
}

比较

对比了开篇的sleep实现,limiter设为每秒3次,循环30次,用时分别为:

2016年08月31日 14:38:26 my limiter spends seconds: 9.00220346s
2016年08月31日 14:38:35 sleep limiter spends seconds: 9.762009934s

这里sleep limiter多sleep了两次,大约 0.67s, 减去这个值等于9.092009934。