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

    话不多说了，接口限流算法有很多种 基于redis实现的例子，

/**
 * 
 */
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;
	}