SpringBoot基于Redis和Lua构建分布式限流防御体系教学

-- keys[1..n]          n 个限流 key
-- argv[1..n]          n 个窗口大小（秒）
-- argv[n+1..2n]       n 个阈值
-- argv[2n+1]          当前时间戳（毫秒）
-- argv[2n+2]          请求唯一前缀（用于生成 member）
for i=1, #keys do
    local window = tonumber(argv[i])
    local threshold = tonumber(argv[i + #keys])
    local current = tonumber(argv[2 * #keys + 1])
    local expired = current - window * 1000
    redis.call('zremrangebyscore', keys[i], 0, expired)
    if redis.call('zcard', keys[i]) >= threshold then
        return 0
    end
end
for i=1, #keys do
    local window = tonumber(argv[i])
    local current = tonumber(argv[2 * #keys + 1])
    local member = argv[2 * #keys + 2] .. '-' .. i
    redis.call('zadd', keys[i], current, member)
    redis.call('expire', keys[i], window)
end
return 1

@target(elementtype.method)
@retention(retentionpolicy.runtime)
@repeatable(ratelimits.class)
public @interface ratelimit {
    /** 限流 key，支持 spel 表达式 */
    string key();

    /** 窗口大小（秒） */
    long window() default 60;

    /** 窗口内最大请求数 */
    long threshold() default 100;

    /** 降级方法名，必须在同类中定义，返回值与原方法一致 */
    string fallback() default "";
}

@target(elementtype.method)
@retention(retentionpolicy.runtime)
public @interface ratelimits {
    ratelimit[] value();
}

@bean
public defaultredisscript<long> ratelimitscript() {
    defaultredisscript<long> script = new defaultredisscript<>();
    script.setscriptsource(new resourcescriptsource(
        new classpathresource("scripts/sliding_window.lua")));
    script.setresulttype(long.class);
    return script;
}

@aspect
@component
public class ratelimitaspect {

    private final stringredistemplate redistemplate;
    private final defaultredisscript<long> ratelimitscript;
    private final spelexpressionparser parser = new spelexpressionparser();
    private final defaultparameternamediscoverer parameternamediscoverer = new defaultparameternamediscoverer();

    public ratelimitaspect(stringredistemplate redistemplate,
                           defaultredisscript<long> ratelimitscript) {
        this.redistemplate = redistemplate;
        this.ratelimitscript = ratelimitscript;
    }

    @around("@annotation(io.ratelimit.ratelimit)")
    public object around(proceedingjoinpoint joinpoint) throws throwable {
        methodsignature signature = (methodsignature) joinpoint.getsignature();
        method method = signature.getmethod();

        ratelimit[] limits = annotatedelementutils.getmergedrepeatableannotations(
                method, ratelimit.class, ratelimits.class);
        if (limits.length == 0) {
            return joinpoint.proceed();
        }

        list<string> keys = new arraylist<>();
        list<long> windows = new arraylist<>();
        list<long> thresholds = new arraylist<>();

        for (ratelimit limit : limits) {
            keys.add("rate:limit:" + evalkey(limit.key(), joinpoint));
            windows.add(limit.window());
            thresholds.add(limit.threshold());
        }

        list<string> args = new arraylist<>();
        windows.foreach(w -> args.add(string.valueof(w)));
        thresholds.foreach(t -> args.add(string.valueof(t)));
        args.add(string.valueof(system.currenttimemillis()));
        args.add(uuid.randomuuid().tostring());

        long result = redistemplate.execute(
            ratelimitscript,
            keys,
            args.toarray(new string[0])
        );

        if (result != null && result == 1l) {
            return joinpoint.proceed();
        }

        ratelimit failed = limits[0];
        if (!failed.fallback().isempty()) {
            return invokefallback(joinpoint, failed.fallback());
        }
        throw new ratelimitexception("too many requests");
    }

    private string evalkey(string expression, proceedingjoinpoint joinpoint) {
        methodsignature signature = (methodsignature) joinpoint.getsignature();
        object[] args = joinpoint.getargs();
        string[] paramnames = parameternamediscoverer.getparameternames(signature.getmethod());

        standardevaluationcontext context = new standardevaluationcontext();
        for (int i = 0; i < paramnames.length; i++) {
            context.setvariable(paramnames[i], args[i]);
        }
        return parser.parseexpression(expression).getvalue(context, string.class);
    }

    private object invokefallback(proceedingjoinpoint joinpoint, string fallback) throws throwable {
        methodsignature signature = (methodsignature) joinpoint.getsignature();
        method targetmethod = joinpoint.gettarget().getclass()
                .getdeclaredmethod(fallback, signature.getparametertypes());
        targetmethod.setaccessible(true);
        return targetmethod.invoke(joinpoint.gettarget(), joinpoint.getargs());
    }
}

@restcontroller
public class seckillcontroller {

    @ratelimit(key = "'seckill:api'", window = 1, threshold = 1000, fallback = "apifallback")
    @ratelimit(key = "'seckill:user:' + #userid", window = 60, threshold = 3, fallback = "userfallback")
    @ratelimit(key = "'seckill:ip:' + #ip", window = 60, threshold = 5, fallback = "ipfallback")
    @postmapping("/seckill")
    public order doseckill(long userid, string ip) {
        // 秒杀业务逻辑
        return new order();
    }

    public order apifallback(long userid, string ip) {
        return order.failed("系统繁忙");
    }

    public order userfallback(long userid, string ip) {
        return order.failed("请勿重复提交");
    }

    public order ipfallback(long userid, string ip) {
        return order.failed("请求过于频繁");
    }
}

public void afterseckillsuccess(long userid, long goodsid) {
    stringredistemplate.executepipelined((rediscallback<object>) connection -> {
        connection.incr("seckill:success:count".getbytes());
        connection.setex("seckill:success:" + userid, 300, "1".getbytes());
        connection.hset("seckill:order".getbytes(), userid.tostring().getbytes(),
                goodsid.tostring().getbytes());
        return null;
    });
}