mirror of
https://github.com/google/nomulus
synced 2026-08-26 00:56:23 +00:00
Split QuotaManager into generic and EppServer classes (#3213)
This separates the logic of "how many tokens should a particular user/group have" from "manage the quota given whatever limits, contacting Valkey". This is in preparation for allowing other types of quota management and throttling besides just on the EPP server (e.g. domain-create throttling). We also convert the expirations from seconds to milliseconds (and use a Duration so that this is masked from users). It's better to have the API use a full-fledged Duration object rather than an int, and this allows for finer control over expiration times. Note that we'll probably want to use a sliding window in the future instead of a fixed window, but that's a problem for future us.
This commit is contained in:
gbrodman
@@ -24,8 +24,9 @@ import google.registry.config.RegistryConfig.Config;
|
||||
import google.registry.config.RegistryConfigSettings;
|
||||
import google.registry.eppserver.Protocol.FrontendProtocol;
|
||||
import google.registry.eppserver.handler.EppServiceHandler;
|
||||
import google.registry.eppserver.quota.QuotaManager;
|
||||
import google.registry.eppserver.quota.EppServerQuotaManager;
|
||||
import google.registry.networking.handler.SslServerInitializer;
|
||||
import google.registry.quota.GenericValkeyQuotaManager;
|
||||
import io.netty.channel.ChannelHandler;
|
||||
import io.netty.channel.socket.nio.NioSocketChannel;
|
||||
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
|
||||
@@ -137,8 +138,9 @@ public final class EppProtocolModule {
|
||||
@Provides
|
||||
@Singleton
|
||||
@CommandQuota
|
||||
static QuotaManager provideCommandQuotaManager(
|
||||
static EppServerQuotaManager provideCommandQuotaManager(
|
||||
@Config("eppServerQuota") RegistryConfigSettings.Quota quota, Optional<UnifiedJedis> jedis) {
|
||||
return new QuotaManager(quota, jedis.orElse(null), "command");
|
||||
return new EppServerQuotaManager(
|
||||
quota, new GenericValkeyQuotaManager(jedis.orElse(null), "command"));
|
||||
}
|
||||
}
|
||||
|
||||
@@ -26,8 +26,8 @@ import com.google.common.flogger.FluentLogger;
|
||||
import google.registry.config.RegistryConfig.Config;
|
||||
import google.registry.eppserver.EppProtocolModule.CommandQuota;
|
||||
import google.registry.eppserver.metric.FrontendMetrics;
|
||||
import google.registry.eppserver.quota.EppServerQuotaManager;
|
||||
import google.registry.eppserver.quota.LocalConnectionLimiter;
|
||||
import google.registry.eppserver.quota.QuotaManager;
|
||||
import google.registry.module.RegistryServlet;
|
||||
import google.registry.request.RequestHandler;
|
||||
import google.registry.util.FakeHttpServletRequest;
|
||||
@@ -76,7 +76,7 @@ public class EppServiceHandler extends SimpleChannelInboundHandler<ByteBuf> {
|
||||
private final byte[] helloBytes;
|
||||
private final FrontendMetrics metrics;
|
||||
private final LocalConnectionLimiter localConnectionLimiter;
|
||||
private final QuotaManager commandQuotaManager;
|
||||
private final EppServerQuotaManager commandQuotaManager;
|
||||
private final Supplier<String> idTokenSupplier;
|
||||
private final String projectId;
|
||||
private final int preLoginReadTimeoutSeconds;
|
||||
@@ -98,7 +98,7 @@ public class EppServiceHandler extends SimpleChannelInboundHandler<ByteBuf> {
|
||||
@Named("hello") byte[] helloBytes,
|
||||
FrontendMetrics metrics,
|
||||
LocalConnectionLimiter localConnectionLimiter,
|
||||
@CommandQuota QuotaManager commandQuotaManager,
|
||||
@CommandQuota EppServerQuotaManager commandQuotaManager,
|
||||
@Named("idToken") Supplier<String> idTokenSupplier,
|
||||
@Config("projectId") String projectId,
|
||||
@Config("eppServerPreLoginReadTimeoutSeconds") int preLoginReadTimeoutSeconds) {
|
||||
@@ -227,7 +227,7 @@ public class EppServiceHandler extends SimpleChannelInboundHandler<ByteBuf> {
|
||||
String throttleId =
|
||||
(authenticatedRegistrarId != null) ? authenticatedRegistrarId : sslClientCertificateHash;
|
||||
if (throttleId != null) {
|
||||
if (!commandQuotaManager.acquireQuota(new QuotaManager.QuotaRequest(throttleId)).success()) {
|
||||
if (!commandQuotaManager.acquireQuota(throttleId)) {
|
||||
metrics.registerQuotaRejection("epp_command", throttleId);
|
||||
closeConnection(ctx);
|
||||
return false;
|
||||
|
||||
@@ -0,0 +1,91 @@
|
||||
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.eppserver.quota;
|
||||
|
||||
import com.google.common.collect.ImmutableMap;
|
||||
import google.registry.config.RegistryConfigSettings.Quota;
|
||||
import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup;
|
||||
import google.registry.quota.GenericValkeyQuotaManager;
|
||||
import java.time.Duration;
|
||||
import javax.annotation.concurrent.ThreadSafe;
|
||||
|
||||
/**
|
||||
* Quota management for the EPP server using Redis/Valkey.
|
||||
*
|
||||
* <p>Handles primarily configuration lookup and delegation to the generic quota manager.
|
||||
*/
|
||||
@ThreadSafe
|
||||
public class EppServerQuotaManager {
|
||||
|
||||
private static final Duration DEFAULT_TTL = Duration.ofHours(1);
|
||||
|
||||
private final GenericValkeyQuotaManager quotaManager;
|
||||
private final QuotaGroup defaultQuota;
|
||||
private final ImmutableMap<String, QuotaGroup> customQuotas;
|
||||
|
||||
public EppServerQuotaManager(Quota quota, GenericValkeyQuotaManager quotaManager) {
|
||||
this.quotaManager = quotaManager;
|
||||
this.defaultQuota = quota.defaultQuota;
|
||||
|
||||
ImmutableMap.Builder<String, QuotaGroup> builder = ImmutableMap.builder();
|
||||
quota.customQuota.forEach(group -> group.userId.forEach(userId -> builder.put(userId, group)));
|
||||
this.customQuotas = builder.build();
|
||||
}
|
||||
|
||||
/** Attempts to acquire a quota token from Redis. */
|
||||
public boolean acquireQuota(String userId) {
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
|
||||
// Unlimited quota check
|
||||
if (group.tokenAmount < 0) {
|
||||
return true;
|
||||
}
|
||||
|
||||
String redisId = getRedisId(group, userId);
|
||||
return quotaManager.acquireQuota(redisId, group.tokenAmount, getTtl(group));
|
||||
}
|
||||
|
||||
/** Refreshes the TTL of an existing quota token. */
|
||||
public void refreshQuota(String userId) {
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
if (group.tokenAmount < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
String redisId = getRedisId(group, userId);
|
||||
quotaManager.refreshQuota(redisId, getTtl(group));
|
||||
}
|
||||
|
||||
/** Returns a token to the pool (used for connection throttling). */
|
||||
public void releaseQuota(String userId) {
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
if (group.tokenAmount < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
String redisId = getRedisId(group, userId);
|
||||
quotaManager.releaseQuota(redisId, group.tokenAmount);
|
||||
}
|
||||
|
||||
private String getRedisId(QuotaGroup group, String userId) {
|
||||
// Use the first ID as the virtual group identity if it's a custom group,
|
||||
// otherwise isolate each default user by their actual ID.
|
||||
return (group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0);
|
||||
}
|
||||
|
||||
private Duration getTtl(QuotaGroup group) {
|
||||
return group.refillSeconds > 0 ? Duration.ofSeconds(group.refillSeconds) : DEFAULT_TTL;
|
||||
}
|
||||
}
|
||||
@@ -1,171 +0,0 @@
|
||||
// Copyright 2024 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.eppserver.quota;
|
||||
|
||||
import com.google.common.collect.ImmutableMap;
|
||||
import com.google.common.flogger.FluentLogger;
|
||||
import google.registry.config.RegistryConfigSettings.Quota;
|
||||
import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup;
|
||||
import java.net.URLEncoder;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
import javax.annotation.Nullable;
|
||||
import javax.annotation.concurrent.ThreadSafe;
|
||||
import redis.clients.jedis.UnifiedJedis;
|
||||
|
||||
/**
|
||||
* Unified manager for distributed quota enforcement using Redis/Valkey.
|
||||
*
|
||||
* <p>Handles both configuration lookup and atomic Redis operations for connection and command-level
|
||||
* throttling.
|
||||
*/
|
||||
@ThreadSafe
|
||||
public class QuotaManager {
|
||||
|
||||
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
|
||||
private static final int DEFAULT_TTL_SECONDS = 3600;
|
||||
|
||||
/** Lua script to atomically decrement a token bucket with a TTL. */
|
||||
private static final String DECR_LUA =
|
||||
"local current = redis.call('GET', KEYS[1]) "
|
||||
+ "if not current then "
|
||||
+ " redis.call('SET', KEYS[1], ARGV[1] - 1, 'EX', ARGV[2]) "
|
||||
+ " return tonumber(ARGV[1]) - 1 "
|
||||
+ "end "
|
||||
+ "if tonumber(current) <= 0 then "
|
||||
+ " return -1 "
|
||||
+ "end "
|
||||
+ "return redis.call('DECR', KEYS[1])";
|
||||
|
||||
/** Lua script to atomically increment back a connection token (capped at max). */
|
||||
private static final String INCR_LUA =
|
||||
"local current = redis.call('GET', KEYS[1]) "
|
||||
+ "if current and tonumber(current) < tonumber(ARGV[1]) then "
|
||||
+ " return redis.call('INCR', KEYS[1]) "
|
||||
+ "end "
|
||||
+ "return nil";
|
||||
|
||||
/** Lua script to refresh the TTL of an existing token bucket. */
|
||||
private static final String EXPIRE_LUA =
|
||||
"if redis.call('EXISTS', KEYS[1]) == 1 then "
|
||||
+ " return redis.call('EXPIRE', KEYS[1], ARGV[1]) "
|
||||
+ "end "
|
||||
+ "return 0";
|
||||
|
||||
private final UnifiedJedis jedis;
|
||||
private final String quotaNamespace;
|
||||
private final QuotaGroup defaultQuota;
|
||||
private final ImmutableMap<String, QuotaGroup> customQuotas;
|
||||
|
||||
public QuotaManager(Quota quota, @Nullable UnifiedJedis jedis, String quotaNamespace) {
|
||||
this.jedis = jedis;
|
||||
this.quotaNamespace = quotaNamespace;
|
||||
this.defaultQuota = quota.defaultQuota;
|
||||
|
||||
ImmutableMap.Builder<String, QuotaGroup> builder = ImmutableMap.builder();
|
||||
quota.customQuota.forEach(group -> group.userId.forEach(userId -> builder.put(userId, group)));
|
||||
this.customQuotas = builder.build();
|
||||
}
|
||||
|
||||
public record QuotaRequest(String userId) {}
|
||||
|
||||
public record QuotaResponse(boolean success) {}
|
||||
|
||||
public record QuotaRebate(String userId) {}
|
||||
|
||||
/** Attempts to acquire a quota token from Redis. */
|
||||
public QuotaResponse acquireQuota(QuotaRequest request) {
|
||||
String userId = request.userId();
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
|
||||
// Unlimited quota check
|
||||
if (group.tokenAmount < 0) {
|
||||
return new QuotaResponse(true);
|
||||
}
|
||||
|
||||
if (jedis == null) {
|
||||
return new QuotaResponse(true); // Fail open if no Valkey configured
|
||||
}
|
||||
|
||||
// Use the first ID as the virtual group identity if it's a custom group,
|
||||
// otherwise isolate each default user by their actual ID.
|
||||
String redisId =
|
||||
(group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0);
|
||||
String key = String.format("%s:%s", quotaNamespace, redisId);
|
||||
int ttl = group.refillSeconds > 0 ? group.refillSeconds : DEFAULT_TTL_SECONDS;
|
||||
|
||||
try {
|
||||
Object result =
|
||||
jedis.eval(DECR_LUA, 1, key, String.valueOf(group.tokenAmount), String.valueOf(ttl));
|
||||
|
||||
return new QuotaResponse(((Long) result) >= 0);
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error for quota key: %s", URLEncoder.encode(key, StandardCharsets.UTF_8));
|
||||
return new QuotaResponse(true); // Fail open
|
||||
}
|
||||
}
|
||||
|
||||
/** Refreshes the TTL of an existing quota token. */
|
||||
public void refreshQuota(QuotaRequest request) {
|
||||
if (jedis == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
String userId = request.userId();
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
if (group.tokenAmount < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Use the first ID as the virtual group identity if it's a custom group,
|
||||
// otherwise isolate each default user by their actual ID.
|
||||
String redisId =
|
||||
(group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0);
|
||||
String key = String.format("%s:%s", quotaNamespace, redisId);
|
||||
int ttl = group.refillSeconds > 0 ? group.refillSeconds : DEFAULT_TTL_SECONDS;
|
||||
|
||||
try {
|
||||
jedis.eval(EXPIRE_LUA, 1, key, String.valueOf(ttl));
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error refreshing quota for: %s", URLEncoder.encode(key, StandardCharsets.UTF_8));
|
||||
}
|
||||
}
|
||||
|
||||
/** Returns a token to the pool (used for connection throttling). */
|
||||
public void releaseQuota(QuotaRebate rebate) {
|
||||
if (jedis == null) {
|
||||
return;
|
||||
}
|
||||
|
||||
String userId = rebate.userId();
|
||||
QuotaGroup group = customQuotas.getOrDefault(userId, defaultQuota);
|
||||
if (group.tokenAmount < 0) {
|
||||
return;
|
||||
}
|
||||
|
||||
// Use the first ID as the virtual group identity if it's a custom group,
|
||||
// otherwise isolate each default user by their actual ID.
|
||||
String redisId =
|
||||
(group == defaultQuota || group.userId.isEmpty()) ? userId : group.userId.get(0);
|
||||
String key = String.format("%s:%s", quotaNamespace, redisId);
|
||||
try {
|
||||
jedis.eval(INCR_LUA, 1, key, String.valueOf(group.tokenAmount));
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error releasing quota for: %s", URLEncoder.encode(key, StandardCharsets.UTF_8));
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.quota;
|
||||
|
||||
import static com.google.common.base.Preconditions.checkArgument;
|
||||
import static java.nio.charset.StandardCharsets.UTF_8;
|
||||
|
||||
import com.google.common.flogger.FluentLogger;
|
||||
import java.net.URLEncoder;
|
||||
import java.time.Duration;
|
||||
import javax.annotation.Nullable;
|
||||
import javax.annotation.concurrent.ThreadSafe;
|
||||
import redis.clients.jedis.UnifiedJedis;
|
||||
|
||||
/** Generic quota manager that uses Redis/Valkey as the backing store. */
|
||||
@ThreadSafe
|
||||
public class GenericValkeyQuotaManager {
|
||||
|
||||
private static final FluentLogger logger = FluentLogger.forEnclosingClass();
|
||||
|
||||
/**
|
||||
* Lua script to atomically decrement a token bucket with a TTL.
|
||||
*
|
||||
* <p>TODO(b/547996770): maybe use a more complex data structure here (SortedSet?) to manage an
|
||||
* actual sliding window. Currently this is a fixed window -- the clock "starts" when the first
|
||||
* request arrives and resets back to 0 entirely once the TTL is hit.
|
||||
*/
|
||||
private static final String DECR_LUA =
|
||||
"""
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if not current then
|
||||
redis.call('SET', KEYS[1], ARGV[1] - 1, 'PX', ARGV[2])
|
||||
return tonumber(ARGV[1]) - 1
|
||||
end
|
||||
if tonumber(current) <= 0 then
|
||||
return -1
|
||||
end
|
||||
return redis.call('DECR', KEYS[1])
|
||||
""";
|
||||
|
||||
/** Lua script to atomically increment back a connection token (capped at max). */
|
||||
private static final String INCR_LUA =
|
||||
"""
|
||||
local current = redis.call('GET', KEYS[1])
|
||||
if current and tonumber(current) < tonumber(ARGV[1]) then
|
||||
return redis.call('INCR', KEYS[1])
|
||||
end
|
||||
return nil
|
||||
""";
|
||||
|
||||
private final UnifiedJedis jedis;
|
||||
private final String namespace;
|
||||
|
||||
public GenericValkeyQuotaManager(@Nullable UnifiedJedis jedis, String namespace) {
|
||||
this.jedis = jedis;
|
||||
this.namespace = namespace;
|
||||
}
|
||||
|
||||
/** Attempts to acquire a quota token from Valkey. */
|
||||
public boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDuration) {
|
||||
if (jedis == null) {
|
||||
return true; // Fail open if no Valkey configured
|
||||
}
|
||||
checkArgument(expirationDuration.isPositive(), "Duration must be positive");
|
||||
checkArgument(maxTokenAmount >= 0, "Max token amount must be non-negative");
|
||||
|
||||
String key = createValkeyKey(id);
|
||||
try {
|
||||
Object result =
|
||||
jedis.eval(
|
||||
DECR_LUA,
|
||||
1,
|
||||
key,
|
||||
String.valueOf(maxTokenAmount),
|
||||
String.valueOf(expirationDuration.toMillis()));
|
||||
return (Long) result >= 0;
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error for quota key: %s", URLEncoder.encode(key, UTF_8));
|
||||
// Fail open
|
||||
return true;
|
||||
}
|
||||
}
|
||||
|
||||
/** Refreshes the TTL of an existing quota token. */
|
||||
public void refreshQuota(String id, Duration expirationDuration) {
|
||||
if (jedis == null) {
|
||||
return;
|
||||
}
|
||||
checkArgument(expirationDuration.isPositive(), "Duration must be positive");
|
||||
|
||||
String key = createValkeyKey(id);
|
||||
try {
|
||||
jedis.pexpire(key, expirationDuration.toMillis());
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error refreshing quota for: %s", URLEncoder.encode(key, UTF_8));
|
||||
}
|
||||
}
|
||||
|
||||
/** Returns a token to the pool (used for connection throttling). */
|
||||
public void releaseQuota(String id, int maxTokenAmount) {
|
||||
if (jedis == null) {
|
||||
return;
|
||||
}
|
||||
checkArgument(maxTokenAmount >= 0, "Max token amount must be non-negative");
|
||||
|
||||
String key = createValkeyKey(id);
|
||||
try {
|
||||
jedis.eval(INCR_LUA, 1, key, String.valueOf(maxTokenAmount));
|
||||
} catch (Exception e) {
|
||||
logger.atSevere().withCause(e).log(
|
||||
"Valkey error releasing quota for: %s", URLEncoder.encode(key, UTF_8));
|
||||
}
|
||||
}
|
||||
|
||||
private String createValkeyKey(String id) {
|
||||
return String.format("%s:%s", namespace, id);
|
||||
}
|
||||
}
|
||||
@@ -30,10 +30,8 @@ import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import google.registry.eppserver.metric.FrontendMetrics;
|
||||
import google.registry.eppserver.quota.EppServerQuotaManager;
|
||||
import google.registry.eppserver.quota.LocalConnectionLimiter;
|
||||
import google.registry.eppserver.quota.QuotaManager;
|
||||
import google.registry.eppserver.quota.QuotaManager.QuotaRequest;
|
||||
import google.registry.eppserver.quota.QuotaManager.QuotaResponse;
|
||||
import google.registry.request.RequestHandler;
|
||||
import google.registry.util.FakeHttpServletRequest;
|
||||
import google.registry.util.FakeHttpServletResponse;
|
||||
@@ -68,7 +66,7 @@ class EppServiceHandlerTest {
|
||||
|
||||
@Mock private FrontendMetrics metrics;
|
||||
@Mock private LocalConnectionLimiter localConnectionLimiter;
|
||||
@Mock private QuotaManager commandQuotaManager;
|
||||
@Mock private EppServerQuotaManager commandQuotaManager;
|
||||
@Mock private Supplier<String> idTokenSupplier;
|
||||
@Mock private ChannelHandlerContext ctx;
|
||||
@Mock private Channel channel;
|
||||
@@ -116,9 +114,7 @@ class EppServiceHandlerTest {
|
||||
.when(executor)
|
||||
.schedule(any(Runnable.class), anyLong(), any(TimeUnit.class));
|
||||
|
||||
lenient()
|
||||
.when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
lenient().when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
}
|
||||
|
||||
private void setUpSuccessfulHandshake() throws Exception {
|
||||
@@ -149,8 +145,7 @@ class EppServiceHandlerTest {
|
||||
.when(requestHandler)
|
||||
.handleRequest(any(FakeHttpServletRequest.class), any(FakeHttpServletResponse.class));
|
||||
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
when(idTokenSupplier.get()).thenReturn("fake_id_token");
|
||||
|
||||
setUpSuccessfulHandshake();
|
||||
@@ -192,8 +187,7 @@ class EppServiceHandlerTest {
|
||||
setUpSuccessfulHandshake();
|
||||
|
||||
when(idTokenSupplier.get()).thenReturn("fake_id_token");
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
|
||||
String eppLoginXml = "<epp><command><login><clID>RegistrarA</clID></login></command></epp>";
|
||||
ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8));
|
||||
@@ -221,8 +215,7 @@ class EppServiceHandlerTest {
|
||||
setUpSuccessfulHandshake();
|
||||
|
||||
when(idTokenSupplier.get()).thenReturn("fake_id_token");
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
|
||||
String eppLoginXml = "<epp><command><login><clID>RegistrarA</clID></login></command></epp>";
|
||||
ByteBuf inFrame = Unpooled.wrappedBuffer(eppLoginXml.getBytes(UTF_8));
|
||||
@@ -248,7 +241,7 @@ class EppServiceHandlerTest {
|
||||
|
||||
// Verify command quota was requested for the cert hash pre-login
|
||||
String certHash = X509Utils.getCertificateHash(certificate);
|
||||
verify(commandQuotaManager, times(2)).acquireQuota(eq(new QuotaRequest(certHash)));
|
||||
verify(commandQuotaManager, times(2)).acquireQuota(eq(certHash));
|
||||
verify(localConnectionLimiter).acquireRegistrar("RegistrarA");
|
||||
verify(scheduledFuture).cancel(eq(false));
|
||||
|
||||
@@ -284,15 +277,14 @@ class EppServiceHandlerTest {
|
||||
handler.channelRead0(ctx, inFrame2);
|
||||
|
||||
// Verify command quota was requested for the authenticated registrar post-login
|
||||
verify(commandQuotaManager).acquireQuota(eq(new QuotaRequest("RegistrarA")));
|
||||
verify(commandQuotaManager).acquireQuota(eq("RegistrarA"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testChannelRead0_commandQuotaRejected() throws Exception {
|
||||
setUpSuccessfulHandshake();
|
||||
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(false));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(false);
|
||||
|
||||
String eppXml = "<epp><command><check></check></command></epp>";
|
||||
ByteBuf inFrame = Unpooled.wrappedBuffer(eppXml.getBytes(UTF_8));
|
||||
@@ -308,8 +300,7 @@ class EppServiceHandlerTest {
|
||||
setUpSuccessfulHandshake();
|
||||
|
||||
when(idTokenSupplier.get()).thenReturn("fake_id_token");
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
|
||||
String eppLogoutXml = "<epp><command><logout/></command></epp>";
|
||||
ByteBuf inFrame = Unpooled.wrappedBuffer(eppLogoutXml.getBytes(UTF_8));
|
||||
@@ -348,8 +339,7 @@ class EppServiceHandlerTest {
|
||||
setUpSuccessfulHandshake();
|
||||
|
||||
when(idTokenSupplier.get()).thenReturn("fake_id_token");
|
||||
when(commandQuotaManager.acquireQuota(any(QuotaRequest.class)))
|
||||
.thenReturn(new QuotaResponse(true));
|
||||
when(commandQuotaManager.acquireQuota(any(String.class))).thenReturn(true);
|
||||
when(localConnectionLimiter.acquireRegistrar("RegistrarA")).thenReturn(true);
|
||||
|
||||
String eppLoginXml = "<epp><command><login><clID>RegistrarA</clID></login></command></epp>";
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.eppserver.quota;
|
||||
|
||||
import static com.google.common.truth.Truth.assertThat;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.verifyNoInteractions;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import google.registry.config.RegistryConfigSettings.Quota;
|
||||
import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup;
|
||||
import google.registry.quota.GenericValkeyQuotaManager;
|
||||
import java.time.Duration;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class EppServerQuotaManagerTest {
|
||||
|
||||
@Mock private GenericValkeyQuotaManager quotaManager;
|
||||
|
||||
private Quota quotaConfig;
|
||||
private EppServerQuotaManager manager;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
quotaConfig = new Quota();
|
||||
QuotaGroup defaultGroup = new QuotaGroup();
|
||||
defaultGroup.tokenAmount = 10;
|
||||
defaultGroup.refillSeconds = 60;
|
||||
quotaConfig.defaultQuota = defaultGroup;
|
||||
|
||||
QuotaGroup customGroup = new QuotaGroup();
|
||||
customGroup.tokenAmount = 5;
|
||||
customGroup.refillSeconds = 30;
|
||||
customGroup.userId = ImmutableList.of("user1");
|
||||
quotaConfig.customQuota = ImmutableList.of(customGroup);
|
||||
|
||||
manager = new EppServerQuotaManager(quotaConfig, quotaManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_defaultQuota() {
|
||||
when(quotaManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).thenReturn(true);
|
||||
|
||||
assertThat(manager.acquireQuota("user2")).isTrue();
|
||||
verify(quotaManager).acquireQuota("user2", 10, Duration.ofMinutes(1));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_customQuota() {
|
||||
when(quotaManager.acquireQuota("user1", 5, Duration.ofSeconds(30))).thenReturn(true);
|
||||
|
||||
assertThat(manager.acquireQuota("user1")).isTrue();
|
||||
verify(quotaManager).acquireQuota("user1", 5, Duration.ofSeconds(30));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_unlimited() {
|
||||
quotaConfig.defaultQuota.tokenAmount = -1;
|
||||
manager = new EppServerQuotaManager(quotaConfig, quotaManager);
|
||||
|
||||
assertThat(manager.acquireQuota("user2")).isTrue();
|
||||
verifyNoInteractions(quotaManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_success() {
|
||||
manager.refreshQuota("user2");
|
||||
verify(quotaManager).refreshQuota("user2", Duration.ofMinutes(1));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_unlimited_noop() {
|
||||
quotaConfig.defaultQuota.tokenAmount = -1;
|
||||
manager = new EppServerQuotaManager(quotaConfig, quotaManager);
|
||||
|
||||
manager.refreshQuota("user2");
|
||||
verifyNoInteractions(quotaManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_success() {
|
||||
manager.releaseQuota("user2");
|
||||
verify(quotaManager).releaseQuota("user2", 10);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_unlimited_noop() {
|
||||
quotaConfig.defaultQuota.tokenAmount = -1;
|
||||
manager = new EppServerQuotaManager(quotaConfig, quotaManager);
|
||||
|
||||
manager.releaseQuota("user2");
|
||||
verifyNoInteractions(quotaManager);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGroupVirtualIdentity_usesFirstIdInList() {
|
||||
// Modify config so "user1" is accompanied by a virtual group ID "my_group"
|
||||
quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user3");
|
||||
manager = new EppServerQuotaManager(quotaConfig, quotaManager);
|
||||
|
||||
when(quotaManager.acquireQuota("my_group", 5, Duration.ofSeconds(30))).thenReturn(true);
|
||||
|
||||
assertThat(manager.acquireQuota("user1")).isTrue();
|
||||
assertThat(manager.acquireQuota("user3")).isTrue();
|
||||
verify(quotaManager, times(2)).acquireQuota("my_group", 5, Duration.ofSeconds(30));
|
||||
}
|
||||
}
|
||||
@@ -1,159 +0,0 @@
|
||||
// Copyright 2024 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.eppserver.quota;
|
||||
|
||||
import static com.google.common.truth.Truth.assertThat;
|
||||
import static org.mockito.ArgumentMatchers.anyInt;
|
||||
import static org.mockito.ArgumentMatchers.anyString;
|
||||
import static org.mockito.ArgumentMatchers.eq;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
|
||||
import com.google.common.collect.ImmutableList;
|
||||
import google.registry.config.RegistryConfigSettings.Quota;
|
||||
import google.registry.config.RegistryConfigSettings.Quota.QuotaGroup;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.extension.ExtendWith;
|
||||
import org.mockito.Mock;
|
||||
import org.mockito.junit.jupiter.MockitoExtension;
|
||||
import redis.clients.jedis.UnifiedJedis;
|
||||
|
||||
@ExtendWith(MockitoExtension.class)
|
||||
class QuotaManagerTest {
|
||||
|
||||
@Mock private UnifiedJedis jedis;
|
||||
|
||||
private Quota quotaConfig;
|
||||
private QuotaManager manager;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
quotaConfig = new Quota();
|
||||
QuotaGroup defaultGroup = new QuotaGroup();
|
||||
defaultGroup.tokenAmount = 10;
|
||||
defaultGroup.refillSeconds = 60;
|
||||
quotaConfig.defaultQuota = defaultGroup;
|
||||
|
||||
QuotaGroup customGroup = new QuotaGroup();
|
||||
customGroup.tokenAmount = 5;
|
||||
customGroup.refillSeconds = 30;
|
||||
customGroup.userId = ImmutableList.of("user1");
|
||||
quotaConfig.customQuota = ImmutableList.of(customGroup);
|
||||
|
||||
manager = new QuotaManager(quotaConfig, jedis, "testQuota");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_success() {
|
||||
when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(5L);
|
||||
|
||||
QuotaManager.QuotaResponse response =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user2"));
|
||||
assertThat(response.success()).isTrue();
|
||||
verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("10"), eq("60"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_failure() {
|
||||
when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(-1L);
|
||||
|
||||
QuotaManager.QuotaResponse response =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user1"));
|
||||
assertThat(response.success()).isFalse();
|
||||
verify(jedis).eval(anyString(), eq(1), eq("testQuota:user1"), eq("5"), eq("30"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_unlimited() {
|
||||
quotaConfig.defaultQuota.tokenAmount = -1;
|
||||
manager = new QuotaManager(quotaConfig, jedis, "testQuota");
|
||||
|
||||
QuotaManager.QuotaResponse response =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user2"));
|
||||
assertThat(response.success()).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_jedisException_failsOpen() {
|
||||
when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString()))
|
||||
.thenThrow(new RuntimeException("Redis error"));
|
||||
|
||||
QuotaManager.QuotaResponse response =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user2"));
|
||||
assertThat(response.success()).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_success() {
|
||||
manager.refreshQuota(new QuotaManager.QuotaRequest("user2"));
|
||||
verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("60"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_success() {
|
||||
manager.releaseQuota(new QuotaManager.QuotaRebate("user2"));
|
||||
verify(jedis).eval(anyString(), eq(1), eq("testQuota:user2"), eq("10"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGroupVirtualIdentity_usesFirstIdInList() {
|
||||
// Modify config so "user1" is accompanied by a virtual group ID "my_group"
|
||||
quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user3");
|
||||
manager = new QuotaManager(quotaConfig, jedis, "testQuota");
|
||||
|
||||
when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString())).thenReturn(5L);
|
||||
|
||||
QuotaManager.QuotaResponse response1 =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user1"));
|
||||
QuotaManager.QuotaResponse response2 =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user3"));
|
||||
|
||||
assertThat(response1.success()).isTrue();
|
||||
assertThat(response2.success()).isTrue();
|
||||
// 5 tokens, 30 seconds ttl
|
||||
verify(jedis, times(2)).eval(anyString(), eq(1), eq("testQuota:my_group"), eq("5"), eq("30"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testGroupVirtualIdentity_exceedsQuota_fails() {
|
||||
// Modify config so "user1", "user2", "user3" share virtual group ID "my_group"
|
||||
quotaConfig.customQuota.get(0).userId = ImmutableList.of("my_group", "user1", "user2", "user3");
|
||||
manager = new QuotaManager(quotaConfig, jedis, "testQuota");
|
||||
|
||||
// Simulate Redis returning 1, 0 for successful decrements, and -1 when empty
|
||||
when(jedis.eval(anyString(), anyInt(), anyString(), anyString(), anyString()))
|
||||
.thenReturn(1L)
|
||||
.thenReturn(0L)
|
||||
.thenReturn(-1L);
|
||||
|
||||
// Act
|
||||
QuotaManager.QuotaResponse response1 =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user1"));
|
||||
QuotaManager.QuotaResponse response2 =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user2"));
|
||||
QuotaManager.QuotaResponse response3 =
|
||||
manager.acquireQuota(new QuotaManager.QuotaRequest("user3"));
|
||||
|
||||
// Assert that the third request to the same group fails
|
||||
assertThat(response1.success()).isTrue();
|
||||
assertThat(response2.success()).isTrue();
|
||||
assertThat(response3.success()).isFalse();
|
||||
|
||||
// Verify all 3 requests went to the shared bucket
|
||||
verify(jedis, times(3)).eval(anyString(), eq(1), eq("testQuota:my_group"), eq("5"), eq("30"));
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,170 @@
|
||||
// Copyright 2026 The Nomulus Authors. All Rights Reserved.
|
||||
//
|
||||
// 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
|
||||
//
|
||||
// http://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 google.registry.quota;
|
||||
|
||||
import static com.google.common.truth.Truth.assertThat;
|
||||
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
|
||||
|
||||
import io.github.ss_bhatt.testcontainers.valkey.ValkeyContainer;
|
||||
import java.time.Duration;
|
||||
import org.junit.jupiter.api.BeforeEach;
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
import redis.clients.jedis.HostAndPort;
|
||||
import redis.clients.jedis.RedisClient;
|
||||
|
||||
@Testcontainers
|
||||
class GenericValkeyQuotaManagerTest {
|
||||
|
||||
@Container private static final ValkeyContainer valkey = new ValkeyContainer();
|
||||
|
||||
private RedisClient jedis;
|
||||
private GenericValkeyQuotaManager quotaManager;
|
||||
|
||||
@BeforeEach
|
||||
void setUp() {
|
||||
jedis =
|
||||
RedisClient.builder()
|
||||
.hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort()))
|
||||
.build();
|
||||
jedis.flushAll();
|
||||
quotaManager = new GenericValkeyQuotaManager(jedis, "testQuota");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_success() {
|
||||
assertThat(quotaManager.acquireQuota("user1", 5, Duration.ofMinutes(1))).isTrue();
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("4");
|
||||
assertThat(jedis.ttl("testQuota:user1")).isGreaterThan(0L);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_exhaustsQuota_thenFails() {
|
||||
assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isTrue();
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("1");
|
||||
|
||||
assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isTrue();
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("0");
|
||||
|
||||
assertThat(quotaManager.acquireQuota("user1", 2, Duration.ofMinutes(1))).isFalse();
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("0");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_resetsAfterExpiration() throws Exception {
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isTrue();
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isFalse();
|
||||
|
||||
Thread.sleep(150);
|
||||
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMillis(50))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_isolatedByNamespaceAndId() {
|
||||
GenericValkeyQuotaManager otherQuotaManager =
|
||||
new GenericValkeyQuotaManager(jedis, "otherQuota");
|
||||
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue();
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse();
|
||||
|
||||
// user2 in same namespace is independent
|
||||
assertThat(quotaManager.acquireQuota("user2", 1, Duration.ofMinutes(1))).isTrue();
|
||||
|
||||
// user1 in other namespace is independent
|
||||
assertThat(otherQuotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_nullJedis_failsOpen() {
|
||||
GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota");
|
||||
assertThat(nullJedisManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testAcquireQuota_jedisException_failsOpen() {
|
||||
jedis.close();
|
||||
assertThat(quotaManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_success() {
|
||||
quotaManager.acquireQuota("user1", 5, Duration.ofSeconds(10));
|
||||
|
||||
quotaManager.refreshQuota("user1", Duration.ofMinutes(5));
|
||||
assertThat(jedis.ttl("testQuota:user1")).isGreaterThan(10L);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_nonexistentKey_noop() {
|
||||
quotaManager.refreshQuota("nonexistent", Duration.ofMinutes(5));
|
||||
assertThat(jedis.exists("testQuota:nonexistent")).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_nullJedis_noop() {
|
||||
GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota");
|
||||
assertDoesNotThrow(() -> nullJedisManager.refreshQuota("user2", Duration.ofMinutes(1)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testRefreshQuota_jedisException_handled() {
|
||||
jedis.close();
|
||||
assertDoesNotThrow(() -> quotaManager.refreshQuota("user2", Duration.ofMinutes(1)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_success() {
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue();
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse();
|
||||
|
||||
quotaManager.releaseQuota("user1", 1);
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("1");
|
||||
|
||||
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_cappedAtMax() {
|
||||
quotaManager.acquireQuota("user1", 3, Duration.ofMinutes(1));
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("2");
|
||||
|
||||
quotaManager.releaseQuota("user1", 3);
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("3");
|
||||
|
||||
// Releasing again when already at max should not increment past max
|
||||
quotaManager.releaseQuota("user1", 3);
|
||||
assertThat(jedis.get("testQuota:user1")).isEqualTo("3");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_nonexistentKey_noop() {
|
||||
quotaManager.releaseQuota("nonexistent", 5);
|
||||
assertThat(jedis.exists("testQuota:nonexistent")).isFalse();
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_nullJedis_noop() {
|
||||
GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota");
|
||||
assertDoesNotThrow(() -> nullJedisManager.releaseQuota("user2", 10));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testReleaseQuota_jedisException_handled() {
|
||||
jedis.close();
|
||||
assertDoesNotThrow(() -> quotaManager.releaseQuota("user2", 10));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user