利用RedissonClient实现访问量控制

题目：

代码部署在多台主机，需要控制访问量。

原理：

创建注解和切面程序，利用RedissonClient实现流量控制并返回指定错误码。

实现：

pom文件：

        <dependency>
            <groupId>org.redisson</groupId>
            <artifactId>redisson-spring-boot-starter</artifactId>
            <version>3.14.1</version>
        </dependency>
        <dependency>
            <groupId>org.aspectj</groupId>
            <artifactId>aspectjrt</artifactId>
            <version>1.9.5</version>
        </dependency>
        <dependency>
            <groupId>org.aspectj</groupId>
            <artifactId>aspectjweaver</artifactId>
            <version>1.9.5</version>
        </dependency>

创建注解：

 

package com.example.demo.annotation;

import com.example.demo.common.api.ResultCode;
import org.redisson.api.RateIntervalUnit;

import java.lang.annotation.*;

/**
 * 切面限流注解
 */
@Retention(RetentionPolicy.RUNTIME)
@Target({ElementType.TYPE,ElementType.METHOD})
@Documented
public @interface DistributedRateLimit {
    /**
     * 生成令牌数
     */
    int rate();

    /**
     * 生成令牌时间
     */
    int rateInterval();

    /**
     * 生成令牌时间单位
     */
    RateIntervalUnit rateIntervalUnit() default RateIntervalUnit.MINUTES;

    /**
     * 令牌桶的key
     */
    String key() default "";

    /**
     * 错误码（enum）
     */
    ResultCode value() default ResultCode.DISTRIBUTED_LIMIT;

    /**
     * 是否是sse流式输出
     */
    boolean s() default true;
}

创建切面程序： 

package com.example.demo.aspect;

import com.example.demo.annotation.DistributedRateLimit;
import com.example.demo.common.api.CommonResult;
import lombok.AllArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.redisson.api.RRateLimiter;
import org.redisson.api.RateType;
import org.redisson.api.RedissonClient;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

/**
 * 利用redis实现多实例下的访问量的限流
 */

@Aspect
@Component
@Slf4j
@AllArgsConstructor
public class DistributedLockAspectConfiguration {
    private final RedissonClient redissonClient;


    @Pointcut("@annotation(com.example.demo.annotation.DistributedRateLimit)")
    public void distributedLimitPointcut(){}

    @Around("distributedLimitPointcut()")
    public Object distributedLimitPointInterceptor(ProceedingJoinPoint pjp) throws  Throwable{
        log.info("distributedLimitPointInterceptor start:{}"+System.currentTimeMillis());
        MethodSignature signature = (MethodSignature) pjp.getSignature();
        Method method = signature.getMethod();
        DistributedRateLimit distributedRateLimit = method.getAnnotation(DistributedRateLimit.class);
        String key = method.getDeclaringClass().getName()+method.getName();
        final RRateLimiter rateLimiter = redissonClient.getRateLimiter(key);
        //设置流速
        if(!rateLimiter.isExists()){
            rateLimiter.trySetRate(RateType.OVERALL,distributedRateLimit.rate(),distributedRateLimit.rateInterval(),
                    distributedRateLimit.rateIntervalUnit());
        }
        rateLimiter.expire(60, TimeUnit.MINUTES);
        final boolean acquireSuccess = rateLimiter.tryAcquire();
        if(acquireSuccess){
            return pjp.proceed();
        }
        log.info("distributedLimitPointInterceptor entry limit:{}"+System.currentTimeMillis());
        if(distributedRateLimit.s()){
            SseEmitter sseEmitter = new SseEmitter();
            sseEmitter.send(SseEmitter.event().data(CommonResult.failed(distributedRateLimit.value())));
            sseEmitter.complete();
            return sseEmitter;
        }
        return CommonResult.failed(distributedRateLimit.value());

    }



}

注意要加@Component注解和加好pom文件的配置项，否则切面不生效。

配置到接口：

这里是配置到流式输出的接口上了，设置的是30分钟产生一个令牌，也就是30分钟内只能有一个访问量，便于测试展示。

package com.example.demo.controller;

import com.example.demo.annotation.DistributedRateLimit;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Controller;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;


/**
 * sse流水调用
 */
@RequestMapping(value = "/sse")
@Controller
@Slf4j
public class SseController {

    @DistributedRateLimit(rate = 1,rateInterval = 30)
    @PostMapping(value = "/getAnswer")
    public SseEmitter getAnswer(String inputParameter){
        SseEmitter sseEmitter = new SseEmitter();
        new Thread(() -> {
            try {
                // Query the database based on the input parameter and send data in batches
                for (int i = 0; i < 10; i++) {
                    String data = "Data batch " + i + " for parameter: " + inputParameter;
                    sseEmitter.send(data);
                    Thread.sleep(1000); // Simulate delay between batches
                }

                sseEmitter.complete(); // Complete the SSE connection
            } catch (Exception e) {
                sseEmitter.completeWithError(e); // Handle errors
            }

        }).start();
        return sseEmitter;
    }
}

效果演示：