最近项目需要用到限流和统计流量的功能, 便用 go 基于计数器的原理简单地实现了这两个功能.
限流
限流的要求是在指定的时间间隔内, server 最多只能服务指定数量的请求. 实现的原理是我们启动一个计数器, 每次服务请求会把计数器加一, 同时到达指定的时间间隔后会把计数器清零; 这个计数器的实现代码如下所示:
type RequestLimitService struct {
Interval time.Duration
MaxCount int
Lock sync.Mutex
ReqCount int
}
func NewRequestLimitService(interval time.Duration, maxCnt int) *RequestLimitService {
reqLimit := &RequestLimitService{
Interval: interval,
MaxCount: maxCnt,
}
go func() {
ticker := time.NewTicker(interval)
for {
<-ticker.C
reqLimit.Lock.Lock()
fmt.Println("Reset Count...")
reqLimit.ReqCount = 0
reqLimit.Lock.Unlock()
}
}()
return reqLimit
}
func (reqLimit *RequestLimitService) Increase() {
reqLimit.Lock.Lock()
defer reqLimit.Lock.Unlock()
reqLimit.ReqCount += 1
}
func (reqLimit *RequestLimitService) IsAvailable() bool {
reqLimit.Lock.Lock()
defer reqLimit.Lock.Unlock()
return reqLimit.ReqCount < reqLimit.MaxCount
}
在服务请求的时候, 我们会对当前计数器和阈值进行比较, 只有未超过阈值时才进行服务:
var RequestLimit = NewRequestLimitService(10 * time.Second, 5)
func helloHandler(w http.ResponseWriter, r *http.Request) {
if RequestLimit.IsAvailable() {
RequestLimit.Increase()
fmt.Println(RequestLimit.ReqCount)
io.WriteString(w, "Hello world!\n")
} else {
fmt.Println("Reach request limiting!")
io.WriteString(w, "Reach request limit!\n")
}
}
func main() {
fmt.Println("Server Started!")
http.HandleFunc("/", helloHandler)
http.ListenAndServe(":8000", nil)
}
完整的代码放在了 Github 上.
功能测试
在代码中我们的默认设定是在 10 秒钟内最多只服务 5 个请求. 我们可以每次并行发送 3 个请求看返回结果:
➜ JustDoIt git:(master) seq 3 | xargs -P10 -I% curl localhost:8000
Hello world!
Hello world!
Hello world!
➜ JustDoIt git:(master) seq 3 | xargs -P10 -I% curl localhost:8000
Hello world!
Hello world!
Reach request limit!
可以看到发送到第 6 个请求时就触发了限流操作, 和我们预期的行为是一致的.
流量统计
流量统计的实现原理也是类似, 先启动一个计数器, 每次请求都会把计数器加一, 同时再启动一个定时器, 每隔一秒就会把当前计数器的值保存下来, 然后再把计数器清零. 代码如下:
var QPS []CountQPS
type CountQPS struct {
CountPerSecond int
Timestamp int64
}
type CounterService struct {
CountQPS
CountAll int
Lock sync.Mutex
}
func NewCounterService() *CounterService {
counter := &CounterService{}
go func() {
ticker := time.NewTicker(time.Second)
for {
<-ticker.C
counter.Lock.Lock()
counter.Timestamp = time.Now().Unix()
if counter.CountPerSecond > 0 {
QPS = append(QPS, CountQPS{counter.CountPerSecond, counter.Timestamp})
}
counter.CountPerSecond = 0
counter.Lock.Unlock()
}
}()
return counter
}
func (counter *CounterService) Increase() {
counter.Lock.Lock()
defer counter.Lock.Unlock()
counter.CountAll++
counter.CountPerSecond++
}
完整的代码保存在 Github 上.
在上面的代码中我们只是把每秒的统计值保存在一个 slice 中, 在实际的项目时间中我们可以把这个信息保存在 influxdb 这样的数据库或者 Kafka 中.
功能测试
在上面的完整代码中我们还实现了一个 get_cnt 的 api , 通过它可以打印出当前的所有流量统计值. 让我们来先发送一些请求, 然后看结果是否符合预期:
➜ JustDoIt git:(master) seq 2 | xargs -P10 -I% curl localhost:8000
Hello world!
Hello world!
➜ JustDoIt git:(master) seq 5 | xargs -P10 -I% curl localhost:8000
Hello world!
Hello world!
Hello world!
Hello world!
Hello world!
➜ JustDoIt git:(master) curl localhost:8000/get_cnt
timestamp,query_per_second
1525496731,2
1525496734,5
total: 7
可以看到这个结果也是符合预期的!
总结
本文给出的限流和流量统计的方法是比较简单基础的实现, 在某些情况下会有问题, 如参考文献中提到的在限流时间间隔的特定时间点 发送请, 可能会导致实际流量是设计限制流量的两倍. 但本文的方法用在 demo 类的项目中应该也是没什么问题的 :)
欢迎大家多多留言交流!
