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.

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<UnifiedJedis> 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);
   }
 }

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;
+    }
+  }
 }

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");
+  }
 }