Upload Lua scripts to Valkey for quota management

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<NioSocketChannel> provideSslServerInitializer(
   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/GenericValkeyQuotaManager.java → 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 @@
       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 @@
     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 @@
 
     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 @@
   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/GenericValkeyQuotaManagerTest.java → 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");
+  }
 }