限流注解:limit:请求的次数 timeout:时间 以秒为单位 @Target({ ElementType.METHOD }) @Retention(RetentionPolicy.RUNTIME) @Documented public @interface RateLimiter { int limit() default 5; int timeout() default 1000; } 限流拦截器: public class RateLimiterInterceptor implements Interceptor { private String cacheName; public RateLimiterInterceptor() { } public RateLimiterInterceptor(String cacheName) { this.cacheName = cacheName; } @Override public void intercept(Invocation inv) { Controller c = inv.getController(); RateLimiter rateLimiter = inv.getMethod().getAnnotation(RateLimiter.class); if (rateLimiter != null) { int limit = rateLimiter.limit(); int timeout = rateLimiter.timeout(); Cache cache = Redis.use(cacheName); Jedis jedis = cache.getJedis(); String token = RedisRateLimite.acquireTokenFromBucket(jedis, limit, timeout); if (token == null) { try { c.getResponse().sendError(500); } catch (IOException e) { e.printStackTrace(); } return; } cache.close(jedis); inv.invoke(); } else { inv.invoke(); } } } 限流工具类:
public class RedisRateLimite { private static final String BUCKET = "RATELIMITE"; private static final String BUCKET_COUNT = "RATELIMITE_COUNT"; private static final String BUCKET_MONITOR = "RATELIMITE_MONITOR"; static String acquireTokenFromBucket(Jedis jedis, int limit, long timeout) { String identifier = UUID.randomUUID().toString(); long now = System.currentTimeMillis(); Transaction transaction = jedis.multi(); // 删除信号量 transaction.zremrangeByScore(BUCKET_MONITOR.getBytes(), "-inf".getBytes(), String.valueOf(now - timeout).getBytes()); ZParams params = new ZParams(); params.weightsByDouble(1.0, 0.0); transaction.zinterstore(BUCKET, params, BUCKET, BUCKET_MONITOR); // 计数器自增 transaction.incr(BUCKET_COUNT); List<Object> results = transaction.exec(); long counter = (Long) results.get(results.size() - 1); transaction = jedis.multi(); transaction.zadd(BUCKET_MONITOR, now, identifier); transaction.zadd(BUCKET, counter, identifier); transaction.zrank(BUCKET, identifier); results = transaction.exec(); // 获取排名,判断请求是否取得了信号量 long rank = (Long) results.get(results.size() - 1); if (rank < limit) { return identifier; } else {// 没有获取到信号量,清理之前放入redis 中垃圾数据 transaction = jedis.multi(); transaction.zrem(BUCKET_MONITOR, identifier); transaction.zrem(BUCKET, identifier); transaction.exec(); } return null; } }
然后在 config中添加限流拦截器并且指定Redis的cacheName就行了。
方法中使用:
@Controller @RequestMapping(value = "/test") public class TestController extends BaseController { @RateLimiter(limit = 1, timeout = 1000) public void test1() { renderJson(Okv.create().set("msg", "ok")); } }