Add expiration/deletion to Valkey cache (#3024)

This will be necessary for the batch job that runs however often to
"catch up" to the existing state.

We also only store "real" domains in Valkey.

core/src/main/java/google/registry/cache/MultilayerDomainCache.java

@@ -17,6 +17,7 @@ package google.registry.cache;
 import com.google.common.collect.ImmutableList;
 import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
+import google.registry.model.tld.Tld;
 import google.registry.util.Clock;
 import java.time.Instant;
 import java.util.Optional;
@@ -57,6 +58,11 @@ public class MultilayerDomainCache extends MultilayerEppResourceCache<Domain>
 
   @Override
   protected String getJedisPrefix() {
-    return "Domain__";
+    return "d_";
+  }
+
+  @Override
+  protected boolean shouldPersistToRemoteCache(Domain domain) {
+    return Tld.get(domain.getTld()).getTldType().equals(Tld.TldType.REAL);
   }
 }

core/src/main/java/google/registry/cache/MultilayerEppResourceCache.java

@@ -45,6 +45,10 @@ public abstract class MultilayerEppResourceCache<V extends EppResource> {
 
   protected abstract String getJedisPrefix();
 
+  protected boolean shouldPersistToRemoteCache(V value) {
+    return true;
+  }
+
   protected Optional<V> loadFromCaches(String key) {
     // hopefully the resource is in the local cache
     Optional<V> possibleValue = Optional.ofNullable(localCache.getIfPresent(key));
@@ -65,7 +69,9 @@ public abstract class MultilayerEppResourceCache<V extends EppResource> {
         .map(
             v -> {
               // Optional has no direct "peek" functionality to fill the caches
-              jedisClient.set(jedisKey, v);
+              if (shouldPersistToRemoteCache(v)) {
+                jedisClient.set(new SimplifiedJedisClient.JedisResource<>(jedisKey, v));
+              }
               localCache.put(key, v);
               return v;
             });

core/src/main/java/google/registry/cache/MultilayerHostCache.java

@@ -44,6 +44,6 @@ public class MultilayerHostCache extends MultilayerEppResourceCache<Host> implem
 
   @Override
   protected String getJedisPrefix() {
-    return "Host__";
+    return "h_";
   }
 }

core/src/main/java/google/registry/cache/SimplifiedJedisClient.java

@@ -16,6 +16,9 @@ package google.registry.cache;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
+import com.google.common.collect.ImmutableCollection;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Streams;
 import google.registry.model.EppResource;
 import io.protostuff.LinkedBuffer;
 import io.protostuff.ProtostuffIOUtil;
@@ -23,7 +26,9 @@ import io.protostuff.Schema;
 import io.protostuff.runtime.RuntimeSchema;
 import java.nio.charset.StandardCharsets;
 import java.util.Optional;
+import redis.clients.jedis.AbstractPipeline;
 import redis.clients.jedis.UnifiedJedis;
+import redis.clients.jedis.params.SetParams;
 
 /**
  * A {@link UnifiedJedis} client that handles serialization/deserialization.
@@ -35,6 +40,10 @@ import redis.clients.jedis.UnifiedJedis;
  */
 public class SimplifiedJedisClient<V extends EppResource> {
 
+  public record JedisResource<V extends EppResource>(String key, V value) {}
+
+  private static final int BATCH_SIZE = 500;
+
   private final Schema<V> valueSchema;
   private final UnifiedJedis jedis;
 
@@ -57,10 +66,50 @@ public class SimplifiedJedisClient<V extends EppResource> {
   }
 
   /** Sets the value in the remote cache. */
-  public void set(String key, V value) {
-    checkNotNull(key, "Key cannot be null");
-    checkNotNull(value, "Value cannot be null");
-    jedis.set(key.getBytes(StandardCharsets.UTF_8), serialize(value));
+  public void set(JedisResource<V> resource) {
+    checkNotNull(resource.key, "Key cannot be null");
+    checkNotNull(resource.value, "Value cannot be null");
+    jedis.set(
+        resource.key.getBytes(StandardCharsets.UTF_8),
+        serialize(resource.value),
+        new SetParams().pxAt(resource.value.getDeletionTime().toEpochMilli()));
+  }
+
+  /** Sets multiple values in the remote cache using a Jedis {@link AbstractPipeline}. */
+  public void setAll(ImmutableCollection<JedisResource<V>> resources) {
+    for (Iterable<JedisResource<V>> batch : Iterables.partition(resources, BATCH_SIZE)) {
+      AbstractPipeline pipeline = jedis.pipelined();
+      batch.forEach(
+          resource ->
+              pipeline.set(
+                  resource.key.getBytes(StandardCharsets.UTF_8),
+                  serialize(resource.value),
+                  new SetParams().pxAt(resource.value.getDeletionTime().toEpochMilli())));
+      pipeline.sync();
+    }
+  }
+
+  /**
+   * Deletes all values associated with the given keys in Valkey.
+   *
+   * <p>If any given key does not exist, it does nothing.
+   *
+   * <p>Note: we use {@code unlink} here instead of {@code del} so that the actual deletion can
+   * happen in the background whenever the server wants. The keys are removed from the namespace
+   * immediately, and we don't need the memory to be reclaimed this instant.
+   *
+   * <p>This could also be accomplished by using {@link #setAll(ImmutableCollection)} with
+   * expiration times that are in the past, but this is clearer.
+   */
+  public void deleteAll(ImmutableCollection<String> keys) {
+    // we use a reasonably small batch size to avoid overwhelming the network
+    for (Iterable<String> batch : Iterables.partition(keys, BATCH_SIZE)) {
+      byte[][] keysToUnlink =
+          Streams.stream(batch)
+              .map(key -> key.getBytes(StandardCharsets.UTF_8))
+              .toArray(byte[][]::new);
+      jedis.unlink(keysToUnlink);
+    }
   }
 
   private byte[] serialize(V value) {
@@ -73,6 +122,7 @@ public class SimplifiedJedisClient<V extends EppResource> {
   }
 
   private V deserialize(byte[] data) {
+    // We use protobufs because other deserializers don't play nicely with immutable collections
     V value = valueSchema.newMessage();
     ProtostuffIOUtil.mergeFrom(data, value, valueSchema);
     return value;

core/src/test/java/google/registry/cache/MultilayerDomainCacheTest.java

@@ -17,6 +17,7 @@ package google.registry.cache;
 import static com.google.common.truth.Truth.assertThat;
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
+import static google.registry.testing.DatabaseHelper.persistResource;
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
@@ -24,6 +25,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
 import static org.mockito.Mockito.when;
 
 import google.registry.model.domain.Domain;
+import google.registry.model.tld.Tld;
 import google.registry.persistence.transaction.JpaTestExtensions;
 import google.registry.persistence.transaction.JpaTestExtensions.JpaIntegrationTestExtension;
 import google.registry.testing.DatabaseHelper;
@@ -56,8 +58,8 @@ public class MultilayerDomainCacheTest {
     assertThat(cache.loadByDomainName("example.tld")).hasValue(domain);
 
     // We should have filled the caches after one attempt to load from Valkey
-    verify(jedisClient).get("Domain__example.tld");
-    verify(jedisClient).set("Domain__example.tld", domain);
+    verify(jedisClient).get("d_example.tld");
+    verify(jedisClient).set(new SimplifiedJedisClient.JedisResource<>("d_example.tld", domain));
 
     // Further loads hit the local cache
     assertThat(cache.loadByDomainName("example.tld")).hasValue(domain);
@@ -69,10 +71,22 @@ public class MultilayerDomainCacheTest {
     // Note: we don't save the domain to SQL
     Domain domain = DatabaseHelper.newDomain("example.tld");
     // We hit the Valkey cache first
-    when(jedisClient.get(eq("Domain__example.tld"))).thenReturn(Optional.of(domain));
+    when(jedisClient.get(eq("d_example.tld"))).thenReturn(Optional.of(domain));
     assertThat(cache.loadByDomainName("example.tld")).hasValue(domain);
   }
 
+  @Test
+  void testSkipsTestTld() {
+    persistResource(Tld.get("tld").asBuilder().setTldType(Tld.TldType.TEST).build());
+
+    Domain domain = persistActiveDomain("example.tld");
+    assertThat(cache.loadByDomainName("example.tld")).hasValue(domain);
+
+    // This time, we don't populate the remote cache because it's prober data
+    verify(jedisClient).get("d_example.tld");
+    verifyNoMoreInteractions(jedisClient);
+  }
+
   @Test
   void testLoad_missing() {
     assertThat(cache.loadByDomainName("nonexistent.tld")).isEmpty();

core/src/test/java/google/registry/cache/MultilayerHostCacheTest.java

@@ -52,8 +52,9 @@ public class MultilayerHostCacheTest {
     assertThat(cache.loadByRepoId(host.getRepoId())).hasValue(host);
 
     // We should have filled the caches after one attempt to load from Valkey
-    verify(jedisClient).get("Host__" + host.getRepoId());
-    verify(jedisClient).set("Host__" + host.getRepoId(), host);
+    verify(jedisClient).get("h_" + host.getRepoId());
+    verify(jedisClient)
+        .set(new SimplifiedJedisClient.JedisResource<>("h_" + host.getRepoId(), host));
 
     // Further loads hit the local cache
     assertThat(cache.loadByRepoId(host.getRepoId())).hasValue(host);
@@ -65,7 +66,7 @@ public class MultilayerHostCacheTest {
     // Note: we don't save the host to SQL
     Host host = DatabaseHelper.newHost("ns1.example.tld");
     // We hit the Valkey cache first
-    when(jedisClient.get(eq("Host__" + host.getRepoId()))).thenReturn(Optional.of(host));
+    when(jedisClient.get(eq("h_" + host.getRepoId()))).thenReturn(Optional.of(host));
     assertThat(cache.loadByRepoId(host.getRepoId())).hasValue(host);
   }
 

core/src/test/java/google/registry/cache/SimplifiedJedisClientTest.java

@@ -19,7 +19,10 @@ import static google.registry.model.ImmutableObjectSubject.assertAboutImmutableO
 import static google.registry.testing.DatabaseHelper.createTld;
 import static google.registry.testing.DatabaseHelper.persistActiveDomain;
 import static google.registry.testing.DatabaseHelper.persistActiveHost;
+import static google.registry.testing.DatabaseHelper.persistDeletedDomain;
+import static org.joda.time.DateTimeZone.UTC;
 
+import com.google.common.collect.ImmutableList;
 import google.registry.model.EppResource;
 import google.registry.model.domain.Domain;
 import google.registry.model.host.Host;
@@ -58,10 +61,10 @@ public class SimplifiedJedisClientTest {
   void testClient_roundTrip_domain() {
     Domain domain = persistActiveDomain("example.tld");
     SimplifiedJedisClient<Domain> client = createSimplifiedClient(Domain.class);
-    client.set("Domain__example.tld", domain);
+    client.set(new SimplifiedJedisClient.JedisResource<>("d_example.tld", domain));
     // dsData and gracePeriods get serialized as null instead of the empty set, which is fine
     assertAboutImmutableObjects()
-        .that(client.get("Domain__example.tld").get())
+        .that(client.get("d_example.tld").get())
         .isEqualExceptFields(domain, "dsData", "gracePeriods");
   }
 
@@ -69,16 +72,86 @@ public class SimplifiedJedisClientTest {
   void testClient_roundTrip_host() {
     Host host = persistActiveHost("ns1.example.tld");
     SimplifiedJedisClient<Host> client = createSimplifiedClient(Host.class);
-    client.set("Host__ns1.example.tld", host);
-    assertThat(client.get("Host__ns1.example.tld")).hasValue(host);
+    client.set(new SimplifiedJedisClient.JedisResource<>("h_repoId1", host));
+    assertThat(client.get("h_repoId1")).hasValue(host);
+  }
+
+  @Test
+  void testSet_withExpiration() throws Exception {
+    SimplifiedJedisClient<Domain> client = createSimplifiedClient(Domain.class);
+    Domain pendingDelete = persistDeletedDomain("example.tld", DateTime.now(UTC).plusMillis(100));
+    client.set(new SimplifiedJedisClient.JedisResource<>("d_example1.tld", pendingDelete));
+    Thread.sleep(101);
+    assertThat(client.get("d_example1.tld")).isEmpty();
+  }
+
+  @Test
+  void testPipeline_domain() {
+    Domain domain1 = persistActiveDomain("example1.tld");
+    Domain domain2 = persistActiveDomain("example2.tld");
+    Domain domain3 = persistActiveDomain("example3.tld");
+    SimplifiedJedisClient<Domain> client = createSimplifiedClient(Domain.class);
+
+    client.setAll(
+        ImmutableList.of(
+            new SimplifiedJedisClient.JedisResource<>("d_example1.tld", domain1),
+            new SimplifiedJedisClient.JedisResource<>("d_example2.tld", domain2),
+            new SimplifiedJedisClient.JedisResource<>("d_example3.tld", domain3)));
+
+    assertAboutImmutableObjects()
+        .that(client.get("d_example1.tld").get())
+        .isEqualExceptFields(domain1, "dsData", "gracePeriods");
+    assertAboutImmutableObjects()
+        .that(client.get("d_example2.tld").get())
+        .isEqualExceptFields(domain2, "dsData", "gracePeriods");
+    assertAboutImmutableObjects()
+        .that(client.get("d_example3.tld").get())
+        .isEqualExceptFields(domain3, "dsData", "gracePeriods");
+  }
+
+  @Test
+  void testPipeline_host() {
+    Host host1 = persistActiveHost("ns1.example.tld");
+    Host host2 = persistActiveHost("ns2.example.tld");
+    Host host3 = persistActiveHost("ns3.example.tld");
+    SimplifiedJedisClient<Host> client = createSimplifiedClient(Host.class);
+
+    client.setAll(
+        ImmutableList.of(
+            new SimplifiedJedisClient.JedisResource<>("h_repoId1", host1),
+            new SimplifiedJedisClient.JedisResource<>("h_repoId2", host2),
+            new SimplifiedJedisClient.JedisResource<>("h_repoId3", host3)));
+
+    assertThat(client.get("h_repoId1")).hasValue(host1);
+    assertThat(client.get("h_repoId2")).hasValue(host2);
+    assertThat(client.get("h_repoId3")).hasValue(host3);
+  }
+
+  @Test
+  void testDelete() {
+    Host host1 = persistActiveHost("ns1.example.tld");
+    Host host2 = persistActiveHost("ns2.example.tld");
+    Host host3 = persistActiveHost("ns3.example.tld");
+    SimplifiedJedisClient<Host> client = createSimplifiedClient(Host.class);
+
+    client.setAll(
+        ImmutableList.of(
+            new SimplifiedJedisClient.JedisResource<>("h_repoId1", host1),
+            new SimplifiedJedisClient.JedisResource<>("h_repoId2", host2),
+            new SimplifiedJedisClient.JedisResource<>("h_repoId3", host3)));
+
+    client.deleteAll(ImmutableList.of("h_repoId1", "h_repoId2", "h_nonexistent"));
+    assertThat(client.get("h_repoId1")).isEmpty();
+    assertThat(client.get("h_repoId2")).isEmpty();
+    assertThat(client.get("h_repoId3")).hasValue(host3);
   }
 
   @Test
   void testClient_nonexistent() {
     SimplifiedJedisClient<Domain> domainClient = createSimplifiedClient(Domain.class);
     SimplifiedJedisClient<Host> hostClient = createSimplifiedClient(Host.class);
-    assertThat(domainClient.get("Domain__nonexistent.tld")).isEmpty();
-    assertThat(hostClient.get("Host__ns1.nonexistent.tld")).isEmpty();
+    assertThat(domainClient.get("d_nonexistent.tld")).isEmpty();
+    assertThat(hostClient.get("h_ns1.nonexistent.tld")).isEmpty();
   }
 
   private <T extends EppResource> SimplifiedJedisClient<T> createSimplifiedClient(Class<T> clazz) {