Redis + Lua实现简单的限流

1. 漏桶:

package cn.gw.common.traffic;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;

/**
 * @author gaowu
 * @date 2022/9/20 15:11
 */
@Component
public class RedisLeakyBucketTrafficLimiter implements ILeakyBucketTrafficLimiter {

    /**
     * $lb = LeakyBucket
     **/
    private static final String GROUP_KEY = "$lb";

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * KEYS[1] = GROUP_KEY
     * KEYS[2] = water    上次水量
     * KEYS[3] = time     上次请求时间(s)
     * ARGV[1] = capacity 最大蓄水量
     * ARGV[2] = rate     水流速率/s
     * ARGV[3] = now_time 当前请求时间(s)
     * ARGV[4] = acquire  请求个数
     */
    private static final String LEAKEY_BUCKET_SCRIPT = " local capacity = tonumber(ARGV[1])" +
            " local rate = tonumber(ARGV[2])" +
            " local now = tonumber(ARGV[3])" +
            " local acquire = tonumber(ARGV[4])" +
            // 上次请求后桶中水量
            " local water = tonumber(redis.call('hget', KEYS[1] , KEYS[2]) or 0) " +
            // 上次请求时间(/s)
            " local time = tonumber(redis.call('hget', KEYS[1] , KEYS[3]) or now) " +
            // 当前桶中的水量 = 上次请求后桶中水量 - 上次到现在流出的水量
            " water = math.max(0, water - (now - time) * rate)" +
            // 设置请求时间(/s)
            " redis.call('hset' , KEYS[1] ,KEYS[3] , now)" +
            // 判断桶中水是否满了
            " if (water + acquire <= capacity) then" +
            // 未满注水
            " redis.call('hset' , KEYS[1] , KEYS[2] , water + acquire)" +
            " return 1" +
            " else" +
            // 满了直接返回
            " return 0" +
            " end";


    @Override
    public Long limit(Long capacity, Long rate, String key) {
        return limit(capacity, rate, 1L, key);
    }

    public Long limit(Long capacity, Long rate, Long acquire, String key) {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>(LEAKEY_BUCKET_SCRIPT, Long.class);
        List<String> keys = Arrays.asList(GROUP_KEY, key + ":water", key + ":time");
        return redisTemplate.execute(redisScript, keys, capacity, rate, System.currentTimeMillis() / 1000, acquire);
    }
}

2.令牌桶：

package cn.gw.common.traffic;

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.Arrays;
import java.util.List;

/**
 * @author gaowu
 * @date 2022/9/21 15:48
 */
@Component
public class RedisTokenBucketTrafficLimiter implements ITokenBucketTrafficLimiter {

    /**
     * $tb = TokenBucket
     **/
    private static final String GROUP_KEY = "$tb";

    @Resource
    private RedisTemplate<String, Object> redisTemplate;

    /**
     * KEYS[1] = GROUP_KEY
     * KEYS[2] = residual  上次剩余token数量
     * KEYS[3] = time     上次请求时间
     * ARGV[1] = capacity token容量
     * ARGV[2] = rate     token生成速率/s
     * ARGV[3] = now_time 当时间请求(s)
     * ARGV[4] = default_token 默认token数量
     * AVG[5] = acquire_token  获取token数据量
     */
    private static final String TOKEN_BUCKET_SCRIPT = " local capacity = tonumber(ARGV[1])" +
            " local rate = tonumber(ARGV[2])" +
            " local now = tonumber(ARGV[3])" +
            " local default_token = tonumber(ARGV[4])" +
            " local acquire_token = tonumber(ARGV[5])" +
            // 上次请求后桶中剩余的令牌数量
            " local residual_token = tonumber(redis.call('hget', KEYS[1] , KEYS[2]) or default_token)" +
            // 上次请求时间
            " local time = tonumber(redis.call('hget', KEYS[1] , KEYS[3]) or now)" +
            // 设置请求时间为now
            " redis.call('hset' , KEYS[1] , KEYS[3] , now)" +
            // 当前令牌数 = 上次请求后桶中剩余的令牌数量 + 上次到现在生成的令牌数量
            " local current_token = math.min(residual_token + (now - time) * rate , capacity)" +
            // 判断令牌数量是否够本次申请
            " if (current_token -acquire_token >= 0) then" +
            // 够减去桶中令牌数量
            " redis.call('hset' , KEYS[1] , KEYS[2] , current_token - acquire_token)" +
            " return 1" +
            " else" +
            // 不够直接返回
            " return 0" +
            " end";

    public Long limit(Long capacity, Long rate, Long defaultToken, Long acquire, String key) {
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>(TOKEN_BUCKET_SCRIPT, Long.class);
        List<String> keys = Arrays.asList(GROUP_KEY, key + ":token", key + ":time");
        return redisTemplate.execute(redisScript, keys, capacity, rate, System.currentTimeMillis() / 1000, defaultToken, acquire);
    }

    public Long limit(Long capacity, Long rate, Long defaultToken, String key) {
        return limit(capacity, rate, defaultToken, 1L, key);
    }
}

3.滑块(单机)：

/**
 * @author gaowu
 * @date 2023/2/18 16:07
 */
public class FlowControl {

    private int capacity = 1000;

    private int index = 0;

    private long[] controls;

    private int timeIntervals = 1000;

    public boolean limit(long time) {
        if (controls[index] == 0 || time - controls[index] < timeIntervals) {
            controls[index] = time;
            index = (index + 1) % capacity;
            return true;
        }
        return false;
    }

    public FlowControl(int capacity, int timeIntervals) {
        if (capacity <= 0) {
            capacity = this.capacity;
        }
        if (timeIntervals <= 0) {
            timeIntervals = this.timeIntervals;
        }
        this.capacity = capacity;
        this.controls = new long[capacity];
        this.timeIntervals = timeIntervals;
    }

    public FlowControl() {

    }