redis 漏桶算法—Lua

柔情只为你懂 2023-07-10 08:25 32阅读 0赞

漏桶算法

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI_size_16_color_FFFFFF_t_70][]

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"));
    	}

测试结果：

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI_size_16_color_FFFFFF_t_70 1][]

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI_size_16_color_FFFFFF_t_70]: https://img-blog.csdnimg.cn/20200301132308929.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI=,size_16,color_FFFFFF,t_70
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI_size_16_color_FFFFFF_t_70 1]: https://img-blog.csdnimg.cn/20200301133455820.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI=,size_16,color_FFFFFF,t_70