我們可以使用策略模式來統(tǒng)一單機限流和分布式限流的實現(xiàn),提高代碼的可擴展性和可維護性。

思路是定義一個 RateLimitStrategy 接口,并分別實現(xiàn)單機限流策略 LocalRateLimitStrategy 和分布式限流策略 DistributedRateLimitStrategy。在 AOP 切面中,根據(jù)配置決定使用哪種限流策略。

定義策略接口

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

實現(xiàn)單機限流策略

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();
        }
    }
}

實現(xiàn)分布式限流策略

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 then\n" +
                "    redis.call('expire', key, window)\n" +
                "end\n" +
                "if current > qps then\n" +
                "    return redis.call('decrBy', key, 1)\n" +
                "else\n" +
                "    return current\n" +
                "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 的實現(xiàn)類。根據(jù)配置決定使用單機限流還是分布式限流策略。

使用示例

@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 配置中,我們可以根據(jù)需求注入不同的 RateLimitStrategy 實現(xiàn)類:

// 單機限流配置
@Bean
public RateLimitStrategy localRateLimitStrategy() {
    return new LocalRateLimitStrategy();
}

// 分布式限流配置
@Bean
public RateLimitStrategy distributedRateLimitStrategy(RedisTemplate<String, Object> redisTemplate) {
    return new DistributedRateLimitStrategy(redisTemplate);
}

通過使用策略模式,我們將限流算法與具體的限流策略解耦,提高了代碼的可擴展性和可維護性。未來如果需要新的限流策略,只需要實現(xiàn) RateLimitStrategy 接口并配置即可,無需修改核心的限流邏輯。文章來源地址http://www.zghlxwxcb.cn/news/detail-856858.html

到了這里，關于使用策略模式實現(xiàn) Spring 分布式和單機限流的文章就介紹完了。如果您還想了解更多內容，請在右上角搜索TOY模板網(wǎng)以前的文章或繼續(xù)瀏覽下面的相關文章，希望大家以后多多支持TOY模板網(wǎng)！