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 JedisGetRateLimiter {
    private static final String IP = "192.168.0.163";
    private String luaScript;

    private String key;

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

    }

    public boolean acquire() {
        try (Jedis jedis = new Jedis(IP, 6379)) {
            return (Long) jedis.eval(luaScript, 1, key) == 1L;
        }
    }
}


import java.io.IOException;
import java.nio.charset.Charset;
import java.util.Timer;
import java.util.TimerTask;
import java.util.concurrent.TimeUnit;

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 JedisInitRateLimiter implements AutoCloseable {
    private static final String IP = "192.168.0.163";
    private String luaScript;

    private Timer timer;

    private final Jedis jedis = new Jedis(IP, 6379);

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

        timer = new Timer();

        // 放入令牌的时间间隔
        long period = 1000L / Long.valueOf(limit);
        // 通过定时器，定时放入令牌
        timer.scheduleAtFixedRate(new TimerTask() {
            @Override
            public void run() {
                System.out.println(
                        System.currentTimeMillis() + " 放入令牌：" + ((Long) jedis.eval(luaScript, 1, key, limit) == 1L));
            }

        }, period, period);
    }

    @Override
    public void close() throws Exception {
        this.jedis.close();
        this.timer.cancel();
    }

    public static void main(String[] args) throws Exception {
        JedisInitRateLimiter jrls = new JedisInitRateLimiter("initRateLimit.lua", "test1", "5");
        TimeUnit.SECONDS.sleep(10L);
        jrls.close();
    }
}

Lua脚本

--令牌桶初始化操作  initRateLimit.lua
local key = KEYS[1] 			--限流KEY
local limit = tonumber(ARGV[1]) --容量
-- 获取当前令牌数
local current = tonumber(redis.call('get', key) or "0")
if current + 1 > limit then --如果超出容量
    return 0
else
	redis.call("INCRBY", key, "1") --令牌数+1
end
return 1  --返回1代表不限流


--取令牌  getRateLimit.lua
local key = KEYS[1] 			--限流KEY
-- 获取当前可用令牌数
local current = tonumber(redis.call('get', key) or "0")
if current <= 0 then --没有令牌了
    return 0
else
	redis.call("DECRBY", key, "1") --令牌数-1
end
return 1  --返回1代表不限流

测试类

//Controller层测试类	
//初始化给入的容器数量
	JedisInitRateLimiter initRateLimiter = new JedisInitRateLimiter("initRateLimit.lua","test","5");
	JedisGetRateLimiter jtwl = new JedisGetRateLimiter("getRateLimit.lua","test");
	public String doQuery2(String name) throws Exception {
		// 从redis 上获得 自增后的值
		if (!jtwl.acquire()) {
			return System.currentTimeMillis() / 1000 + "秒杀结束，谢谢参与！";
		}
		return System.currentTimeMillis() / 1000 + "恭喜，秒杀成功！";
	}



//漏桶算法
	 @Test
	public void smoothRequestTest() 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.doQuery2("test"));
				} 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.doQuery2("Mike"));
	}

测试结果：