From a5db76cd3acb54b65e6bf46e0f28a930ebcad057 Mon Sep 17 00:00:00 2001 From: Gus Brodman Date: Wed, 19 Aug 2026 16:29:52 -0400 Subject: [PATCH] Upload Lua scripts to Valkey for quota management We don't want to have to pass the entire script to Valkey every single time. Instead, we compute the SHA-1 hash of the script and refer to it by hash, significantly reducing the amount of bytes we need to send to the server. If the script is not loaded in Valkey (NOSCRIPT error), we reload it and retry. --- .../registry/eppserver/EppProtocolModule.java | 6 +-- ...taManager.java => ValkeyQuotaManager.java} | 43 +++++++++++++------ ...rTest.java => ValkeyQuotaManagerTest.java} | 38 ++++++++++++---- 3 files changed, 61 insertions(+), 26 deletions(-) rename core/src/main/java/google/registry/quota/{GenericValkeyQuotaManager.java => ValkeyQuotaManager.java} (73%) rename core/src/test/java/google/registry/quota/{GenericValkeyQuotaManagerTest.java => ValkeyQuotaManagerTest.java} (80%) diff --git a/core/src/main/java/google/registry/eppserver/EppProtocolModule.java b/core/src/main/java/google/registry/eppserver/EppProtocolModule.java index 576ea8f0796..becf052877d 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.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 @@ static SslServerInitializer provideSslServerInitializer( 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 246d9e79ac8..21aa75762d3 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 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 boolean acquireQuota(String id, int maxTokenAmount, Duration expirationDu 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 void releaseQuota(String id, int maxTokenAmount) { 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 void releaseQuota(String id, int maxTokenAmount) { 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 62bac2711cf..1c5118af96b 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.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 @@ void setUp() { .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 @@ void testAcquireQuota_resetsAfterExpiration() throws Exception { @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 @@ void testAcquireQuota_isolatedByNamespaceAndId() { } @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 @@ void testReleaseQuota_jedisException_handled() { 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"); + } }