io.github.dengchen2020.ratelimiter.local.LocalRateLimiterAop Maven / Gradle / Ivy

Show more of this group Show more artifacts with this name
Show all versions of dc-spring-boot-starter-ratelimiter Show documentation
限流
There is a newer version: 0.0.28
Show newest version
package io.github.dengchen2020.ratelimiter.local;

import io.github.dengchen2020.core.utils.AnnotationUtils;
import io.github.dengchen2020.core.utils.IPUtils;
import io.github.dengchen2020.core.utils.RequestUtils;
import io.github.dengchen2020.ratelimiter.annotation.RateLimit;
import io.github.dengchen2020.ratelimiter.annotation.RateLimitStrategy;
import io.github.dengchen2020.ratelimiter.exception.RateLimitException;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.util.ConcurrentReferenceHashMap;

import java.lang.reflect.Method;
import java.security.Principal;
import java.time.Duration;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.TimeUnit;

/**
 * 单机限流实现
 *
 * @author dengchen
 * @since 2024/4/18
 */
@Aspect
public class LocalRateLimiterAop {

    RateLimiter minuteRateLimiter = new RateLimiter(Duration.ofMinutes(1));

    RateLimiter secondRateLimiter = new RateLimiter(Duration.ofSeconds(1));

    private static final ConcurrentMap ipRateLimiterMap = new ConcurrentReferenceHashMap<>();

    @Around(value = "@within(rateLimit) || @annotation(rateLimit)")
    public Object handle(ProceedingJoinPoint joinPoint, RateLimit rateLimit) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        return limit(ipRateLimiterMap.computeIfAbsent(method, key -> AnnotationUtils.annotation(key, RateLimit.class)), signature.toString(), joinPoint);
    }

    private Object limit(RateLimit rateLimit, String methodSignature, ProceedingJoinPoint joinPoint) throws Throwable {
        RateLimitStrategy strategy = rateLimit.strategy();
        String limitKey;
        if (strategy == RateLimitStrategy.userAndMethod) {
            Principal principal = RequestUtils.getCurrRequest().getUserPrincipal();
            if (principal == null) {
                limitKey = IPUtils.getIpAddr() + methodSignature;
            } else {
                limitKey = principal.getName() + methodSignature;
            }
        } else if (strategy == RateLimitStrategy.ip) {
            limitKey = IPUtils.getIpAddr();
        } else if (strategy == RateLimitStrategy.ipAndMethod) {
            limitKey = IPUtils.getIpAddr() + methodSignature;
        } else if (strategy == RateLimitStrategy.user) {
            Principal principal = RequestUtils.getCurrRequest().getUserPrincipal();
            if (principal == null) {
                limitKey = IPUtils.getIpAddr();
            } else {
                limitKey = principal.getName();
            }
        } else if (strategy == RateLimitStrategy.method) {
            limitKey = methodSignature;
        } else {
            limitKey = IPUtils.getIpAddr() + methodSignature;
        }
        RateLimiter rateLimiter;
        if (rateLimit.timeUnit() == TimeUnit.MINUTES) {
            rateLimiter = minuteRateLimiter;
        } else {
            rateLimiter = secondRateLimiter;
        }
        if (rateLimiter.limit(limitKey, rateLimit.value())) {
            throw new RateLimitException(rateLimit.errorMsg());
        }
        try {
            return joinPoint.proceed();
        } catch (Throwable e) {
            if (!rateLimit.errorLimit()) {
                rateLimiter.quashLimit(limitKey);
            }
            throw e;
        }
    }

}