redis实现限流

令牌桶逻辑

计算逻辑：

代码：

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;/*** @ClassName RedisRateLimiterTokenBucket* @Description TODO* @Author zhang zhengdong* @DATE 2025/1/17 20:22* @Version 1.0*/
public class RedisRateLimiterTokenBucket {private JedisPool jedisPool;private String rateLimitKey;private int timeWindow; // 限制的时间窗口，单位为秒private int maxTokens; // 每个时间窗口的最大令牌数private int refillRate; // 每秒钟生成的令牌数public RedisRateLimiterTokenBucket(JedisPool jedisPool, String rateLimitKey, int timeWindow, int maxTokens, int refillRate) {this.jedisPool = jedisPool;this.rateLimitKey = rateLimitKey;this.timeWindow = timeWindow;this.maxTokens = maxTokens;this.refillRate = refillRate;}/*** 尝试获取令牌*/public boolean acquire() {try (Jedis jedis = jedisPool.getResource()) {long currentTime = System.currentTimeMillis() / 1000;String key = rateLimitKey + ":tokens";String lastRefillKey = rateLimitKey + ":lastRefillTime";// 获取最后一次令牌补充的时间String lastRefillTime = jedis.get(lastRefillKey);if (lastRefillTime == null) {jedis.set(lastRefillKey, String.valueOf(currentTime));jedis.set(key, String.valueOf(maxTokens));lastRefillTime = String.valueOf(currentTime);}// 计算补充令牌的数量long elapsed = currentTime - Long.parseLong(lastRefillTime);int tokensToAdd = (int) (elapsed * refillRate);// 计算当前令牌数int tokens = Integer.parseInt(jedis.get(key));tokens = Math.min(tokens + tokensToAdd, maxTokens);// 如果令牌数大于 0，则消费一个令牌if (tokens > 0) {jedis.set(key, String.valueOf(tokens - 1));jedis.set(lastRefillKey, String.valueOf(currentTime));return true;}return false; // 没有令牌，拒绝请求}}
}

验证代码

import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;/*** @ClassName RedisRateLimiterTokenBucketTest* @Description TODO* @Author zhang zhengdong* @DATE 2025/1/17 20:25* @Version 1.0*/
public class RedisRateLimiterTokenBucketTest {public static void main(String[] args) {// 初始化 Jedis 连接池JedisPoolConfig poolConfig = new JedisPoolConfig();poolConfig.setMaxTotal(10);JedisPool jedisPool = new JedisPool(poolConfig, "localhost", 6379);// 限流设置：10 秒内最多 20 次请求int timeWindow = 10;int maxRequestsPerWindow = 10;// 使用 令牌桶 算法实现限流RedisRateLimiterTokenBucket rateLimiterTokenBucket = new RedisRateLimiterTokenBucket(jedisPool,"rate_limit_token_bucket", timeWindow, maxRequestsPerWindow, 1);for (int i = 0; i < 25; i++) {boolean allowed = rateLimiterTokenBucket.acquire();if (allowed) {System.out.println("请求 " + (i + 1) + " 被允许");} else {System.out.println("请求 " + (i + 1) + " 被拒绝");}//			try {
//				Thread.sleep(400);  // 模拟请求间隔
//			} catch (InterruptedException e) {
//				e.printStackTrace();
//			}}}
}

zset实现限流逻辑

import redis.clients.jedis.Jedis;
import redis.clients.jedis.JedisPool;import java.util.UUID;/*** @ClassName RedisRateLimiterZSet* @Description TODO* @Author zhang zhengdong* @DATE 2025/1/17 19:50* @Version 1.0*/
public class RedisRateLimiterZSet {private JedisPool jedisPool;private String rateLimitKey;private int timeWindow; // 限制的时间窗口，单位为秒private int maxRequestsPerWindow; // 每个窗口允许的最大请求次数public RedisRateLimiterZSet(JedisPool jedisPool, String rateLimitKey, int timeWindow, int maxRequestsPerWindow) {this.jedisPool = jedisPool;this.rateLimitKey = rateLimitKey;this.timeWindow = timeWindow;this.maxRequestsPerWindow = maxRequestsPerWindow;}public boolean acquire() {try (Jedis jedis = jedisPool.getResource()) {Long currentTime = System.currentTimeMillis();System.out.println(currentTime);if (jedis.zcard(rateLimitKey) > 0) {Integer count = jedis.zrangeByScore(rateLimitKey, currentTime - 10 * 1000, currentTime).size(); // 注意这里使用zrangeByScore，以时间作为score。zrange key start stop 命令的start和stop是序号。System.out.println(count);if (count != null && count > maxRequestsPerWindow) {return false;}}jedis.zadd(rateLimitKey, Double.valueOf(currentTime), UUID.randomUUID().toString());return true;}}}

测试

import redis.clients.jedis.JedisPool;
import redis.clients.jedis.JedisPoolConfig;/*** @ClassName RateLimiterZSetTest* @Description TODO* @Author zhang zhengdong* @DATE 2025/1/17 20:00* @Version 1.0*/
public class RedisRateLimiterZSetTest {public static void main(String[] args) {JedisPoolConfig jedisConfig = new JedisPoolConfig();jedisConfig.setMaxTotal(10);JedisPool jedisPool = new JedisPool(jedisConfig, "localhost", 6379);// 限流设置：10 秒内最多 20 次请求RedisRateLimiterZSet rateLimiter = new RedisRateLimiterZSet(jedisPool, "rate_limit_zset", 10, 20);// 测试：连续请求for (int i = 0; i < 25; i++) {boolean allowed = rateLimiter.acquire();if (allowed) {System.out.println("请求 " + (i + 1) + " 被允许");} else {System.out.println("请求 " + (i + 1) + " 被拒绝");}// 延迟 400ms 后再次尝试
//			try {
//				Thread.sleep(400);  // 模拟请求间隔
//			} catch (InterruptedException e) {
//				e.printStackTrace();
//			}}}