diff --git a/core/src/main/java/google/registry/eppserver/EppProtocolModule.java b/core/src/main/java/google/registry/eppserver/EppProtocolModule.java index 576ea8f07..becf05287 100644 --- a/core/src/main/java/google/registry/eppserver/EppProtocolModule.java +++ b/core/src/main/java/google/registry/eppserver/EppProtocolModule.java @@ -26,9 +26,9 @@ import google.registry.eppserver.Protocol.FrontendProtocol; import google.registry.eppserver.handler.EppServiceHandler; 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 google.registry.quota.ValkeyQuotaManager; import io.netty.channel.ChannelHandler; import io.netty.channel.socket.nio.NioSocketChannel; import io.netty.handler.codec.LengthFieldBasedFrameDecoder; @@ -143,9 +143,7 @@ public final class EppProtocolModule { static EppServerQuotaManager provideCommandQuotaManager( @Config("eppServerQuota") RegistryConfigSettings.Quota quota, Optional jedis) { QuotaManager quotaManager = - jedis.isPresent() - ? GenericValkeyQuotaManager.create(jedis.get(), "command") - : new NoopQuotaManager(); + jedis.isPresent() ? new ValkeyQuotaManager(jedis.get(), "command") : new NoopQuotaManager(); return new EppServerQuotaManager(quota, quotaManager); } } diff --git a/core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java b/core/src/main/java/google/registry/quota/ValkeyQuotaManager.java similarity index 73% rename from core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java rename to core/src/main/java/google/registry/quota/ValkeyQuotaManager.java index 246d9e79a..21aa75762 100644 --- a/core/src/main/java/google/registry/quota/GenericValkeyQuotaManager.java +++ b/core/src/main/java/google/registry/quota/ValkeyQuotaManager.java @@ -19,14 +19,16 @@ import static com.google.common.base.Preconditions.checkNotNull; import static java.nio.charset.StandardCharsets.UTF_8; import com.google.common.flogger.FluentLogger; +import com.google.common.hash.Hashing; import java.net.URLEncoder; import java.time.Duration; import javax.annotation.concurrent.ThreadSafe; import redis.clients.jedis.UnifiedJedis; +import redis.clients.jedis.exceptions.JedisDataException; /** Generic quota manager that uses Redis/Valkey as the backing store. */ @ThreadSafe -public class GenericValkeyQuotaManager implements QuotaManager { +public class ValkeyQuotaManager implements QuotaManager { private static final FluentLogger logger = FluentLogger.forEnclosingClass(); @@ -60,19 +62,18 @@ public class GenericValkeyQuotaManager implements QuotaManager { return nil """; + // Valkey has the capability to upload scripts and refer to them by their SHA-1 hash (SHA-1 is OK + // because the point isn't security or encryption, just shortening). This means we don't need to + // pass the whole script each time. + private static final String DECR_LUA_HASH = Hashing.sha1().hashString(DECR_LUA, UTF_8).toString(); + private static final String INCR_LUA_HASH = Hashing.sha1().hashString(INCR_LUA, UTF_8).toString(); + private final UnifiedJedis jedis; private final 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; + public ValkeyQuotaManager(UnifiedJedis jedis, String namespace) { + this.jedis = checkNotNull(jedis, "jedis must not be null"); + this.namespace = checkNotNull(namespace, "namespace must not be null"); } /** Attempts to acquire a quota token from Valkey. */ @@ -84,9 +85,9 @@ public class GenericValkeyQuotaManager implements QuotaManager { String key = createValkeyKey(id); try { Object result = - jedis.eval( + runScript( DECR_LUA, - 1, + DECR_LUA_HASH, key, String.valueOf(maxTokenAmount), String.valueOf(expirationDuration.toMillis())); @@ -120,7 +121,7 @@ public class GenericValkeyQuotaManager implements QuotaManager { String key = createValkeyKey(id); try { - jedis.eval(INCR_LUA, 1, key, String.valueOf(maxTokenAmount)); + runScript(INCR_LUA, INCR_LUA_HASH, key, String.valueOf(maxTokenAmount)); } catch (Exception e) { logger.atSevere().withCause(e).log( "Valkey error releasing quota for: %s", URLEncoder.encode(key, UTF_8)); @@ -130,4 +131,18 @@ public class GenericValkeyQuotaManager implements QuotaManager { private String createValkeyKey(String id) { return String.format("%s:%s", namespace, id); } + + private Object runScript(String script, String scriptHash, String... params) { + try { + return jedis.evalsha(scriptHash, 1, params); + } catch (JedisDataException e) { + // Scripts are technically ephemeral and while they aren't evicted during the normal course of + // operations, we can't guarantee that the Valkey server hasn't restarted + if (e.getMessage() != null && e.getMessage().startsWith("NOSCRIPT")) { + jedis.scriptLoad(script); + return jedis.evalsha(scriptHash, 1, params); + } + throw e; + } + } } diff --git a/core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java b/core/src/test/java/google/registry/quota/ValkeyQuotaManagerTest.java similarity index 80% rename from core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java rename to core/src/test/java/google/registry/quota/ValkeyQuotaManagerTest.java index 62bac2711..1c5118af9 100644 --- a/core/src/test/java/google/registry/quota/GenericValkeyQuotaManagerTest.java +++ b/core/src/test/java/google/registry/quota/ValkeyQuotaManagerTest.java @@ -28,12 +28,12 @@ import redis.clients.jedis.HostAndPort; import redis.clients.jedis.RedisClient; @Testcontainers -class GenericValkeyQuotaManagerTest { +class ValkeyQuotaManagerTest { @Container private static final ValkeyContainer valkey = new ValkeyContainer(); private RedisClient jedis; - private GenericValkeyQuotaManager quotaManager; + private ValkeyQuotaManager quotaManager; @BeforeEach void setUp() { @@ -42,7 +42,7 @@ class GenericValkeyQuotaManagerTest { .hostAndPort(new HostAndPort(valkey.getHost(), valkey.getFirstMappedPort())) .build(); jedis.flushAll(); - quotaManager = GenericValkeyQuotaManager.create(jedis, "testQuota"); + quotaManager = new ValkeyQuotaManager(jedis, "testQuota"); } @Test @@ -76,8 +76,7 @@ class GenericValkeyQuotaManagerTest { @Test void testAcquireQuota_isolatedByNamespaceAndId() { - GenericValkeyQuotaManager otherQuotaManager = - GenericValkeyQuotaManager.create(jedis, "otherQuota"); + ValkeyQuotaManager otherQuotaManager = new ValkeyQuotaManager(jedis, "otherQuota"); assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isTrue(); assertThat(quotaManager.acquireQuota("user1", 1, Duration.ofMinutes(1))).isFalse(); @@ -90,9 +89,13 @@ class GenericValkeyQuotaManagerTest { } @Test - void testCreate_nullJedis_throwsNpe() { - assertThrows( - NullPointerException.class, () -> GenericValkeyQuotaManager.create(null, "testQuota")); + void testConstructor_nullJedis_throwsNpe() { + assertThrows(NullPointerException.class, () -> new ValkeyQuotaManager(null, "testQuota")); + } + + @Test + void testConstructor_nullNamespace_throwsNpe() { + assertThrows(NullPointerException.class, () -> new ValkeyQuotaManager(jedis, null)); } @Test @@ -156,4 +159,23 @@ class GenericValkeyQuotaManagerTest { jedis.close(); assertDoesNotThrow(() -> quotaManager.releaseQuota("user2", 10)); } + + @Test + void testScriptFlushed_reloadsScriptAndSucceeds() { + assertThat(quotaManager.acquireQuota("user1", 3, Duration.ofMinutes(1))).isTrue(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("2"); + + // Simulate scripts being evicted or server restart + jedis.scriptFlush(); + + // acquireQuota should reload the scripts and continue working + assertThat(quotaManager.acquireQuota("user1", 3, Duration.ofMinutes(1))).isTrue(); + assertThat(jedis.get("testQuota:user1")).isEqualTo("1"); + + jedis.scriptFlush(); + + // releaseQuota should also reload the scripts and continue working + quotaManager.releaseQuota("user1", 3); + assertThat(jedis.get("testQuota:user1")).isEqualTo("2"); + } }