java 漏桶算法实现,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"));

}

测试结果：