使用策略模式实现 Spring 分布式和单机限流

咱们能够运用战略形式来统一单机限流和分布式限流的完成,提高代码的可扩展性和可维护性。

思路是界说一个 RateLimitStrategy 接口,并别离完成单机限流战略 LocalRateLimitStrategy 和分布式限流战略 DistributedRateLimitStrategy。在 AOP 切面中,依据装备决议运用哪种限流战略。

界说战略接口

public interface RateLimitStrategy {
    boolean tryAcquire(String key, double qps, long timeout, TimeUnit timeUnit);
}

完成单机限流战略

import com.google.common.util.concurrent.RateLimiter;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
public class LocalRateLimitStrategy implements RateLimitStrategy {
    private final Map<String, RateLimiter> rateLimiters = new ConcurrentHashMap<>();
    @Override
    public boolean tryAcquire(String key, double qps, long timeout, TimeUnit timeUnit) {
        RateLimiter limiter = rateLimiters.computeIfAbsent(key, k -> RateLimiter.create(qps));
        if (timeout > 0) {
            return limiter.tryAcquire(timeout, timeUnit);
        } else {
            return limiter.tryAcquire();
        }
    }
}

完成分布式限流战略

import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.TimeUnit;
public class DistributedRateLimitStrategy implements RateLimitStrategy {
    private final RedisTemplate<String, Object> redisTemplate;
    public DistributedRateLimitStrategy(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }
    @Override
    public boolean tryAcquire(String key, double qps, long timeout, TimeUnit timeUnit) {
        long window = timeUnit.toSeconds(timeout);
        List<String> keys = Collections.singletonList(key);
        String luaScript = buildLuaScript();
        RedisScript<Long> redisScript = new DefaultRedisScript<>(luaScript, Long.class);
        Long currentCount = redisTemplate.execute(redisScript, keys, Collections.singletonList(window), Collections.singletonList(qps));
        return currentCount <= qps;
    }
    private String buildLuaScript() {
       return "local key = KEYS[1]n" +
                "local window = tonumber(ARGV[1])n" +
                "local qps = tonumber(ARGV[2])n" +
                "local current = redis.call('incrBy', key, 1)n" +
                "if current == 1 thenn" +
                "    redis.call('expire', key, window)n" +
                "endn" +
                "if current > qps thenn" +
                "    return redis.call('decrBy', key, 1)n" +
                "elsen" +
                "    return currentn" +
                "end";
    }
}

修正切面逻辑

import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.concurrent.TimeUnit;
@Aspect
@Component
public class RateLimitAspect {
    @Autowired
    private RateLimitStrategy rateLimitStrategy;
    @Around("@annotation(rateLimitAnnotation)")
    public Object around(ProceedingJoinPoint joinPoint, RateLimit rateLimitAnnotation) throws Throwable {
        String key = joinPoint.getSignature().toLongString();
        double qps = rateLimitAnnotation.qps();
        long timeout = rateLimitAnnotation.timeout();
        TimeUnit timeUnit = rateLimitAnnotation.timeUnit();
        boolean acquired = rateLimitStrategy.tryAcquire(key, qps, timeout, timeUnit);
        if (!acquired) {
            throw new RuntimeException("Rate limit exceeded");
        }
        return joinPoint.proceed();
    }
}

在切面逻辑中,咱们注入了 RateLimitStrategy 的完成类。依据装备决议运用单机限流还是分布式限流战略。

运用示例

@RestController
public class DemoController {
    @Autowired
    private RateLimitStrategy rateLimitStrategy;
    @GetMapping("/test")
    @ApiRateLimit(qps = 10, timeout = 60, timeUnit = TimeUnit.SECONDS)
    public String test() {
        return "hello world";
    }
}

在运用时,咱们只需求在办法上标示 @RateLimit 注解即可,而不需求关心底层运用的是单机限流还是分布式限流。

装备限流战略

在 Spring 装备中,咱们能够依据需求注入不同的 RateLimitStrategy 完成类:

// 单机限流装备
@Bean
public RateLimitStrategy localRateLimitStrategy() {
    return new LocalRateLimitStrategy();
}
// 分布式限流装备
@Bean
public RateLimitStrategy distributedRateLimitStrategy(RedisTemplate<String, Object> redisTemplate) {
    return new DistributedRateLimitStrategy(redisTemplate);
}

通过运用战略形式,咱们将限流算法与详细的限流战略解耦,提高了代码的可扩展性和可维护性。未来假如需求新的限流战略,只需求完成 RateLimitStrategy 接口并装备即可,无需修正核心的限流逻辑。