From 98b6ffda32941a313b97519ed24f7ab83f96e287 Mon Sep 17 00:00:00 2001 From: jizhuozhi Date: Fri, 25 Feb 2022 17:51:34 +0800 Subject: [PATCH] Add tests for request_rate_limiter.lua Fixes gh-2532 --- .../META-INF/scripts/request_rate_limiter.lua | 11 +- .../RedisRateLimiterLuaScriptTests.java | 156 ++++++++++++++++++ 2 files changed, 164 insertions(+), 3 deletions(-) create mode 100644 spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java diff --git a/spring-cloud-gateway-server/src/main/resources/META-INF/scripts/request_rate_limiter.lua b/spring-cloud-gateway-server/src/main/resources/META-INF/scripts/request_rate_limiter.lua index 36a01733..7959d95e 100644 --- a/spring-cloud-gateway-server/src/main/resources/META-INF/scripts/request_rate_limiter.lua +++ b/spring-cloud-gateway-server/src/main/resources/META-INF/scripts/request_rate_limiter.lua @@ -6,11 +6,16 @@ local timestamp_key = KEYS[2] local rate = tonumber(ARGV[1]) local capacity = tonumber(ARGV[2]) -local now = redis.call('TIME')[1] +local now = tonumber(ARGV[3]) local requested = tonumber(ARGV[4]) -local fill_time = capacity/rate -local ttl = math.floor(fill_time*2) +local fill_time = capacity / rate +local ttl = math.floor(fill_time * 2) + +-- for testing, it should use redis system time in production +if now == nil then + now = redis.call('TIME')[1] +end --redis.log(redis.LOG_WARNING, "rate " .. ARGV[1]) --redis.log(redis.LOG_WARNING, "capacity " .. ARGV[2]) diff --git a/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java new file mode 100644 index 00000000..dc36f929 --- /dev/null +++ b/spring-cloud-gateway-server/src/test/java/org/springframework/cloud/gateway/filter/ratelimit/RedisRateLimiterLuaScriptTests.java @@ -0,0 +1,156 @@ +/* + * Copyright 2013-2022 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.cloud.gateway.filter.ratelimit; + +import java.util.Arrays; +import java.util.List; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.GenericContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.test.context.SpringBootTest; +import org.springframework.data.redis.core.ReactiveStringRedisTemplate; +import org.springframework.data.redis.core.script.RedisScript; +import org.springframework.test.annotation.DirtiesContext; +import org.springframework.test.context.DynamicPropertyRegistry; +import org.springframework.test.context.DynamicPropertySource; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Zhuozhi Ji + */ +@SpringBootTest +@DirtiesContext +@Testcontainers +@Tag("DockerRequired") +public class RedisRateLimiterLuaScriptTests { + + static final String KEY_PREFIX = "redis-rate-limiter-lua-script-tests"; + + @Container + public static GenericContainer redis = new GenericContainer<>("redis:5.0.14-alpine").withExposedPorts(6379); + + @Autowired + ReactiveStringRedisTemplate redisTemplate; + + @Autowired + RedisScript> redisScript; + + @DynamicPropertySource + static void containerProperties(DynamicPropertyRegistry registry) { + registry.add("spring.data.redis.host", redis::getContainerIpAddress); + registry.add("spring.data.redis.port", redis::getFirstMappedPort); + } + + static List getKeys(String id) { + String prefix = KEY_PREFIX + ".{" + id; + String tokens = prefix + "}.tokens"; + String timestamp = prefix + "}.timestamp"; + return Arrays.asList(tokens, timestamp); + } + + static List getArgs(long rate, long capacity, long now, long requested) { + return Arrays.asList(rate + "", capacity + "", now + "", requested + ""); + } + + @Test + public void testNewAccess() { + long rate = 1; + long capacity = 10; + long now = System.currentTimeMillis(); + long requested = 1; + List keys = getKeys("new_access"); + List args = getArgs(rate, capacity, now, requested); + List result = redisTemplate.execute(redisScript, keys, args).blockFirst(); + assertThat(result.get(0)).isEqualTo(1); + assertThat(result.get(1)).isEqualTo(9); + + for (String key : keys) { + long ttl = redisTemplate.getExpire(key).map(duration -> duration.getSeconds()).block(); + long fillTime = (capacity / rate); + assertThat(ttl).isGreaterThanOrEqualTo(fillTime); + } + } + + @Test + public void testTokenFilled() { + long rate = 1; + long capacity = 10; + long now = System.currentTimeMillis(); + long requested = 5; + List keys = getKeys("token_filled"); + List args = getArgs(rate, capacity, now, requested); + redisTemplate.execute(redisScript, keys, args).blockFirst(); + + now = now + 3; + args = getArgs(rate, capacity, now, requested); + List result = redisTemplate.execute(redisScript, keys, args).blockFirst(); + assertThat(result.get(0)).isEqualTo(1); + assertThat(result.get(1)).isEqualTo(3); + + for (String key : keys) { + long ttl = redisTemplate.getExpire(key).map(duration -> duration.getSeconds()).block(); + long fillTime = (capacity / rate); + assertThat(ttl).isGreaterThanOrEqualTo(fillTime); + } + } + + @Test + public void testAfterTillTime() { + long rate = 1; + long capacity = 10; + long now = System.currentTimeMillis(); + long requested = 1; + List keys = getKeys("after_fill_time"); + List args = getArgs(rate, capacity, now, requested); + redisTemplate.execute(redisScript, keys, args).blockFirst(); + + long fillTime = capacity / rate; + now = now + fillTime; + args = getArgs(rate, capacity, now, requested); + List result = redisTemplate.execute(redisScript, keys, args).blockFirst(); + assertThat(result.get(0)).isEqualTo(1); + assertThat(result.get(1)).isEqualTo(9); + } + + @Test + public void testTokensNotEnough() { + long rate = 1; + long capacity = 10; + long now = System.currentTimeMillis(); + long requested = 20; + List keys = getKeys("tokens_not_enough"); + List args = getArgs(rate, capacity, now, requested); + List result = redisTemplate.execute(redisScript, keys, args).blockFirst(); + assertThat(result.get(0)).isEqualTo(0); + assertThat(result.get(1)).isEqualTo(10); + } + + @EnableAutoConfiguration + @SpringBootConfiguration + public static class TestConfig { + + } + +}