java学习之使用redis简单实现限流功能

事先声明基于redis实现滑动时间窗实现限流这种方式我们工作中一般不会使用,因为目前市面上对于大体量的服务使用的都是springCloud netFlix或者spring CloundAlibaba(主流)他们都有各自的限流组件( Netflix Hystrix , sentinel), 小体量服务又基本用不到限流,这里可以当成是时间窗口算法的一种学习

滑动时间窗口限流是一种基于时间的流量控制策略,本质就是规定一个时间大小的窗口,值统计在这个窗口内的请求,然后滑动这个窗口;至于怎么滑动窗口,可以那请求的当前时间-设置的窗口大小来完成

具体实现:

主要步骤:

1.定义窗口的大小,如30秒

2.每次收到请求,就定义一个zset存储到redis中

3.然后通过ZREMVRANGBYSCORE命令来删除窗口起始时间之前的数据

4.最后通过ZCARD命令获取zset集合中的成员数量,及窗口内的请求量

public class SlidingWindowRateLimiter {
    private final Jedis jedis;
    public SlidingWindowRateLimiter(Jedis jedis) {
        this.jedis = jedis;
    }
    public boolean isAllowed(String key, int limit, long windowSize, long currentTime) {
        // 计算时间窗口的起始时间
        long windowStart = currentTime - windowSize;
        // 移除时间窗口之前的数据
        jedis.zremrangeByScore(key, "-inf", String.valueOf(windowStart));
        // 获取当前时间窗口内的请求数量
        long currentRequests = jedis.zcount(key, "-inf", "+inf");
        // 判断请求数量是否超过限制
        if (currentRequests < limit) {
            // 添加当前请求的时间戳到有序集合
            jedis.zadd(key, currentTime, String.valueOf(currentTime));
            return true;  // 请求通过限流
        } else {
            return false;  // 请求被限流
        }
    }
}

这里会存在并发问题建议使用lua脚本实现

package com.tuxing.re.annotation;

import redis.clients.jedis.Jedis;

public class SlidingWindowRateLimiter {
    private Jedis jedis;
    private String key;
    private int limit;
    public SlidingWindowRateLimiter(Jedis jedis, String key, int limit) {
        this.jedis = jedis;
        this.key = key;
        this.limit = limit;
    }
    public boolean isAllowed(String key) {
        // 当前时间
        long currentTime = System.currentTimeMillis();
        // 使用lua脚本保证原子性
        String script = "local window_start = ARGV[1] - 30000\n" +
                "redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', window_start)\n" +
                "local current_requests = redis.call('ZCARD', KEYS[1])\n" +
                "if current_requests < tonumber(ARGV[2]) then\n" +
                "    redis.call('ZADD', KEYS[1], ARGV[1], ARGV[1])\n" +
                "    return 1\n" +
                "else\n" +
                "    return 0\n" +
                "end";
        Object result = jedis.eval(script, 1, key, String.valueOf(currentTime), String.valueOf(limit));

        return (int) result == 1;
    }
}

使用示例:

package com.tuxing.re.annotation;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;

/**
 * @author ln
 * @version 1.0
 * @className RateLimitFilter
 * @description 基础描述:
 * @date 2024/5/31 10:49
 */
@Component
public class RateLimitFilter implements Filter {

    public static final Integer LIMIT = 100;// 限制次数

    public static final Integer WINDOW_SIZE_SECONDS = 30;// 窗口大小
    @Autowired
    private SlidingWindowRateLimiter slidingWindowRateLimiter;

    @Resource(name = "againRedisTemplate")
    private RedisTemplate<String, Object> redisTemplate;


    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        HttpServletRequest httpRequest = (HttpServletRequest) request;
        HttpServletResponse httpResponse = (HttpServletResponse) response;
        String apiKey = httpRequest.getHeader("X-API-KEY");
        boolean allowed = slidingWindowRateLimiter.isAllowed(apiKey);
        if (allowed) {
            httpResponse.setStatus(HttpStatus.TOO_MANY_REQUESTS.value());
            httpResponse.getWriter().write("too many request");
        }
    }
}