AOP实现方法访问限制

对于一些方法尤其是对接某个API的通常会有对于请求次数，文件大小等的限制，通过java的面向切面变成实现一个访问限制器
1.使用元注解自定义一个注解，并包含使用所需要的一些属性

@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
public @interface RateLimited {
    int maxTokens(); // 单位时间允许最大请求次数
    long refillInterval(); // 请求刷新时间
    TimeUnit unit(); // 请求刷新时间单位
    long maxFileSize() default -1; // 允许最大的文件大小
    long maxExecutionTime() default -1; // 方法允许的最大执行时间
    int retryTime() default 0;  // 方法重试次数
}

2. 定义一个类，主要保存当前对于该方法请求的状态，以及获得请求令牌和刷新请求令牌的逻辑

public class RateLimiter {
    private final int maxTokens;
    private final long refillInterval;
    private final AtomicInteger tokens;
    private long lastRefillTimestamp;

    public RateLimiter(int maxTokens, long refillInterval, TimeUnit unit) {
        this.maxTokens = maxTokens;
        this.refillInterval = unit.toMillis(refillInterval);
        this.tokens = new AtomicInteger(maxTokens);
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    public synchronized boolean tryAcquire() {
        refill();
        if (tokens.get() > 0) {
            tokens.decrementAndGet();
            return true;
        }
        return false;
    }

    private void refill() {
        long now = System.currentTimeMillis();
        if (now - lastRefillTimestamp > refillInterval) {
            tokens.set(maxTokens);
            lastRefillTimestamp = now;
        } else {
            long interval = refillInterval / maxTokens;
            if (now - lastRefillTimestamp > interval){
                tokens.addAndGet(1);
            }
        }
    }
}

3.定义切面类，对于请求的过滤的逻辑，以及定义一个环绕通知

@Aspect
@Component
public class RateLimiterAspect {
    private final ConcurrentHashMap<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();

    @Around("@annotation(rateLimited)")
    public Object rateLimit(ProceedingJoinPoint joinPoint, RateLimited rateLimited) {
        String key = joinPoint.getSignature().toShortString();
        RateLimiter rateLimiter = rateLimiters.computeIfAbsent(key, k -> new RateLimiter(
                rateLimited.maxTokens(),
                rateLimited.refillInterval(),
                rateLimited.unit()
        ));

        if (!rateLimiter.tryAcquire()) {
            throw new RuntimeException("请求次数过多，请稍后重试.");
        }

        for (Object arg : joinPoint.getArgs()) {
            if (arg instanceof Document) {
                Document document = (Document) arg;
                if (rateLimited.maxFileSize() != -1 && document.getFileSize() > rateLimited.maxFileSize()) {
                    String format = String.format("文件大小超过限制,文件大小%skb,允许最大文件为%skb",document.getFileSize(), rateLimited.maxFileSize());
                    throw new RuntimeException(format);
                }
            }
        }

        int retryTimes = rateLimited.retryTime();
        long maxExecutionTime = rateLimited.maxExecutionTime();
        for (int attempt = 0; attempt <= retryTimes; attempt++) {
            ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
            Future<Object> future = null;
            ScheduledFuture<?> timeoutFuture = null;
            try {
                future = Executors.newSingleThreadExecutor().submit(() -> {
                    try {
                        return joinPoint.proceed();
                    } catch (Throwable e) {
                        throw new RuntimeException(e);
                    }
                });
                if (maxExecutionTime != -1) {
                    Future<Object> finalFuture = future;
                    timeoutFuture = scheduler.schedule(() -> finalFuture.cancel(true), maxExecutionTime, TimeUnit.SECONDS);
                }
                return future.get();
            } catch (Throwable throwable) {
                if (attempt == retryTimes) {
                    throw new RuntimeException("方法执行时间超过限制,方法允许执行最大时长" + maxExecutionTime + "s");
                }
            } finally {
                if (timeoutFuture != null) {
                    timeoutFuture.cancel(true);
                }
                if (future != null) {
                    future.cancel(true);
                }
                scheduler.shutdown();
            }
        }
        throw new RuntimeException("方法执行出错,请联系管理员");
    }
}

使用的时候只需要在指定方法添加注解@RateLimited就可以
 

@RateLimited(maxTokens = 10, refillInterval = 1, unit = TimeUnit.HOURS, maxFileSize = 10 * 1024, maxExecutionTime = 30 * 60)
public void testRateLimited(Object arg) {
    // your code。。。
}

这个例子就是表示每个小时允许十次对于方法的访问，允许文件最大为10MB，方法允许最长执行时间为30分钟，可重试次数为默认的0次