import java.lang.annotation.*;

/**
 * 功能:分布式接口限流注解
 * @author love ice
 * @create 2023-09-18 15:43
 */
@Target({ElementType.TYPE,ElementType.METHOD})
@Documented
@Retention(RetentionPolicy.RUNTIME)
public @interface RedisLimit {
    /**
     * redis中唯一key,一般用方法名字做區(qū)分
     * 作用: 針對不同接口，做不同的限流控制
     */
    String key() default "";

    /**
     * 限流時(shí)間內(nèi)允許訪問次數(shù) 默認(rèn)1
     */
    long permitsPerSecond() default 1;

    /**
     * 限流時(shí)間,單位秒 默認(rèn)60秒
     */
    long expire() default 60;

    /**
     * 限流提示信息
     */
    String msg() default "接口限流,請稍后重試";
}

import com.aliyuncs.utils.StringUtils;
import com.test.redis.Infrastructure.annotation.RedisLimit;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.JoinPoint;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;
import javax.annotation.PostConstruct;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;

/**
 * redis限流切面
 *
 * @author love ice
 * @create 2023-09-18 15:44
 */
@Slf4j
@Aspect
@Component
public class RedisLimitAop {
    @Autowired
    private StringRedisTemplate stringRedisTemplate;
    private DefaultRedisScript<Long> redisScript;

    @PostConstruct
    public void init() {
        redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        // 執(zhí)行 lua 腳本
        ResourceScriptSource resourceScriptSource = new ResourceScriptSource(new ClassPathResource("rateLimiter.lua"));
        redisScript.setScriptSource(resourceScriptSource);
    }

    @Pointcut("@annotation(com.test.redis.Infrastructure.annotation.RedisLimit)")
    private void check() {

    }

    @Before("check()")
    private void before(JoinPoint joinPoint) {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        // 拿到 RedisLimit 注解,如果存在則說明需要限流
        RedisLimit redisLimit = method.getAnnotation(RedisLimit.class);

        if (redisLimit != null) {
            // 獲取 redis 的 key
            String key = redisLimit.key();
            String className = method.getDeclaringClass().getName();
            String name = method.getName();
            String limitKey = key + className + name;
            log.info("限流的key:{}", limitKey);

            if (StringUtils.isEmpty(key)) {
                // 這里是自定義異常,為了方便寫成了 RuntimeException
                throw new RuntimeException("code:101,msg:接口中 key 參數(shù)不能為空");
            }

            long limit = redisLimit.permitsPerSecond();
            long expire = redisLimit.expire();
            // 把 key 放入 List 中
            List<String> keys = new ArrayList<>(Collections.singletonList(key));
            Long count = stringRedisTemplate.execute(redisScript, keys, String.valueOf(limit), String.valueOf(expire));
            log.info("Access try count is {} for key ={}", count, keys);

            if (count != null && count == 0) {
                log.debug("令牌桶={}, 獲取令牌失效,接口觸發(fā)限流", key);
                throw new RuntimeException("code:10X, redisLimit.msg()");
            }
        }
    }
}

--獲取KEY
local key = KEYS[1]

local limit = tonumber(ARGV[1])

local curentLimit = tonumber(redis.call('get', key) or "0")

if curentLimit + 1 > limit
    then return 0
else
    -- 自增長 1
    redis.call('INCRBY', key, 1)
    -- 設(shè)置過期時(shí)間
    redis.call('EXPIRE', key, ARGV[2])
    return curentLimit + 1
end