基于redis实现限流器

package org.jsirenia.ratelimit;

import java.io.FileNotFoundException;
import java.util.List;

import org.jsirenia.file.MyFile;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.ResourceUtils;

import com.google.common.collect.Lists;

import redis.clients.jedis.Jedis;
/**
 * 限流器
 */
public class RateLimiter {
	protected Logger logger = LoggerFactory.getLogger(RateLimiter.class);

	private Jedis jedis;
	private String script;
	private static List emptyList = Lists.newArrayList();

	public RateLimiter(Jedis jedis){
		this.jedis = jedis;
		String file = "classpath:ratelimit2.lua";
		try {
			script = new MyFile(ResourceUtils.getFile(file)).readText();
		} catch (FileNotFoundException e) {
			throw new RuntimeException(e);
		}
	}
	/**
	 * 申请令牌
	 * @return
	 */
	public boolean applyToken(String service){
		Object result = jedis.eval(script, emptyList,
				Lists.newArrayList(service, System.currentTimeMillis()/1000 + ""));
		return result.equals(1);
	}
	/**
	 * 清理token
	 * @param args
	 * @throws FileNotFoundException
	 * @throws InterruptedException
	 */
	public void clear(String service){
		jedis.del(service+":tokens");
		jedis.del(service+":timestamp");
	}
	public static void main(String[] args) throws FileNotFoundException, InterruptedException {
		Jedis jedis = new Jedis("localhost", 6379);
		try {
			//清理环境
			jedis.del("myservice:tokens");
			jedis.del("myservice:timestamp");
			String script = new MyFile(ResourceUtils.getFile("classpath:ratelimit2.lua")).readText();
			for(int i=0;i<500;i++){
				Object result = jedis.eval(script, Lists.newArrayList(),
						Lists.newArrayList("myservice", System.currentTimeMillis()/1000 + ""));
				System.out.println(result);
				Thread.sleep(1000);
			}
		} finally {
			jedis.close();
		}
	}
}

ratelimit.lua

--redis调用 eval script keysnumber keys ARGV
--例子：假设服务名称为service
--redis.eval(luascript,[],["service",System.currentTimeMillis()/1000])
--令牌桶容量，常量
local capacity = 10 
--最后一次取令牌的时间，从redis中读写
local timestampkey = ARGV[1]..':timestamp'
local tokenskey = ARGV[1]..':tokens'
local timestamp = redis.call('get',timestampkey)
local now = tonumber(ARGV[2])

if not timestamp then
	redis.call('set',timestampkey,now)
	redis.call('set',tokenskey,capacity-1)
	return 1
end
--生成令牌的速率,单位秒，常量（也可以参数传入，暂时定为常量）
local rate = 5 
--最后一次取令牌时剩余令牌数，从redis中读写
local tokens = redis.call('get',tokenskey)
--被限流的服务名称，参数传入
local service = ARGV[1]

local function grant() --获取一个令牌
	tokens = tokens+(now-timestamp)*rate
	if tokens > capacity then
		tokens = capacity
	end
	redis.call('set',ARGV[1]..':timestamp',now)
	if tokens<1 then
		redis.call('set',tokenskey,tokens)
		return 0
	else
		tokens=tokens-1
		redis.call('set',ARGV[1]..':timestamp',now)
		redis.call('set',tokenskey,tokens)
		return 1
	end
end
return grant()
--[[
多行注释
注意返回值，不要使用true和false，因为redis执行lua脚本，返回true时，在java里面是1（long），
返回false时，在java里面是null。
所以，这里设计成，取到令牌返回1，没取到令牌返回0。
--]]