Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
import google.registry.eppserver.quota.EppServerQuotaManager;
import google.registry.networking.handler.SslServerInitializer;
import google.registry.quota.GenericValkeyQuotaManager;
import google.registry.quota.NoopQuotaManager;
import google.registry.quota.QuotaManager;
import io.netty.channel.ChannelHandler;
import io.netty.channel.socket.nio.NioSocketChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
Expand Down Expand Up @@ -140,7 +142,10 @@ static SslServerInitializer<NioSocketChannel> provideSslServerInitializer(
@CommandQuota
static EppServerQuotaManager provideCommandQuotaManager(
@Config("eppServerQuota") RegistryConfigSettings.Quota quota, Optional<UnifiedJedis> jedis) {
return new EppServerQuotaManager(
quota, new GenericValkeyQuotaManager(jedis.orElse(null), "command"));
QuotaManager quotaManager =
jedis.isPresent()
? GenericValkeyQuotaManager.create(jedis.get(), "command")
: new NoopQuotaManager();
return new EppServerQuotaManager(quota, quotaManager);
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@
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 google.registry.quota.QuotaManager;
import java.time.Duration;
import javax.annotation.concurrent.ThreadSafe;

Expand All @@ -31,11 +31,11 @@ public class EppServerQuotaManager {

private static final Duration DEFAULT_TTL = Duration.ofHours(1);

private final GenericValkeyQuotaManager quotaManager;
private final QuotaManager quotaManager;
private final QuotaGroup defaultQuota;
private final ImmutableMap<String, QuotaGroup> customQuotas;

public EppServerQuotaManager(Quota quota, GenericValkeyQuotaManager quotaManager) {
public EppServerQuotaManager(Quota quota, QuotaManager quotaManager) {
this.quotaManager = quotaManager;
this.defaultQuota = quota.defaultQuota;

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -15,18 +15,18 @@
package google.registry.quota;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;
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 {
public class GenericValkeyQuotaManager implements QuotaManager {

private static final FluentLogger logger = FluentLogger.forEnclosingClass();

Expand Down Expand Up @@ -63,16 +63,21 @@ public class GenericValkeyQuotaManager {
private final UnifiedJedis jedis;
private final String namespace;

public GenericValkeyQuotaManager(@Nullable UnifiedJedis jedis, String namespace) {
public static GenericValkeyQuotaManager create(UnifiedJedis jedis, String namespace) {
// TODO(gbrodman): upload the scripts to Valkey so we can reference them by hash rather than
// uploading the whole script each time
return new GenericValkeyQuotaManager(jedis, namespace);
}

private GenericValkeyQuotaManager(UnifiedJedis jedis, String namespace) {
checkNotNull(jedis, "UnifiedJedis cannot be null");
this.jedis = jedis;
this.namespace = namespace;
}

/** Attempts to acquire a quota token from Valkey. */
@Override
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");

Expand All @@ -95,10 +100,8 @@ public boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDu
}

/** Refreshes the TTL of an existing quota token. */
@Override
public void refreshQuota(String id, Duration expirationDuration) {
if (jedis == null) {
return;
}
checkArgument(expirationDuration.isPositive(), "Duration must be positive");

String key = createValkeyKey(id);
Expand All @@ -111,10 +114,8 @@ public void refreshQuota(String id, Duration expirationDuration) {
}

/** Returns a token to the pool (used for connection throttling). */
@Override
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);
Expand Down
38 changes: 38 additions & 0 deletions core/src/main/java/google/registry/quota/NoopQuotaManager.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
// 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 java.time.Duration;
import javax.annotation.concurrent.ThreadSafe;

/**
* A {@link QuotaManager} that always returns true and performs no actions.
*
* <p>Used when Valkey/Jedis are not available.
*/
@ThreadSafe
public class NoopQuotaManager implements QuotaManager {

@Override
public boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDuration) {
return true;
}

@Override
public void refreshQuota(String id, Duration expirationDuration) {}

@Override
public void releaseQuota(String id, int maxTokenAmount) {}
}
29 changes: 29 additions & 0 deletions core/src/main/java/google/registry/quota/QuotaManager.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
// 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 java.time.Duration;

/** Interface for managing quota. */
public interface QuotaManager {
/** Attempts to acquire a token (out of the given max amount) with the given TTL. */
boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDuration);

/** Refreshes the TTL of an existing token. */
void refreshQuota(String id, Duration expirationDuration);

/** Returns a token to the pool (possibly useful for connection throttling). */
void releaseQuota(String id, int maxTokenAmount);
}
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
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 google.registry.quota.QuotaManager;
import java.time.Duration;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
Expand All @@ -34,7 +34,7 @@
@ExtendWith(MockitoExtension.class)
class EppServerQuotaManagerTest {

@Mock private GenericValkeyQuotaManager quotaManager;
@Mock private QuotaManager quotaManager;

private Quota quotaConfig;
private EppServerQuotaManager manager;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@

import static com.google.common.truth.Truth.assertThat;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertThrows;

import io.github.ss_bhatt.testcontainers.valkey.ValkeyContainer;
import java.time.Duration;
Expand All @@ -41,7 +42,7 @@ void setUp() {
.hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort()))
.build();
jedis.flushAll();
quotaManager = new GenericValkeyQuotaManager(jedis, "testQuota");
quotaManager = GenericValkeyQuotaManager.create(jedis, "testQuota");
}

@Test
Expand Down Expand Up @@ -76,7 +77,7 @@ void testAcquireQuota_resetsAfterExpiration() throws Exception {
@Test
void testAcquireQuota_isolatedByNamespaceAndId() {
GenericValkeyQuotaManager otherQuotaManager =
new GenericValkeyQuotaManager(jedis, "otherQuota");
GenericValkeyQuotaManager.create(jedis, "otherQuota");

assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue();
assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse();
Expand All @@ -89,9 +90,9 @@ void testAcquireQuota_isolatedByNamespaceAndId() {
}

@Test
void testAcquireQuota_nullJedis_failsOpen() {
GenericValkeyQuotaManager nullJedisManager = new GenericValkeyQuotaManager(null, "testQuota");
assertThat(nullJedisManager.acquireQuota("user2", 10, Duration.ofMinutes(1))).isTrue();
void testCreate_nullJedis_throwsNpe() {
assertThrows(
NullPointerException.class, () -> GenericValkeyQuotaManager.create(null, "testQuota"));
}

@Test
Expand All @@ -114,12 +115,6 @@ void testRefreshQuota_nonexistentKey_noop() {
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();
Expand Down Expand Up @@ -156,12 +151,6 @@ void testReleaseQuota_nonexistentKey_noop() {
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();
Expand Down
41 changes: 41 additions & 0 deletions core/src/test/java/google/registry/quota/NoopQuotaManagerTest.java
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
// 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 java.time.Duration;
import org.junit.jupiter.api.Test;

class NoopQuotaManagerTest {

private final NoopQuotaManager quotaManager = new NoopQuotaManager();

@Test
void testAcquireQuota_returnsTrue() {
assertThat(quotaManager.acquireQuota("user1", 10, Duration.ofMinutes(1))).isTrue();
}

@Test
void testRefreshQuota_noop() {
assertDoesNotThrow(() -> quotaManager.refreshQuota("user1", Duration.ofMinutes(1)));
}

@Test
void testReleaseQuota_noop() {
assertDoesNotThrow(() -> quotaManager.releaseQuota("user1", 10));
}
}
Loading