redis 令牌桶算法—Lua

令牌桶算法

场景:秒杀（也可以用于平滑限流某个接口请求）

---

import java.io.IOException;
import java.nio.charset.Charset;
import org.springframework.core.io.ClassPathResource;
import com.google.common.io.Files;
import redis.clients.jedis.Jedis;
/***
 *@author dzb
 *@date 2019/11/3 22:09
 *@Description:  令牌桶
 * */
public class JedisLuaTokenTimeLimiter {

	/**执行脚本**/
	private String luaScript;
	/**值**/
	private String key;
	/**限流次数**/
	private String limit;
	/**时间**/
	private String expire;

	public JedisLuaTokenTimeLimiter(String key, String limit, String expire, String scriptFile) {
		super();
		this.key = key;
		this.limit = limit;
		this.expire = expire;
		try {
			luaScript = Files.asCharSource(new ClassPathResource(scriptFile).getFile(), Charset.defaultCharset())
					.read();
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	/***
	 * 释放获取
	 * */
	public boolean acquire() {
		Jedis jedis = new Jedis("localhost", 6379);
		return (Long) jedis.eval(luaScript, 1, key, limit, expire) == 1L;
	}
}

--令牌秒杀 Lua脚本  tokenTimeLimiter.lua
local key = KEYS[1] 			--限流KEY（一秒一个）
local limit = tonumber(ARGV[1]) --限流大小
local exprie = ARGV[2] 			--过期时间
-- 获取当前计数值
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limit then --如果超出限流大小
    return 0
else
	current = tonumber(redis.call("INCRBY", key, "1")) --请求数+1
	if current == 1 then  --第一次访问需要设置过期时间
    	redis.call("expire", key,exprie) --设置过期时间
    end
end
return 1  --返回1代表不限流

Controller层
//模拟秒杀场景	
JedisLuaTokenTimeLimiter jtwlm = new JedisLuaTokenTimeLimiter("商品A", "5", "1", "tokenTimeLimiter.lua");
	public String doQuery(String name) throws Exception {
		// 从redis 上获得 自增后的值
		if (!jtwlm.acquire()) {
			return System.currentTimeMillis() / 1000 + "秒杀结束，谢谢参与！";
		}
		return System.currentTimeMillis() / 1000 + "恭喜，秒杀成功！";
	}

测试代码

	//模拟场景并发
	 @Test
	public void tokenTimeLimiter() throws Exception {
		CountDownLatch cdl = new CountDownLatch(10);
		CyclicBarrier cyb = new CyclicBarrier(10);
		for (int i = 0; i < 10; i++) {
			new Thread(() -> {
				try {
					cyb.await();
				} catch (InterruptedException | BrokenBarrierException e) {
					e.printStackTrace();
				}
				try {
					System.out.println(Thread.currentThread().getName() + " " + orderc.doQuery("商品A"));
				} catch (Exception e) {
					e.printStackTrace();
				}
				cdl.countDown();
			}).start();
		}

		try {
			cdl.await();
		} catch (InterruptedException e) {
			e.printStackTrace();
		}
		TimeUnit.SECONDS.sleep(1);
		System.out.println(Thread.currentThread().getName() + " " + orderc.doQuery("商品A"));
	}

结果：

注：本人只是简单的模拟了分布式环境并发的场景，其细节还有很多就没有过多给出。