From 7825a92fe07d67d93c6156a1e193606188c252d1 Mon Sep 17 00:00:00 2001 From: gbrodman Date: Thu, 20 Aug 2026 17:46:53 +0000 Subject: [PATCH] Upload Lua scripts to Valkey for quota management (#3215) 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 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"); + } }