API接口限流竟然如此簡單
作者：了不起
我們已經成功實現了一個基于Redisson和Spring AOP的API接口限流功能。這個方案不僅簡單易懂，而且非常靈活，可以通過注解輕松地應用到任意方法上，并且支持多種限流策略（如全局限流、IP限流、集群限流等）。
簡介
API接口限流是一種流量控制技術，其目的是通過設置規則來限制客戶端對API接口的調用速率或總量，從而避免因過載而導致的服務性能下降甚至崩潰。
API限流在各種系統上都會有廣泛的使用場景，本文介紹一種非常簡單的實現API限流的方式。
為什么需要API接口限流？
- 防止惡意攻擊：通過限制請求速率，可以有效抵御DDoS等類型的攻擊。
- 優化資源使用：合理分配有限的計算資源給所有用戶，避免單個用戶占用過多資源。
- 提升服務質量：保持服務響應時間在一個合理的范圍內，提高整體用戶體驗。
令牌桶
常見的API限流策略有令牌桶等算法。
令牌桶算法是一種常用的流量控制和限流機制，它通過模擬一個存放“令牌”的桶來控制請求的速率。
這個算法的核心思想是：系統以恒定的速率向桶中添加令牌，而每個請求在被處理之前必須從桶中獲取一個令牌。如果桶中有足夠的令牌，則請求可以繼續執行；如果沒有足夠的令牌（即桶為空），則請求要么等待直到有新的令牌產生，要么直接被拒絕。
實現API限流
這個算法很容易理解，但是要想手動實現一個令牌桶算法，并不是一個容易的事情。
還需要考慮：時間精度、并發處理、存儲管理、可配置性等問題。
Redis是一個常用的非關系型數據庫，非常適合用于緩存、實現限流等功能。本文介紹一個利用redis非常簡單的實現限流的功能，采用 AOP + 注解 + Redisson 框架實現。
1.定義限流注解
@Target(ElementType.METHOD)
@Retention(RetentionPolicy.RUNTIME)
@Documented
public @interface RateLimiter {
/**
* 限流key,支持使用Spring el表達式來動態獲取方法上的參數值
* 格式類似于 #code.id #{#code}
*/
String key() default "";
/**
* 限流時間,單位秒
*/
int time() default 60;
/**
* 限流次數
*/
int count() default 100;
/**
* 限流類型
*/
LimitType limitType() default LimitType.DEFAULT;
/**
* 提示消息
*/
String message() default "服務器暫無資源處理新的請求，請稍后重試";
}public enum LimitType {
/**
* 默認策略全局限流
*/
DEFAULT,
/**
* 根據請求者IP進行限流
*/
IP,
/**
* 實例限流(集群多后端實例)
*/
CLUSTER
}2.注解切面
@Slf4j
@Aspect
@Order(1)
public class RateLimiterAspect {
private static final String LIMITER_KEY = "global:limiter:";
/**
* 定義spel表達式解析器
*/
private final ExpressionParser parser = new SpelExpressionParser();
/**
* 定義spel解析模版
*/
private final ParserContext parserContext = new TemplateParserContext();
/**
* 方法參數解析器
*/
private final ParameterNameDiscoverer pnd = new DefaultParameterNameDiscoverer();
/**
* \@within(rateLimiter) 和 \@annotation(rateLimiter) 必須按照這個順序，才會優先執行方法上的注解
*/
@Before("@within(rateLimiter) || @annotation(rateLimiter)")
public void doBefore(JoinPoint point, RateLimiter rateLimiter) {
if (rateLimiter == null) {
// 如果方法上沒有，就從類上獲取注解
Class<?> targetClass = point.getTarget().getClass();
rateLimiter = targetClass.getAnnotation(RateLimiter.class);
if (rateLimiter == null) {
// 如果還是沒有獲取到注解，直接返回
return;
}
}
int time = rateLimiter.time();
int count = rateLimiter.count();
try {
String combineKey = getCombineKey(rateLimiter, point);
RateType rateType = RateType.OVERALL;
if (rateLimiter.limitType() == LimitType.CLUSTER) {
rateType = RateType.PER_CLIENT;
}
long number = RedisUtils.rateLimiter(combineKey, rateType, count, time);
if (number == -1) {
throw new RateLimiterException(rateLimiter.message());
}
log.debug("限制令牌 => {}, 剩余令牌 => {}, 緩存key => '{}'", count, number, combineKey);
} catch (Exception e) {
if (e instanceof RateLimiterException) {
throw e;
} else {
throw new RuntimeException("服務器限流異常，請稍候再試", e);
}
}
}
private String getCombineKey(RateLimiter rateLimiter, JoinPoint point) {
String key = rateLimiter.key();
// 判斷 key 不為空 和 不是表達式
if (StringUtils.hasText(key) && key.contains("#")) {
MethodSignature signature = (MethodSignature) point.getSignature();
Method targetMethod = signature.getMethod();
Object[] args = point.getArgs();
MethodBasedEvaluationContext context =
new MethodBasedEvaluationContext(null, targetMethod, args, pnd);
context.setBeanResolver(new BeanFactoryResolver(SpringUtil.getBeanFactory()));
Expression expression;
if (key.startsWith(parserContext.getExpressionPrefix()) && key.endsWith(parserContext.getExpressionSuffix())) {
expression = parser.parseExpression(key, parserContext);
} else {
expression = parser.parseExpression(key);
}
key = expression.getValue(context, String.class);
}
StringBuilder str = new StringBuilder(LIMITER_KEY);
HttpServletRequest request = getRequest();
str.append(request.getRequestURI()).append(":");
if (rateLimiter.limitType() == LimitType.IP) {
// 獲取請求ip
str.append(ServletUtil.getClientIP(request)).append(":");
} else if (rateLimiter.limitType() == LimitType.CLUSTER) {
// 獲取客戶端實例id
str.append(RedisUtils.getClient().getId()).append(":");
}
return str.append(key).toString();
}
/**
* 獲取request
*/
private HttpServletRequest getRequest() {
try {
ServletRequestAttributes attributes = (ServletRequestAttributes) RequestContextHolder.getRequestAttributes();
return attributes.getRequest();
} catch (Exception e) {
return null;
}
}
}3.Redisson 限流工具類
public class RedisUtils {
private static final RedissonClient CLIENT = SpringUtil.getBean(RedissonClient.class);
/**
* 限流
*
* @param key 限流key
* @param rateType 限流類型
* @param rate 速率
* @param rateInterval 速率間隔
* @return -1 表示失敗
*/
public static long rateLimiter(String key, RateType rateType, int rate, int rateInterval) {
RRateLimiter rateLimiter = CLIENT.getRateLimiter(key);
// 如果限流器存在
if (rateLimiter.isExists()) {
// 獲取上次限流的配置信息
RateLimiterConfig rateLimiterConfig = rateLimiter.getConfig();
// 如果rateLimiterConfig的配置跟我們注解上面的值不一致，說明服務器重啟過，程序員又修改了限流的配置
if (TimeUnit.SECONDS.convert(rateLimiterConfig.getRateInterval(), TimeUnit.MILLISECONDS) != rateInterval || rateLimiterConfig.getRate() != rate) {
rateLimiter.delete();
rateLimiter.trySetRate(rateType, rate, rateInterval, RateIntervalUnit.SECONDS);
}
}
rateLimiter.trySetRate(rateType, rate, rateInterval, RateIntervalUnit.SECONDS);
if (rateLimiter.tryAcquire()) {
return rateLimiter.availablePermits();
} else {
return -1L;
}
}
/**
* 獲取客戶端實例
*/
public static RedissonClient getClient() {
return CLIENT;
}
}4.捕獲異常
@Data
@EqualsAndHashCode(callSuper = true)
public class RateLimiterException extends RuntimeException {
/**
* 錯誤提示
*/
private final String message;
public RateLimiterException(String message) {
this.message = message;
}
}@Slf4j
@Order(1)
@RestControllerAdvice
public class LimiterExceptionHandler {
/**
* 限流異常
*/
@ExceptionHandler({RateLimiterException.class})
public Map<String, Object> handleRateLimiterException(RateLimiterException e, HttpServletRequest request) {
log.error("請求地址'{}', 限流異常'{}'", request.getRequestURI(), e.getMessage());
return result(e.getMessage());
}
private Map<String, Object> result(String msg) {
HashMap<String, Object> result = new HashMap<>();
result.put("code", 500);
result.put("msg", msg);
return result;
}
}到這里，已經實現了一個完整的API接口限流功能。
可以將之進一步封裝，作為一個springboot的starter，用于任意一個項目中。
小結
通過上述步驟，我們已經成功實現了一個基于Redisson和Spring AOP的API接口限流功能。這個方案不僅簡單易懂，而且非常靈活，可以通過注解輕松地應用到任意方法上，并且支持多種限流策略（如全局限流、IP限流、集群限流等）。
責任編輯：武曉燕
來源：
Java技術指北
