基于Redis+LUA脚本的令牌桶算法限流策略实现

/**
 * 
 */
package com.matrix.cloud.service.redis.ratelimit;

import java.util.Collections;

import javax.annotation.Resource;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.data.redis.core.RedisCallback;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.data.redis.serializer.StringRedisSerializer;
import org.springframework.stereotype.Service;

/**
 * 1 使用redis+lua脚本实现接口限流
 *   基于Redis的令牌桶算法限流策略实现  测试通过
 * @author lmc
 * 2018年12月21日
 */
@Service
public class RateLimitClient {
	
	@Autowired
	StringRedisTemplate  redisTemplate;
	  
	//这里需要到redis配置文件下配置相关bean
	@Qualifier(\"ratelimitInitLua\")
	@Resource
	RedisScript<Long> ratelimitInitLua;
	
	
	 public boolean initToken(String key){
	        boolean token;
	      //此序列化redis会把数字类字符串存为数字类型
		  redisTemplate.setKeySerializer(new StringRedisSerializer());
		  redisTemplate.setValueSerializer(new StringRedisSerializer());
	        Long currMillSecond = redisTemplate.execute(
	                (RedisCallback<Long>) redisConnection -> redisConnection.time()
	        );
	        /** 初始化接收到的参数
	         * redis.pcall(\"HMSET\",KEYS[1],  \"last_mill_second\",ARGV[1],  \"curr_permits\",ARGV[2],
	         * \"max_burst\",ARGV[3], \"rate\",ARGV[4], \"app\",ARGV[5])
	         */
	        StringBuffer ratelimitStr= new StringBuffer();
	        ratelimitStr.append(\" local result=1 \");
	        ratelimitStr.append(\" redis.pcall(\'HMSET\',KEYS[1], \");
	        ratelimitStr.append(\" \'last_mill_second\',ARGV[1], \");
	        ratelimitStr.append(\" \'curr_permits\',ARGV[2], \");
	        ratelimitStr.append(\" \'max_burst\',ARGV[3], \");
	        ratelimitStr.append(\" \'rate\',ARGV[4], \");
	        ratelimitStr.append(\" \'app\',ARGV[5]) \");
	        ratelimitStr.append(\" return result \");
	     	DefaultRedisScript<Long> ratelimitLua = new DefaultRedisScript<>(ratelimitStr.toString(), Long.class);
	     	
			String last_mill_second =  String.valueOf(currMillSecond);//上一次添加令牌的毫秒数
			String curr_permits = \"3\";//令牌桶的最少令牌数
			String max_permits = \"200\";//令牌桶的最大令牌数
			String rate = \"100\";//向令牌桶中添加令牌的速率  , 令牌消耗速率
			String app = \"skynet\";//定义标记，比如哪些是被限流的
	        Long accquire = redisTemplate.execute(ratelimitLua,Collections.singletonList(\"ratelimit:\"+key), currMillSecond.toString(), curr_permits, max_permits, rate, app);
	        if (accquire == 1) {
	            token = true;
	        } else if (accquire == 0) {
	            token = true;
	        } else {
	            token = false;
	        }
	        return token;
	    }
	 
		/*
		 *  last_mill_second 最后时间毫秒 
			curr_permits 当前可用的令牌 
			max_burst 令牌桶最大值 
			rate 每秒生成几个令牌 
			app 应用 
			令牌桶内令牌生成借鉴Guava-RateLimiter类的设计 
			每次根据时间戳生成token，不超过最大值
			permits 每次请求令牌数
		 */
	     public boolean accquireToken(String key, Integer permits) {
	        boolean token;
	        redisTemplate.setKeySerializer(new StringRedisSerializer());
			redisTemplate.setValueSerializer(new StringRedisSerializer());
	        Long currMillSecond = redisTemplate.execute(
	                (RedisCallback<Long>) redisConnection -> redisConnection.time()
	        );

	        Long accquire = redisTemplate.execute(ratelimitInitLua,
	                Collections.singletonList(\"ratelimit:\"+key), permits.toString(), currMillSecond.toString());
	        if (accquire == 1) {
	            token = true;
	        } else {
	            token = false;
	        }
	        return token;
	    }
	
}

ratelimitInit.lua   配置文件加载进来

local ratelimit_info=redis.pcall(\"HMGET\",KEYS[1],\"last_mill_second\",\"curr_permits\",\"max_burst\",\"rate\",\"app\")
local last_mill_second=ratelimit_info[1]
local curr_permits=tonumber(ratelimit_info[2])
local max_burst=tonumber(ratelimit_info[3])
local rate=tonumber(ratelimit_info[4])
local app=tostring(ratelimit_info[5])
if app == nil then
    return 0
end

local local_curr_permits=max_burst;

if(type(last_mill_second) ~=\'boolean\' and last_mill_second ~=nil) then
    local reverse_permits=math.floor((ARGV[2]-last_mill_second)/1000)*rate
    if(reverse_permits>0) then
        redis.pcall(\"HMSET\",KEYS[1],\"last_mill_second\",ARGV[2])
    end

    local expect_curr_permits=reverse_permits+curr_permits
    local_curr_permits=math.min(expect_curr_permits,max_burst);

else
    redis.pcall(\"HMSET\",KEYS[1],\"last_mill_second\",ARGV[2])
end

local result=-1
if(local_curr_permits-ARGV[1]>0) then
    result=1
    redis.pcall(\"HMSET\",KEYS[1],\"curr_permits\",local_curr_permits-ARGV[1])
else
    redis.pcall(\"HMSET\",KEYS[1],\"curr_permits\",local_curr_permits)
end

return result

调用 方式  方式简单点

	@RequestMapping(value = \"/getRatelimitlua\", method = RequestMethod.POST)
	@ResponseBody
	public Object RatelimitRedisLua() {//@RequestParam(value=\"uid\") String uid
		
		String key=\"lmc168\";
		
		rateLimitClient.initToken(key);//初如化开始
		
		if (!rateLimitClient.accquireToken(key, 1)) {
	        	System.out.println(\"触发限流API:调用太忙了,请休息下\");
				//throw new Exception();
	    }else{
		  System.out.println(\"没有触发限流策略\");
	    }
		return key;
	}