redis 令牌桶算法—Lua

迈不过友情╰ 2023-07-10 08:24 11阅读 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 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"));
    	}

结果：

![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/20200301130415607.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/202003011317541.png?x-oss-process=image/watermark,type_ZmFuZ3poZW5naGVpdGk,shadow_10,text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L3NkbXhkemI=,size_16,color_FFFFFF,t_70