基于SpringBoot AOP+Redis实现接口限流注解的示例代码_Java

限流是保护后端系统的第一道防线。用 aop + redis 实现一个自定义限流注解，加到接口上就能限流。

一、定义注解

@target(elementtype.method)
@retention(retentionpolicy.runtime)
public @interface ratelimit {

    /** 限流 key */
    string key() default "";

    /** 限流维度：user / ip / api */
    string type() default "user";

    /** 窗口内最大请求数 */
    long max() default 10;

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

    /** 被限后的提示 */
    string message() default "请求太频繁，请稍后重试";
}

二、实现 aop

@aspect
@component
public class ratelimitaspect {

    @autowired
    private stringredistemplate redistemplate;

    @around("@annotation(ratelimit)")
    public object around(proceedingjoinpoint pjp, ratelimit ratelimit) throws throwable {
        string key = buildkey(ratelimit);

        // 滑动窗口限流
        boolean allowed = tryacquire(key, ratelimit.max(), ratelimit.window());

        if (!allowed) {
            throw new businessexception(429, ratelimit.message());
        }

        return pjp.proceed();
    }

    private boolean tryacquire(string key, long max, long windowsec) {
        long now = system.currenttimemillis();
        long windowms = windowsec * 1000l;
        long windowstart = now - windowms;

        // lua 脚本保证原子性
        string lua =
            "redis.call('zremrangebyscore', keys[1], 0, argv[1]) " +
            "local count = redis.call('zcard', keys[1]) " +
            "if count < tonumber(argv[2]) then " +
            "  redis.call('zadd', keys[1], argv[3], argv[3]) " +
            "  redis.call('expire', keys[1], argv[4]) " +
            "  return 1 " +
            "else " +
            "  return 0 " +
            "end";

        long result = redistemplate.execute(
            new defaultredisscript<>(lua, long.class),
            collections.singletonlist(key),
            string.valueof(windowstart),
            string.valueof(max),
            string.valueof(now),
            string.valueof(windowsec + 1)
        );

        return long.valueof(1).equals(result);
    }

    private string buildkey(ratelimit ratelimit) {
        string prefix = "rate:";

        switch (ratelimit.type()) {
            case "ip":
                httpservletrequest request = ((servletrequestattributes)
                    requestcontextholder.getrequestattributes()).getrequest();
                return prefix + "ip:" + getip(request);
            case "api":
                return prefix + "api:" + stream.of(
                    thread.currentthread().getstacktrace())
                    .filter(s -> s.getmethodname().contains("$"))
                    .findfirst().orelse(new stacktraceelement("", "", "", 0))
                    .getmethodname();
            default:
                // type = user
                return prefix + "user:" + stputil.getloginidasstring();
        }
    }

    public static string getip(httpservletrequest request) {
        string ip = request.getheader("x-forwarded-for");
        if (ip == null || ip.isempty()) ip = request.getheader("x-real-ip");
        if (ip == null || ip.isempty()) ip = request.getremoteaddr();
        if (ip != null && ip.contains(",")) ip = ip.split(",")[0].trim();
        return ip;
    }
}

三、使用

@restcontroller
@requestmapping("/api")
public class testcontroller {

    @getmapping("/test")
    @ratelimit(key = "test", type = "ip", max = 5, window = 10)
    public resultvo<?> test() {
        return resultvo.success("成功");
    }

    @postmapping("/seckill/{productid}")
    @ratelimit(type = "user", max = 1, window = 10)
    public resultvo<?> seckill(@pathvariable long productid) {
        return resultvo.success("秒杀成功");
    }
}

四、异常处理

@restcontrolleradvice
public class ratelimitexceptionhandler {

    @exceptionhandler(businessexception.class)
    public resultvo<?> handlebusiness(businessexception e) {
        if (e.getmessage().contains("请求太频繁")) {
            return resultvo.error(429, e.getmessage());
        }
        return resultvo.error(e.getcode(), e.getmessage());
    }
}

五、压测效果

# 模拟 20 个并发请求
for i in {1..20}; do
    curl -x get http://localhost:9090/api/test &
done

# 正常响应：{"code":200,"message":"成功"}
# 被限响应：{"code":429,"message":"请求太频繁，请稍后重试"}

到此这篇关于基于springboot aop+redis实现接口限流注解的示例代码的文章就介绍到这了,更多相关springboot aop redis接口限流注解内容请搜索代码网以前的文章或继续浏览下面的相关文章希望大家以后多多支持代码网！