基于Redis有序集合實現滑動窗口限流

滑動窗口算法是一種基于時間窗口的限流算法，它將時間劃分為若干個固定大小的窗口，每個窗口內記錄了該時間段內的請求次數。通過動態地滑動窗口，可以動態調整限流的速率，以應對不同的流量變化。

整個限流可以概括為兩個主要步驟：

統計窗口內的請求數量
應用限流規則

Redis有序集合每個value有一個score（分數），基于score我們可以定義一個時間窗口，然后每次一個請求進來就設置一個value，這樣就可以統計窗口內的請求數量。key可以是資源名，比如一個url，或者ip+url，用戶標識+url等。value在這里不那么重要，因為我們只需要統計數量，因此value可以就設置成時間戳，但是如果value相同的話就會被覆蓋，所以我們可以把請求的數據做一個hash，將這個hash值當value，或者如果每個請求有流水號的話，可以用請求流水號當value，總之就是要能唯一標識一次請求的。

所以，簡化后的命令就變成了：

ZADD  資源標識   時間戳   請求標識

Java代碼

public boolean isAllow(String key) {
    ZSetOperations<String, String> zSetOperations = stringRedisTemplate.opsForZSet();
    //  獲取當前時間戳
    long currentTime = System.currentTimeMillis();
    //  當前時間 - 窗口大小 = 窗口開始時間
    long windowStart = currentTime - period;
    //  刪除窗口開始時間之前的所有數據
    zSetOperations.removeRangeByScore(key, 0, windowStart);
    //  統計窗口中請求數量
    Long count = zSetOperations.zCard(key);
    //  如果窗口中已經請求的數量超過閾值，則直接拒絕
    if (count >= threshold) {
        return false;
    }
    //  沒有超過閾值，則加入集合
    String value = "請求唯一標識（比如：請求流水號、哈希值、MD5值等）";
    zSetOperations.add(key, String.valueOf(currentTime), currentTime);
    //  設置一個過期時間，及時清理冷數據
    stringRedisTemplate.expire(key, period, TimeUnit.MILLISECONDS);
    //  通過
    return true;
}

上面代碼中涉及到三條Redis命令，并發請求下可能存在問題，所以我們把它們寫成Lua腳本

local key = KEYS[1]
local current_time = tonumber(ARGV[1])
local window_size = tonumber(ARGV[2])
local threshold = tonumber(ARGV[3])
redis.call('ZREMRANGEBYSCORE', key, 0, current_time - window_size)
local count = redis.call('ZCARD', key)
if count >= threshold then
    return tostring(0)
else
    redis.call('ZADD', key, tostring(current_time), current_time)
    return tostring(1)
end

完整的代碼如下：

package com.example.demo.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.stereotype.Service;

import java.util.Collections;
import java.util.concurrent.TimeUnit;

/**
 * 基于Redis有序集合實現滑動窗口限流
 * @Author: ChengJianSheng
 * @Date: 2024/12/26
 */
@Service
public class SlidingWindowRatelimiter {

    private long period = 60*1000;  //  1分鐘
    private int threshold = 3;      //  3次

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

    /**
     * RedisTemplate
     */
    public boolean isAllow(String key) {
        ZSetOperations<String, String> zSetOperations = stringRedisTemplate.opsForZSet();
        //  獲取當前時間戳
        long currentTime = System.currentTimeMillis();
        //  當前時間 - 窗口大小 = 窗口開始時間
        long windowStart = currentTime - period;
        //  刪除窗口開始時間之前的所有數據
        zSetOperations.removeRangeByScore(key, 0, windowStart);
        //  統計窗口中請求數量
        Long count = zSetOperations.zCard(key);
        //  如果窗口中已經請求的數量超過閾值，則直接拒絕
        if (count >= threshold) {
            return false;
        }
        //  沒有超過閾值，則加入集合
        String value = "請求唯一標識（比如：請求流水號、哈希值、MD5值等）";
        zSetOperations.add(key, String.valueOf(currentTime), currentTime);
        //  設置一個過期時間，及時清理冷數據
        stringRedisTemplate.expire(key, period, TimeUnit.MILLISECONDS);
        //  通過
        return true;
    }

    /**
     * Lua腳本
     */
    public boolean isAllow2(String key) {
        String luaScript = "local key = KEYS[1]\n" +
                "local current_time = tonumber(ARGV[1])\n" +
                "local window_size = tonumber(ARGV[2])\n" +
                "local threshold = tonumber(ARGV[3])\n" +
                "redis.call('ZREMRANGEBYSCORE', key, 0, current_time - window_size)\n" +
                "local count = redis.call('ZCARD', key)\n" +
                "if count >= threshold then\n" +
                "    return tostring(0)\n" +
                "else\n" +
                "    redis.call('ZADD', key, tostring(current_time), current_time)\n" +
                "    return tostring(1)\n" +
                "end";

        long currentTime = System.currentTimeMillis();

        DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(luaScript, String.class);

        String result = stringRedisTemplate.execute(redisScript, Collections.singletonList(key), String.valueOf(currentTime), String.valueOf(period), String.valueOf(threshold));
        //  返回1表示通過，返回0表示拒絕
        return "1".equals(result);
    }
}

這里用StringRedisTemplate執行Lua腳本，先把Lua腳本封裝成DefaultRedisScript對象。注意，千萬注意，Lua腳本的返回值必須是字符串，參數也最好都是字符串，用整型的話可能類型轉換錯誤。

String requestId = UUID.randomUUID().toString();

DefaultRedisScript<String> redisScript = new DefaultRedisScript<>(luaScript, String.class);

String result = stringRedisTemplate.execute(redisScript,
        Collections.singletonList(key),
        requestId,
        String.valueOf(period),
        String.valueOf(threshold));

好了，上面就是基于Redis有序集合實現的滑動窗口限流。順帶提一句，Redis List類型也可以用來實現滑動窗口。

接下來，我們來完善一下上面的代碼，通過AOP來攔截請求達到限流的目的

為此，我們必須自定義注解，然后根據注解參數，來個性化的控制限流。那么，問題來了，如果獲取注解參數呢？

舉例說明：

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.METHOD)
public @interface MyAnnotation {
    String value();
}


@Aspect
@Component
public class MyAspect {

    @Before("@annotation(myAnnotation)")
    public void beforeMethod(JoinPoint joinPoint, MyAnnotation myAnnotation) {
        // 獲取注解參數
        String value = myAnnotation.value();
        System.out.println("Annotation value: " + value);

        // 其他業務邏輯...
    }
}

注意看，切點是怎么寫的 @Before("@annotation(myAnnotation)")

是@Before("@annotation(myAnnotation)")，而不是@Before("@annotation(MyAnnotation)")

myAnnotation，是參數，而MyAnnotation則是注解類

此處參考

http://www.rzrgm.cn/javaxubo/p/16556924.html

https://blog.csdn.net/qq_40977118/article/details/119488358

https://blog.51cto.com/knifeedge/5529885

言歸正傳，我們首先定義一個注解

package com.example.demo.controller;

import java.lang.annotation.*;

/**
 * 請求速率限制
 * @Author: ChengJianSheng
 * @Date: 2024/12/26
 */
@Documented
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimit {
    /**
     * 窗口大小（默認：60秒）
     */
    long period() default 60;

    /**
     * 閾值（默認：3次）
     */
    long threshold() default 3;
}

定義切面

package com.example.demo.controller;

import jakarta.servlet.http.HttpServletRequest;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.ZSetOperations;
import org.springframework.stereotype.Component;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;
import org.springframework.web.servlet.support.RequestContextUtils;

import java.util.concurrent.TimeUnit;

/**
 * @Author: ChengJianSheng
 * @Date: 2024/12/26
 */
@Slf4j
@Aspect
@Component
public class RateLimitAspect {

    @Autowired
    private StringRedisTemplate stringRedisTemplate;

//    @Autowired
//    private SlidingWindowRatelimiter slidingWindowRatelimiter;

    @Before("@annotation(rateLimit)")
    public void doBefore(JoinPoint joinPoint, RateLimit rateLimit) {
        //  獲取注解參數
        long period = rateLimit.period();
        long threshold = rateLimit.threshold();

        //  獲取請求信息
        ServletRequestAttributes servletRequestAttributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
        HttpServletRequest httpServletRequest = servletRequestAttributes.getRequest();
        String uri = httpServletRequest.getRequestURI();
        Long userId = 123L;     //  模擬獲取用戶ID
        String key = "limit:" + userId + ":" + uri;
        /*
        if (!slidingWindowRatelimiter.isAllow2(key)) {
            log.warn("請求超過速率限制！userId={}, uri={}", userId, uri);
            throw new RuntimeException("請求過于頻繁！");
        }*/

        ZSetOperations<String, String> zSetOperations = stringRedisTemplate.opsForZSet();
        //  獲取當前時間戳
        long currentTime = System.currentTimeMillis();
        //  當前時間 - 窗口大小 = 窗口開始時間
        long windowStart = currentTime - period * 1000;
        //  刪除窗口開始時間之前的所有數據
        zSetOperations.removeRangeByScore(key, 0, windowStart);
        //  統計窗口中請求數量
        Long count = zSetOperations.zCard(key);
        //  如果窗口中已經請求的數量超過閾值，則直接拒絕
        if (count < threshold) {
            //  沒有超過閾值，則加入集合
            zSetOperations.add(key, String.valueOf(currentTime), currentTime);
            //  設置一個過期時間，及時清理冷數據
            stringRedisTemplate.expire(key, period, TimeUnit.SECONDS);
        } else {
            throw new RuntimeException("請求過于頻繁！");
        }
    }

}

加注解

@RestController
@RequestMapping("/hello")
public class HelloController {

    @RateLimit(period = 30, threshold = 2)
    @GetMapping("/sayHi")
    public void sayHi() {

    }
}

最后，看Redis中的數據結構

最后的最后，流量控制建議看看阿里巴巴 Sentinel

https://sentinelguard.io/zh-cn/

posted @ 2024-12-30 18:11 廢物大師兄閱讀(1716) 評論(0) 收藏舉報

刷新頁面返回頂部

狂亂的貴公子

人生南北多歧路，君向瀟湘我向秦。此后人間多寬敞，南來北往不遇卿。

基于Redis有序集合實現滑動窗口限流

公告