This is an automated email from the ASF dual-hosted git repository.
Aias00 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git
The following commit(s) were added to refs/heads/master by this push:
new 707fb7d348 [type:fix] give the token limiter counter its window only
once (#7166)
707fb7d348 is described below
commit 707fb7d3486810170aff948ff3774ca782a01045
Author: HY-love-sleep <[email protected]>
AuthorDate: Thu Sep 24 10:35:45 2026 +0800
[type:fix] give the token limiter counter its window only once (#7166)
* fix: give the token limiter counter its window only once (#6649)
recordTokensUsage re-issued the expiration after every increment, so as
long as requests kept
arriving within the window the counter never expired and the token budget
of the window was never
reset: "N tokens per window" behaved as a lifetime cap.
The counter now gets its TTL when it is created (SET NX with a timeout) and
the increment only adds
the usage, so an existing counter keeps its remaining window and expires on
schedule. The TTL is
attached by the command that creates the key, which also avoids a counter
without a window if the
process dies in between.
* fix: write the token counter as a String and keep its window (#6649)
Review follow-up on #7166.
- the counter is created with the String "0": the template the handler
builds uses StringRedisSerializer
for values, so a Long reached the serializer's bridge method and threw
ClassCastException at runtime,
on the first request that records token usage
- the raw ReactiveRedisTemplate that let that compile is gone: the plugin
declares
ReactiveRedisTemplate<String, String>, so a wrong value type no longer
compiles
- a counter that exists without any expiration (written by an earlier
version or by another path) is
given one from getExpire, once, instead of on every increment
- tests: typed mocks, the three branches (created, existing with a window,
existing without one) and a
case that pushes the value through the real stringSerializationContext()
value pair
---------
Co-authored-by: aias00 <[email protected]>
---
.../ai/token/limiter/AiTokenLimiterPlugin.java | 26 +++---
.../ai/token/limiter/AiTokenLimiterPluginTest.java | 93 ++++++++++++++++++++++
2 files changed, 110 insertions(+), 9 deletions(-)
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
index 0094a14313..4b9398d1a9 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/main/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPlugin.java
@@ -38,6 +38,7 @@ import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.data.redis.core.ReactiveRedisTemplate;
+import org.springframework.data.redis.core.ReactiveValueOperations;
import org.springframework.http.HttpCookie;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
@@ -82,7 +83,7 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
return chain.execute(exchange);
}
- ReactiveRedisTemplate reactiveRedisTemplate =
AiTokenLimiterPluginHandler.REDIS_CACHED_HANDLE.get().obtainHandle(PluginEnum.AI_TOKEN_LIMITER.getName());
+ ReactiveRedisTemplate<String, String> reactiveRedisTemplate =
AiTokenLimiterPluginHandler.REDIS_CACHED_HANDLE.get().obtainHandle(PluginEnum.AI_TOKEN_LIMITER.getName());
Assert.notNull(reactiveRedisTemplate, "reactiveRedisTemplate is null");
// generate redis key - include rule id to scope counters per rule
@@ -127,12 +128,12 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
* @param tokenLimit the token limit for the request
* @return whether the request is allowed
*/
- private Mono<Boolean> isAllowed(final ReactiveRedisTemplate
reactiveRedisTemplate, final String cacheKey, final Long tokenLimit) {
+ private Mono<Boolean> isAllowed(final ReactiveRedisTemplate<String,
String> reactiveRedisTemplate, final String cacheKey, final Long tokenLimit) {
return reactiveRedisTemplate.opsForValue().get(cacheKey)
- .defaultIfEmpty(0L)
+ .defaultIfEmpty("0")
.flatMap(currentTokens -> {
- if (Long.parseLong(currentTokens.toString()) >=
tokenLimit) {
+ if (Long.parseLong(currentTokens) >= tokenLimit) {
return Mono.just(false);
}
return Mono.just(true);
@@ -168,11 +169,18 @@ public class AiTokenLimiterPlugin extends
AbstractShenyuPlugin {
return StringUtils.isBlank(key) ? "" : key;
}
- private void recordTokensUsage(final ReactiveRedisTemplate
reactiveRedisTemplate, final String cacheKey, final Long tokens, final Long
windowSeconds) {
- // Record token usage with expiration
- reactiveRedisTemplate.opsForValue()
- .increment(cacheKey, tokens)
- .flatMap(currentValue ->
reactiveRedisTemplate.expire(cacheKey, Duration.ofSeconds(windowSeconds)))
+ private void recordTokensUsage(final ReactiveRedisTemplate<String, String>
reactiveRedisTemplate, final String cacheKey, final Long tokens, final Long
windowSeconds) {
+ // The counter is given its window when it is created: re-issuing the
expiration after every increment would
+ // push the window forward, so a sustained traffic would never reset
the token budget. An existing counter
+ // keeps the window it was created with; only one without any
expiration (written by an earlier version or
+ // by another path) is given one, and just once.
+ final Duration window = Duration.ofSeconds(windowSeconds);
+ final ReactiveValueOperations<String, String> valueOperations =
reactiveRedisTemplate.opsForValue();
+ valueOperations.setIfAbsent(cacheKey, "0", window)
+ .flatMap(created -> created ? Mono.just(Boolean.TRUE) :
reactiveRedisTemplate.getExpire(cacheKey)
+ .filter(timeToLive -> timeToLive.isNegative() ||
timeToLive.isZero())
+ .flatMap(timeToLive ->
reactiveRedisTemplate.expire(cacheKey, window)))
+ .then(valueOperations.increment(cacheKey, tokens))
.subscribe();
}
diff --git
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
index 3142a60497..277427a091 100644
---
a/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
+++
b/shenyu-plugin/shenyu-plugin-ai/shenyu-plugin-ai-token-limiter/src/test/java/org/apache/shenyu/plugin/ai/token/limiter/AiTokenLimiterPluginTest.java
@@ -19,14 +19,22 @@ package org.apache.shenyu.plugin.ai.token.limiter;
import org.apache.shenyu.common.dto.RuleData;
import org.apache.shenyu.common.dto.convert.rule.AiTokenLimiterHandle;
+import
org.apache.shenyu.infra.redis.serializer.ShenyuRedisSerializationContext;
import
org.apache.shenyu.plugin.ai.token.limiter.handler.AiTokenLimiterPluginHandler;
import org.apache.shenyu.plugin.base.utils.CacheKeyUtils;
import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+import org.springframework.data.redis.core.ReactiveRedisTemplate;
+import org.springframework.data.redis.core.ReactiveValueOperations;
+import org.springframework.data.redis.serializer.RedisSerializationContext;
+import reactor.core.publisher.Mono;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
+import java.lang.reflect.Method;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
+import java.time.Duration;
import java.util.zip.GZIPOutputStream;
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -167,6 +175,91 @@ class AiTokenLimiterPluginTest {
assertTrue(decompressed.toString().contains("completion_tokens\":75"));
}
+ @Test
+ void testRecordTokensUsageGivesTheCounterItsWindowWhenItIsCreated() throws
Exception {
+ // Regression for #6649: the window is applied when the counter is
created and never re-issued, otherwise a
+ // sustained traffic keeps pushing it forward and the token budget of
the window is never reset.
+ ReactiveRedisTemplate<String, String> redisTemplate =
mockRedisTemplate();
+ ReactiveValueOperations<String, String> valueOperations =
mockValueOperations();
+ Mockito.when(redisTemplate.opsForValue()).thenReturn(valueOperations);
+ Mockito.when(valueOperations.setIfAbsent("key", "0",
Duration.ofSeconds(60L))).thenReturn(Mono.just(true));
+ Mockito.when(valueOperations.increment("key",
10L)).thenReturn(Mono.just(30L));
+
+ recordTokensUsage(new AiTokenLimiterPlugin(), redisTemplate, "key",
10L, 60L);
+
+ // the counter is created together with its window, and its value goes
through the template's String value
+ // serializer, so it has to be a String (a Long would throw at runtime)
+ Mockito.verify(valueOperations).setIfAbsent("key", "0",
Duration.ofSeconds(60L));
+ Mockito.verify(valueOperations).increment("key", 10L);
+ Mockito.verify(redisTemplate,
Mockito.never()).getExpire(Mockito.anyString());
+ Mockito.verify(redisTemplate,
Mockito.never()).expire(Mockito.anyString(), Mockito.any(Duration.class));
+ }
+
+ @Test
+ void testRecordTokensUsageKeepsTheWindowOfAnExistingCounter() throws
Exception {
+ // An existing counter that already has a window keeps it: re-issuing
the expiration would make the window
+ // slide forward on every request, which is the bug this change fixes.
+ ReactiveRedisTemplate<String, String> redisTemplate =
mockRedisTemplate();
+ ReactiveValueOperations<String, String> valueOperations =
mockValueOperations();
+ Mockito.when(redisTemplate.opsForValue()).thenReturn(valueOperations);
+ Mockito.when(valueOperations.setIfAbsent("key", "0",
Duration.ofSeconds(60L))).thenReturn(Mono.just(false));
+
Mockito.when(redisTemplate.getExpire("key")).thenReturn(Mono.just(Duration.ofSeconds(30L)));
+ Mockito.when(valueOperations.increment("key",
10L)).thenReturn(Mono.just(30L));
+
+ recordTokensUsage(new AiTokenLimiterPlugin(), redisTemplate, "key",
10L, 60L);
+
+ Mockito.verify(redisTemplate).getExpire("key");
+ Mockito.verify(redisTemplate,
Mockito.never()).expire(Mockito.anyString(), Mockito.any(Duration.class));
+ Mockito.verify(valueOperations).increment("key", 10L);
+ }
+
+ @Test
+ void testRecordTokensUsageGivesAWindowToACounterThatHasNone() throws
Exception {
+ // A counter written by an earlier version, or by another path, can
exist without an expiration: it is given
+ // one once, instead of on every increment.
+ ReactiveRedisTemplate<String, String> redisTemplate =
mockRedisTemplate();
+ ReactiveValueOperations<String, String> valueOperations =
mockValueOperations();
+ Mockito.when(redisTemplate.opsForValue()).thenReturn(valueOperations);
+ Mockito.when(valueOperations.setIfAbsent("key", "0",
Duration.ofSeconds(60L))).thenReturn(Mono.just(false));
+
Mockito.when(redisTemplate.getExpire("key")).thenReturn(Mono.just(Duration.ofSeconds(-1L)));
+ Mockito.when(redisTemplate.expire("key",
Duration.ofSeconds(60L))).thenReturn(Mono.just(true));
+ Mockito.when(valueOperations.increment("key",
10L)).thenReturn(Mono.just(30L));
+
+ recordTokensUsage(new AiTokenLimiterPlugin(), redisTemplate, "key",
10L, 60L);
+
+ Mockito.verify(redisTemplate).expire("key", Duration.ofSeconds(60L));
+ Mockito.verify(valueOperations).increment("key", 10L);
+ }
+
+ @Test
+ void testTheCounterValueSurvivesTheTemplateValueSerializer() {
+ // The template built by the handler serializes values with
StringRedisSerializer, so the counter is written
+ // as a String: the serializer's bridge method casts to String, and a
Long would throw at runtime. This is the
+ // class of mistake a mock-only test cannot catch.
+ RedisSerializationContext.SerializationPair<String> valuePair =
+
ShenyuRedisSerializationContext.stringSerializationContext().getValueSerializationPair();
+
+ assertEquals("0", valuePair.read(valuePair.write("0")));
+ }
+
+ @SuppressWarnings("unchecked")
+ private ReactiveValueOperations<String, String> mockValueOperations() {
+ return Mockito.mock(ReactiveValueOperations.class);
+ }
+
+ @SuppressWarnings("unchecked")
+ private ReactiveRedisTemplate<String, String> mockRedisTemplate() {
+ return Mockito.mock(ReactiveRedisTemplate.class);
+ }
+
+ private void recordTokensUsage(final AiTokenLimiterPlugin plugin, final
ReactiveRedisTemplate<String, String> redisTemplate,
+ final String cacheKey, final Long tokens,
final Long windowSeconds) throws Exception {
+ Method method =
AiTokenLimiterPlugin.class.getDeclaredMethod("recordTokensUsage",
+ ReactiveRedisTemplate.class, String.class, Long.class,
Long.class);
+ method.setAccessible(true);
+ method.invoke(plugin, redisTemplate, cacheKey, tokens, windowSeconds);
+ }
+
private String buildSseContentWithTokens(final int tokens) {
return "data:
{\"id\":\"chatcmpl-1\",\"choices\":[{\"delta\":{\"content\":\"Hello\"}}]}\n\n"
+ "data:
{\"id\":\"chatcmpl-1\",\"choices\":[{\"delta\":{\"content\":\" World\"}}]}\n\n"
