diff --git a/src/main/java/io/jboot/components/limiter/LimitScope.java b/src/main/java/io/jboot/components/limiter/LimitScope.java new file mode 100644 index 0000000000000000000000000000000000000000..7e4dd25f16452c0ae569b03aa3cc04ac1f21e238 --- /dev/null +++ b/src/main/java/io/jboot/components/limiter/LimitScope.java @@ -0,0 +1,15 @@ +package io.jboot.components.limiter; + +public enum LimitScope { + + /** + * 整个集群限次,多实例共享 + */ + CLUSTER, + + /** + * 每个实例单独限次 + */ + NODE; + +} diff --git a/src/main/java/io/jboot/components/limiter/annotation/EnableLimit.java b/src/main/java/io/jboot/components/limiter/annotation/EnableLimit.java index 582cf8a0857ccbf9377163d8a87453d146f858b2..248069a138995928f3a684bcfa0e4100ab2a4098 100644 --- a/src/main/java/io/jboot/components/limiter/annotation/EnableLimit.java +++ b/src/main/java/io/jboot/components/limiter/annotation/EnableLimit.java @@ -15,6 +15,7 @@ */ package io.jboot.components.limiter.annotation; +import io.jboot.components.limiter.LimitScope; import io.jboot.components.limiter.LimitType; import java.lang.annotation.*; @@ -38,6 +39,12 @@ public @interface EnableLimit { */ String type() default LimitType.TOKEN_BUCKET; + + /** + * 作用域,默认为单节点本地限次 + */ + LimitScope scope() default LimitScope.NODE; + /** * 频率 * diff --git a/src/main/java/io/jboot/components/limiter/interceptor/BaseLimiterInterceptor.java b/src/main/java/io/jboot/components/limiter/interceptor/BaseLimiterInterceptor.java index 39fb08777f2bc044e984bb9a6ddffec30511ad4c..2ab9177d5047f0ec5ccfa151556ab6d1914f4331 100644 --- a/src/main/java/io/jboot/components/limiter/interceptor/BaseLimiterInterceptor.java +++ b/src/main/java/io/jboot/components/limiter/interceptor/BaseLimiterInterceptor.java @@ -19,6 +19,7 @@ package io.jboot.components.limiter.interceptor; import com.google.common.util.concurrent.RateLimiter; import com.jfinal.aop.Invocation; import io.jboot.components.limiter.LimiterManager; +import io.jboot.components.limiter.redis.RedisRateLimitUtil; import io.jboot.utils.ClassUtil; import io.jboot.utils.StrUtil; @@ -60,6 +61,17 @@ public abstract class BaseLimiterInterceptor { } } + protected void doInterceptForTokenBucketWithCluster(int rate, String resource, String fallback, Invocation inv) { + //允许通行 + if (RedisRateLimitUtil.tryAcquire(resource, rate)) { + inv.invoke(); + } + //不允许通行 + else { + doExecFallback(resource, fallback, inv); + } + } + protected void doExecFallback(String resource, String fallback, Invocation inv) { LimiterManager.me().processFallback(resource, fallback, inv); } diff --git a/src/main/java/io/jboot/components/limiter/interceptor/LimiterInterceptor.java b/src/main/java/io/jboot/components/limiter/interceptor/LimiterInterceptor.java index c400a53dfb11a1894d7c7c2e5f511187d5b70bf6..58defb320ee748eb34399b007e9d444cc7a0879f 100644 --- a/src/main/java/io/jboot/components/limiter/interceptor/LimiterInterceptor.java +++ b/src/main/java/io/jboot/components/limiter/interceptor/LimiterInterceptor.java @@ -18,6 +18,7 @@ package io.jboot.components.limiter.interceptor; import com.jfinal.aop.Interceptor; import com.jfinal.aop.Invocation; +import io.jboot.components.limiter.LimitScope; import io.jboot.components.limiter.LimitType; import io.jboot.components.limiter.annotation.EnableLimit; import io.jboot.utils.AnnotationUtil; @@ -40,10 +41,17 @@ public class LimiterInterceptor extends BaseLimiterInterceptor implements Interc String type = AnnotationUtil.get(enableLimit.type()); switch (type) { case LimitType.CONCURRENCY: + if (LimitScope.CLUSTER == enableLimit.scope()) { + throw new IllegalArgumentException("Concurrency limit for cluster not implement!"); + } doInterceptForConcurrency(enableLimit.rate(), resource, enableLimit.fallback(), inv); break; case LimitType.TOKEN_BUCKET: - doInterceptForTokenBucket(enableLimit.rate(), resource, enableLimit.fallback(), inv); + if (LimitScope.CLUSTER == enableLimit.scope()) { + doInterceptForTokenBucketWithCluster(enableLimit.rate(), resource, enableLimit.fallback(), inv); + } else { + doInterceptForTokenBucket(enableLimit.rate(), resource, enableLimit.fallback(), inv); + } break; } } diff --git a/src/main/java/io/jboot/components/limiter/redis/RedisRateLimitUtil.java b/src/main/java/io/jboot/components/limiter/redis/RedisRateLimitUtil.java new file mode 100644 index 0000000000000000000000000000000000000000..2152ce7d4b1293e95284492782101a54853e1da4 --- /dev/null +++ b/src/main/java/io/jboot/components/limiter/redis/RedisRateLimitUtil.java @@ -0,0 +1,50 @@ +package io.jboot.components.limiter.redis; + +import io.jboot.support.redis.JbootRedis; +import io.jboot.support.redis.JbootRedisManager; + +/** + * 通过lua脚本来进行限次 + */ +public class RedisRateLimitUtil { + + private static final String RATE_LIMIT_SCRIPT = "local c" + + "\nc = redis.call('get',KEYS[1])" + + // 调用量已经超过最大值,直接返回 + "\nif c and tonumber(c) > tonumber(ARGV[1]) then" + + "\nreturn tonumber(c);" + + "\nend" + + // 自增 + "\nc = redis.call('incr',KEYS[1])" + + "\nif tonumber(c) == 1 then" + + // 从第一次调用开始限流,设置对应键值的过期 + "\nredis.call('expire',KEYS[1],ARGV[2])" + + "\nend" + + "\nreturn c;"; + + private static JbootRedis redis; + + /** + * 限制时长默认为1秒 + */ + public static boolean tryAcquire(String resource, int rate) { + return tryAcquire(resource, rate, 1); + } + + /** + * 尝试是否能正常执行 + * + * @param resource 资源名 + * @param rate 限制次数 + * @param periodSeconds 限制时长,单位为秒 + * @return true 可以执行 + * false 限次,禁止 + */ + public static boolean tryAcquire(String resource, int rate, int periodSeconds) { + if (redis == null) { + redis = JbootRedisManager.me().getRedis(); + } + Long count = (Long) redis.eval(RATE_LIMIT_SCRIPT, 1, resource, String.valueOf(rate), String.valueOf(periodSeconds)); + return count <= rate; + } +} diff --git a/src/main/java/io/jboot/support/redis/JbootRedis.java b/src/main/java/io/jboot/support/redis/JbootRedis.java index 9d9f14b9da8230a1062bcb3f0dd991bb23e0ec75..a6ec93f106d9e323302e5fe2c7a7f485754f3fab 100644 --- a/src/main/java/io/jboot/support/redis/JbootRedis.java +++ b/src/main/java/io/jboot/support/redis/JbootRedis.java @@ -659,6 +659,7 @@ public interface JbootRedis { public List valueListFromBytesList(Collection data); + Object eval(String script, int keyCount, String... params); } diff --git a/src/main/java/io/jboot/support/redis/jedis/JbootJedisClusterImpl.java b/src/main/java/io/jboot/support/redis/jedis/JbootJedisClusterImpl.java index c43fc4bb7468ef67cbda6bb0fcc4a84c52f95eb5..f68fb6c9fba1ad802c0a15d7880781705cf09bed 100644 --- a/src/main/java/io/jboot/support/redis/jedis/JbootJedisClusterImpl.java +++ b/src/main/java/io/jboot/support/redis/jedis/JbootJedisClusterImpl.java @@ -52,8 +52,8 @@ public class JbootJedisClusterImpl extends JbootRedisBase { if (timeout != null) { this.timeout = timeout; } - if(maxAttempts == null) { - maxAttempts = this.maxAttempts; + if (maxAttempts == null) { + maxAttempts = this.maxAttempts; } GenericObjectPoolConfig poolConfig = new GenericObjectPoolConfig(); @@ -1227,6 +1227,10 @@ public class JbootJedisClusterImpl extends JbootRedisBase { return new RedisScanResult<>(scanResult.getStringCursor(), scanResult.getResult()); } + @Override + public Object eval(String script, int keyCount, String... params) { + return jedisCluster.eval(script, keyCount, params); + } public JedisCluster getJedisCluster() { return jedisCluster; diff --git a/src/main/java/io/jboot/support/redis/jedis/JbootJedisImpl.java b/src/main/java/io/jboot/support/redis/jedis/JbootJedisImpl.java index 9cf5a4a1699887c1ebab650f63fdddd8b3be7f99..a7181b9807ef62bfaa07e0832d7b66c31420d3e5 100644 --- a/src/main/java/io/jboot/support/redis/jedis/JbootJedisImpl.java +++ b/src/main/java/io/jboot/support/redis/jedis/JbootJedisImpl.java @@ -1563,6 +1563,15 @@ public class JbootJedisImpl extends JbootRedisBase { } } + @Override + public Object eval(String script, int keyCount, String... params) { + Jedis jedis = getJedis(); + try { + return jedis.eval(script, keyCount, params); + } finally { + returnResource(jedis); + } + } public Jedis getJedis() { try { diff --git a/src/main/java/io/jboot/support/redis/lettuce/JbootLettuceImpl.java b/src/main/java/io/jboot/support/redis/lettuce/JbootLettuceImpl.java index 830ad87a6518cefc8f8e8419a08fc955de6e5cd3..653c336f4c771c9e2311f045de71dd345e6f9d61 100644 --- a/src/main/java/io/jboot/support/redis/lettuce/JbootLettuceImpl.java +++ b/src/main/java/io/jboot/support/redis/lettuce/JbootLettuceImpl.java @@ -525,4 +525,9 @@ public class JbootLettuceImpl implements JbootRedis { public List valueListFromBytesList(Collection data) { return null; } + + @Override + public Object eval(String script, int keyCount, String... params) { + return null; + } } diff --git a/src/main/java/io/jboot/support/redis/redisson/JbootRedissonImpl.java b/src/main/java/io/jboot/support/redis/redisson/JbootRedissonImpl.java index 7c20d105e6a8aea76e7018dcbf4ac4384c78dbf1..d76259040f8e5d458c442d1661b664ed6012ecd8 100644 --- a/src/main/java/io/jboot/support/redis/redisson/JbootRedissonImpl.java +++ b/src/main/java/io/jboot/support/redis/redisson/JbootRedissonImpl.java @@ -514,4 +514,9 @@ public class JbootRedissonImpl implements JbootRedis { public List valueListFromBytesList(Collection data) { return null; } + + @Override + public Object eval(String script, int keyCount, String... params) { + return null; + } } diff --git a/src/test/java/io/jboot/test/aop/inject/AopController.java b/src/test/java/io/jboot/test/aop/inject/AopController.java index 5026f25691598b12790c828aeb5775d48aede764..4a234bae0fbaeed45af9030557a33e8075a54972 100644 --- a/src/test/java/io/jboot/test/aop/inject/AopController.java +++ b/src/test/java/io/jboot/test/aop/inject/AopController.java @@ -2,6 +2,7 @@ package io.jboot.test.aop.inject; import com.jfinal.aop.Inject; import io.jboot.aop.annotation.ConfigValue; +import io.jboot.components.limiter.LimitScope; import io.jboot.components.limiter.annotation.EnableLimit; import io.jboot.test.aop.staticconstruct.StaticConstructManager; import io.jboot.web.controller.JbootController; @@ -36,6 +37,11 @@ public class AopController extends JbootController { renderText("host:" + host + " port:" + port + " xxx:" + xxx); } + @EnableLimit(rate = 1, fallback = "aaa", scope = LimitScope.CLUSTER) + public void bbb() { + renderText("host:" + host + " port:" + port + " xxx:" + xxx); + } + public void aaa(){ renderText("aaa"); } diff --git a/src/test/java/io/jboot/test/redis/RedisTester.java b/src/test/java/io/jboot/test/redis/RedisTester.java index a591373824a8783d4439eecf107d604a8da8c056..45bb80db81c168e528a88d1e60b81226e6e57e30 100644 --- a/src/test/java/io/jboot/test/redis/RedisTester.java +++ b/src/test/java/io/jboot/test/redis/RedisTester.java @@ -1,15 +1,49 @@ package io.jboot.test.redis; import io.jboot.app.JbootApplication; +import io.jboot.components.limiter.redis.RedisRateLimitUtil; import io.jboot.support.redis.JbootRedis; import io.jboot.support.redis.JbootRedisManager; import io.jboot.support.redis.RedisScanResult; +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; import java.util.ArrayList; import java.util.List; public class RedisTester { + @Before + public void config() { + JbootApplication.setBootArg("jboot.redis.host", "127.0.0.1"); + JbootApplication.setBootArg("jboot.redis.port", "6379"); + } + + @Test + public void testGetAndSet() { + JbootRedis redis = JbootRedisManager.me().getRedis(); + String key = "JbootRedisValue"; + Assert.assertEquals("OK", redis.set(key, "10")); + Assert.assertEquals("10", redis.get(key)); + redis.del(key); + } + + @Test + public void testEval() { + JbootRedis redis = JbootRedisManager.me().getRedis(); + String response = (String) redis.eval("return KEYS[1]", 1, "key1"); + Assert.assertEquals("key1", response); + } + + @Test + public void testRateLimit() { + String resource = "limited-resource"; + Assert.assertTrue(RedisRateLimitUtil.tryAcquire(resource, 2, 1)); + Assert.assertTrue(RedisRateLimitUtil.tryAcquire(resource, 2, 1)); + Assert.assertFalse(RedisRateLimitUtil.tryAcquire(resource, 2, 1)); + } + public static void main(String[] args) { JbootApplication.setBootArg("jboot.redis.host", "127.0.0.1"); JbootApplication.setBootArg("jboot.redis.database", "3");