可以选择普通模式和LUA脚本模式,建议选择普通模式,实际上不需要控制的那么精确。
package Middlewaresimport ('github.com/gin-gonic/gin''strconv''time''voteapi/pkg/app/response''voteapi/pkg/gredis''voteapi/pkg/util')const IP_LIMIT_NUM_KEY = 'ipLimit:ipLimitNum'const IP_BLACK_LIST_KEY = 'ipLimit:ipBlackList'var prefix = '{gateway}'var delaySeconds int64 = 60 // 观察时间跨度,秒var maxAttempts int64 = 10000 // 限制请求数var blackSeconds int64 = 0 // 封禁时长,秒,0-不封禁func GateWayPlus() gin.HandlerFunc {return func(c *gin.Context) {path := c.FullPath()clientIp := c.ClientIP()// redis配置集群时必须param := make(map[string]string)param['path'] = pathparam['clientIp'] = clientIpif !main(param) {c.Abort()response.JsonResponseError(c, '当前IP请求过于频繁,暂时被封禁~')}}}func main(param map[string]string) bool {// 预知的IP黑名单var blackList []stringif util.InStringArray(param['clientIp'], blackList) {return false}// 预知的IP白名单var whiteList []stringif util.InStringArray(param['clientIp'], whiteList) {return false}blackKey := prefix + ':' + IP_BLACK_LIST_KEYlimitKey := prefix + ':' + IP_LIMIT_NUM_KEYcurr := time.Now().Unix()item := util.Md5(param['path'] + '|' + param['clientIp'])return normal(blackKey, limitKey, item, curr)}// 普通模式func normal(blackKey string, limitKey string, item string, time int64) (res bool) {if blackSeconds > 0 {timeout, _ := gredis.RawCommand('HGET', blackKey, item)if timeout != nil {to, _ := strconv.Atoi(string(timeout.([]uint8)))if int64(to) > time {// 未解封return false}// 已解封,移除黑名单gredis.RawCommand('HDEL', blackKey, item)}}l, _ := gredis.RawCommand('HGET', limitKey, item)if l != nil {last, _ := strconv.Atoi(string(l.([]uint8)))if int64(last) >= maxAttempts {return false}}num, _ := gredis.RawCommand('HINCRBY', limitKey, item, 1)if ttl, _ := gredis.TTLKey(limitKey); ttl == int64(-1) {gredis.Expire(limitKey, int64(delaySeconds))}if num.(int64) >= maxAttempts && blackSeconds > 0 {// 加入黑名单gredis.RawCommand('HSET', blackKey, item, time+blackSeconds)// 删除记录gredis.RawCommand('HDEL', limitKey, item)}return true}// LUA脚本模式// 支持redis集群部署func luaScript(blackKey string, limitKey string, item string, time int64) (res bool) {script := `local blackSeconds = tonumber(ARGV[5])if(blackSeconds > 0)then local timeout = redis.call(’hget’, KEYS[1], ARGV[1]) if(timeout ~= false) then if(tonumber(timeout) > tonumber(ARGV[2])) then return false end redis.call(’hdel’, KEYS[1], ARGV[1]) endendlocal last = redis.call(’hget’, KEYS[2], ARGV[1])if(last ~= false and tonumber(last) >= tonumber(ARGV[3]))then return falseendlocal num = redis.call(’hincrby’, KEYS[2], ARGV[1], 1)local ttl = redis.call(’ttl’, KEYS[2])if(ttl == -1)then redis.call(’expire’, KEYS[2], ARGV[4])endif(tonumber(num) >= tonumber(ARGV[3]) and blackSeconds > 0)then redis.call(’hset’, KEYS[1], ARGV[1], ARGV[2] + ARGV[5]) redis.call(’hdel’, KEYS[2], ARGV[1])endreturn true`result, err := gredis.RawCommand('EVAL', script, 2, blackKey, limitKey, item, time, maxAttempts, delaySeconds, blackSeconds)if err != nil {return false}if result == int64(1) {return true} else {return false}}
补充:golang实现限制每秒多少次的限频操作
前言一些函数的执行可能会限制频率,比如某个api接口要求每秒最大请求30次。下面记录了自己写的限频和官方的限频
代码
// 加锁限频,输出次数大概率小于最大值func ExecLimit(lastExecTime *time.Time, l *sync.RWMutex ,maxTimes int, perDuration time.Duration, f func()) { l.Lock() defer l.Unlock() // per times cost time(s) SecondsPerTimes := float64(perDuration) / float64(time.Second) / float64(maxTimes) now := time.Now() interval := now.Sub(*lastExecTime).Seconds() if interval < SecondsPerTimes { time.Sleep(time.Duration(int64((SecondsPerTimes-interval)*1000000000)) * time.Nanosecond) } f() *lastExecTime = time.Now()}// 官方的,需要引用 'golang.org/x/time/rate'// 基本上可以达到满值,比自己写的更优func ExecLimit2(l *rate.Limiter, f func()) { go func() { l.Wait(context.Background()) f() }()}
使用
func TestExecLimit(t *testing.T) { runtime.GOMAXPROCS(runtime.NumCPU()) go func() { var lastExecTime time.Time var l sync.RWMutex for { ExecLimit(&lastExecTime, &l, 10, time.Second, func() { fmt.Println('do') }) } }() select { case <-time.After(1 * time.Second): fmt.Println('1秒到时') }}func TestExecLimit2(t *testing.T) { runtime.GOMAXPROCS(runtime.NumCPU()) l := rate.NewLimiter(1, 30) go func() { for { ExecLimit2(l, func() { fmt.Println('do') }) } }() select { case <-time.After(1 * time.Second): fmt.Println('1秒到时') }}
输出:
一秒内输出了<=10次 'do'
如何在多节点服务中限制频上述使用,定义在某个服务节点的全局变量lastExecTime仅仅会对该服务的函数f()操作限频,如果在负载均衡后,多个相同服务的节点,对第三方的接口累计限频,比如三个服务共同拉取第三方接口,合计限频为30次/s.
则,必须将lastExecTime的获取,从redis等共享中间件中获取,而不应该从任何一个单点服务获取。
以上为个人经验,希望能给大家一个参考,也希望大家多多支持乐呵呵网。如有错误或未考虑完全的地方,望不吝赐教。