golang gin框架中實現一個簡單的不是特別精確的秒級限流器

起因

看了兩篇關於golang中限流器的帖子：

• Gin 開發實踐：如何實現限流中間件
• 常用限流策略——漏桶與令牌桶介紹

我照着用，居然沒效果……

時間有限沒有深究。這實在是一個很簡單的功能，我的需求是每分鍾限制大約xxx次請求，平均每秒限制到XXX除60次左右的請求也就夠了。

限流器代碼

於是下面用了個及其簡單的方式來實現限流器：
utils/limiter.go

package utils

import (
	"sync/atomic"
	"time"
)

//Limiter 限流器對象
type Limiter struct {
	value  int64
	max    int64
	ts     int64
}

//NewLimiter 產生一個限流器
func NewLimiter(cnt int64) *Limiter {
	return &Limiter{
		value:  0,
		max:    cnt,
		ts:     time.Now().Unix(),
	}
}

//Ok 是否可以通過
func (l *Limiter) Ok() bool {
	ts := time.Now().Unix()
	tsOld := atomic.LoadInt64(&l.ts)
	if ts != tsOld {
		atomic.StoreInt64(&l.ts, ts)
		atomic.StoreInt64(&l.value, 1)
		return true
	}
	return atomic.AddInt64(&(l.value), 1) < l.max
}

//SetMax 設置最大限制
func (l *Limiter) SetMax(m int64) {
	l.max = m
}

gin框架中使用

main.go

//MaxAllowed 限流器
func MaxAllowed(limitValue int64) func(c *gin.Context) {
	limiter := utils.NewLimiter(limitValue)
	log.Println("limiter.SetMax:", limitValue)
	// 返回限流邏輯
	return func(c *gin.Context) {
		if !limiter.Ok() {
			c.AbortWithStatus(http.StatusServiceUnavailable)  //超過每秒200，就返回503錯誤碼
			return
		}
		c.Next()
	}
}

func main(){
    router := gin.New()
    router.Use(MaxAllowed(200))  //限制每秒最多允許200個請求
}

have fun!