26.redis实现日限流、周限流（含黑名单、白名单）

代码地址：https://gitee.com/lymgoforIT/golang-trick/tree/master/18-redis-limit

一：简介

在日常工作中，经常会遇到对某种操作进行频次控制的需求，此时常用的做法是采用redis的incr来递增,记录访问次数, 以及 expire 来设置失效时间.

比如有一个活动，用户完成后可以领取奖励，但是对日和周有一定的频次限制，并且对某些特殊用户，开通白名单和黑名单通道，流程如下

二：go实现

package main

import (
	"context"
	"fmt"
	"github.com/go-redis/redis/v8"
	"strconv"
	"time"
)

var redisClient *redis.Client
var ctx = context.Background()
var DayLimitKey = "RewardKey_%d_%s" // 奖励key_用户id_年月日   redis string类型 用于是否存在该key即可
var WeekLimitKey = "RewardKey_%d"   // 奖励key_用户id redis string类型，用于计数
var DayExpireTime = time.Duration(86400) * time.Second
var WeekExpireTime = time.Duration(7*86400) * time.Second

func init() {
	config := &redis.Options{
		Addr:         "localhost:6379",
		Password:     "",
		DB:           0, // 使用默认DB
		PoolSize:     15,
		MinIdleConns: 10, //在启动阶段创建指定数量的Idle连接，并长期维持idle状态的连接数不少于指定数量；。
		//超时
		//DialTimeout:  5 * time.Second, //连接建立超时时间，默认5秒。
		//ReadTimeout:  3 * time.Second, //读超时，默认3秒， -1表示取消读超时
		//WriteTimeout: 3 * time.Second, //写超时，默认等于读超时
		//PoolTimeout:  4 * time.Second, //当所有连接都处在繁忙状态时，客户端等待可用连接的最大等待时长，默认为读超时+1秒。
	}
	redisClient = redis.NewClient(config)
}

func main() {
	var userId int64 = 123       // 使用123作为测试用户
	var weekLimitCount int64 = 2 // 假定一周只能发两次奖励

	// 校验
	res, err := limitValidation(userId, weekLimitCount)
	if err != nil {
		fmt.Printf("校验过程出现错误，err:%v", err)
		return
	}

	if !res {
		fmt.Println("校验未通过，无法发奖")
		return
	}

	// 发奖成功后，设置相关redisKey
	fmt.Println("模拟发奖成功。。。")

	today := time.Now().Format("2006-01-02")
	redisClient.Set(ctx, fmt.Sprintf(DayLimitKey, userId, today), 1, DayExpireTime)

	// 周限制由于是要计数，所以需要先判断key是否已经设置过
	weekLimitKey := fmt.Sprintf(WeekLimitKey, userId)
	exists, err := redisClient.Exists(ctx, weekLimitKey).Result()
	if err != nil {
		return
	}
	if exists == 1 { // key存在，计数加1
		redisClient.Incr(ctx, weekLimitKey)
	} else { // 本周首次下发，设置key与过期时间
		redisClient.Set(ctx, weekLimitKey, 1, WeekExpireTime)
	}

}

func limitValidation(userId int64, weekLimitCount int64) (bool, error) {
	// 是否在白名单中，实际工作中，白名单一般配置到远程配置中心、或者rpc接口、DB等
	// 这里为了演示，直接给定一个列表
	var whiteList = []int64{111, 222}
	if isInList(whiteList, userId) {
		// 在白名单中，可以直接发奖
		return true, nil
	}

	// 是否在黑名单中
	var blackList = []int64{333, 444}
	if isInList(blackList, userId) {
		// 在黑名单中，直接拒绝发奖
		fmt.Printf("在黑名单中，直接拒绝发奖，userId:%v\n", userId)
		return false, nil
	}

	// 今日是否已经发过
	today := time.Now().Format("2006-01-02")
	//  存在返回1
	exists, err := redisClient.Exists(ctx, fmt.Sprintf(DayLimitKey, userId, today)).Result()
	if err != nil {
		fmt.Printf("访问redis错误，err:%v", userId)
		return false, err
	}
	if exists == 1 { // 今日已经发过,不可以再发了
		fmt.Printf("今日已经发过,不可以再发了，userId:%v\n", userId)
		return false, nil
	}

	// 本周下发次数是否已经达到上限
	exists, err = redisClient.Exists(ctx, fmt.Sprintf(WeekLimitKey, userId)).Result()
	if err != nil {
		fmt.Printf("访问redis错误，err:%v\n", userId)
		return false, err
	}
	if exists != 1 { // 本周没有发过，可以发
		return true, nil
	}

	result, err := redisClient.Get(ctx, fmt.Sprintf(WeekLimitKey, userId)).Result()
	if err != nil {
		fmt.Printf("访问redis错误，err:%v\n", userId)
		return false, err
	}
	count, _ := strconv.ParseInt(result, 10, 64)
	if count >= weekLimitCount {
		fmt.Printf("本周下发次数是否已经达到上限,不可以再发了，userId:%v\n", userId)
		return false, err
	}

	// 以上校验都通过，可以发奖
	return true, nil
}

func isInList(list []int64, userId int64) bool {
	for _, val := range list {
		if val == userId {
			return true
		}
	}
	return false
}

三：测试

1. 日限流

首次执行，肯定显示当日可以发成功

redis客户端查看日限制的key以及过期时间

今日想再次下发

2. 周限流

在日限流测试时，因为发过奖励了，所以也设置了周任务的限流key了的

判断周限流时需要先通过日限流，这里由于是在同一天测试，会被日限流拦住，为了方便测试，直接从redis客户端删除日限流的key，从而模拟为今日还没有发过奖励。删除操作如下

注：ttl命令返回值是键的剩余时间（单位是秒）。当键不存在时，ttl命令会返回-2。没有为键设置过期时间（即永久存在，这是建立一个键后的默认情况）返回-1。

再次执行程序

查看redis如下，可以看到日限流的key又设置上了，符合预期，同时周限流的key增加了1变为了2

注：redis中的incr命令是不会改变key的过期时间的

继续把日限流的key删除，通过日限流的检查，去判断周限流，再一次执行程序

四：lua脚本

存在多个Redis操作的时候，最好还是使用Lua进行操作保证原子性，这里提供一个比较通用的计数Lua脚本。

-- 脚本第一个参数，用作限流的key
local key = KEYS[1]
-- 三个参数分别是上限、过期时间、步长
local upper = ARGV[1]
local expireSecond = ARGV[2]
local step = ARGV[3]

-- 取出current的值，不存在为0
local current = tonumber(redis.call('get', key) or "0")

-- 达到上限，返回-1（达到上限）
-- 当前请求是第一个，设置incrby再设置expire
-- 当前请求不是第一个，直接设置incrby
if current >= tonumber(upper) then
return -1
elseif current == 0 then
redis.call("INCRBY", key, step)
redis.call("EXPIRE", key, expireSecond)
return current + step
else
redis.call("INCRBY", key, step)
return current + step
end