Modify the way we load resources via foreign keys

Previously, we would have separate database calls for mapping from
foreign key to repo ID and then from repo ID to object. This PR modifies
those calls to load the resource directly (the old system was an
artifact of the Datastore key-value storage system).

In this PR, we merge the load-resource-by-foreign-key calls into a
single database load, as well as adding a separate cache object for
(foreign key) -> (resource). Now we cache, and have separate cleaner
code paths, for fk -> resource, fk -> repo ID, and repo ID -> resource.

Also removes the unused RdeFragmenter class

Author: gbrodman

core/src/main/java/google/registry/bsa/BsaValidateAction.java

@@ -28,7 +28,6 @@
 import static google.registry.bsa.persistence.Queries.queryUnblockableDomainByLabels;
 import static google.registry.model.tld.Tld.isEnrolledWithBsa;
 import static google.registry.model.tld.Tlds.getTldEntitiesOfType;
-import static google.registry.persistence.transaction.TransactionManagerFactory.replicaTm;
 import static google.registry.request.Action.Method.GET;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.util.BatchedStreams.toBatches;
@@ -53,7 +52,6 @@
 import google.registry.model.domain.Domain;
 import google.registry.model.tld.Tld;
 import google.registry.model.tld.Tld.TldType;
-import google.registry.persistence.VKey;
 import google.registry.request.Action;
 import google.registry.request.Action.GaeService;
 import google.registry.request.Response;
@@ -185,8 +183,8 @@ ImmutableList<String> checkWronglyReportedUnblockableDomains() {
     ImmutableList<UnblockableDomain> batch;
     do {
       batch = Queries.batchReadUnblockableDomains(lastRead, transactionBatchSize);
-      ImmutableMap<String, VKey<Domain>> activeDomains =
-          ForeignKeyUtils.load(
+      ImmutableMap<String, Domain> activeDomains =
+          ForeignKeyUtils.loadResources(
               Domain.class,
               batch.stream().map(UnblockableDomain::domainName).collect(toImmutableList()),
               clock.nowUtc());
@@ -201,7 +199,7 @@ ImmutableList<String> checkWronglyReportedUnblockableDomains() {
   }
 
   Optional<String> verifyDomainStillUnblockableWithReason(
-      UnblockableDomain domain, ImmutableMap<String, VKey<Domain>> activeDomains) {
+      UnblockableDomain domain, ImmutableMap<String, Domain> activeDomains) {
     DateTime now = clock.nowUtc();
     boolean isRegistered = activeDomains.containsKey(domain.domainName());
     boolean isReserved = isReservedDomain(domain.domainName(), now);
@@ -230,10 +228,8 @@ Optional<String> verifyDomainStillUnblockableWithReason(
             domain.reason()));
   }
 
-  boolean isStalenessAllowed(VKey<Domain> domainVKey) {
-    Domain domain = bsaQuery(() -> replicaTm().loadByKey(domainVKey));
-    var now = clock.nowUtc();
-    return domain.getCreationTime().plus(maxStaleness).isAfter(now);
+  boolean isStalenessAllowed(Domain domain) {
+    return domain.getCreationTime().plus(maxStaleness).isAfter(clock.nowUtc());
   }
 
   /** Returns unique labels across all block lists in the download specified by {@code jobName}. */

core/src/main/java/google/registry/bsa/persistence/DomainsRefresher.java

@@ -156,7 +156,7 @@ ImmutableSet<UnblockableDomainChange> recheckStaleDomainsBatch(
             .collect(toImmutableSet());
     ImmutableSet<String> currRegistered =
         ImmutableSet.copyOf(
-            ForeignKeyUtils.load(Domain.class, nameToEntity.keySet(), now).keySet());
+            ForeignKeyUtils.loadKeys(Domain.class, nameToEntity.keySet(), now).keySet());
     SetView<String> noLongerRegistered = Sets.difference(prevRegistered, currRegistered);
     SetView<String> newlyRegistered = Sets.difference(currRegistered, prevRegistered);
 

core/src/main/java/google/registry/bsa/persistence/LabelDiffUpdates.java

@@ -145,11 +145,10 @@ static ImmutableList<UnblockableDomain> tallyUnblockableDomainsForNewLabels(
 
     ImmutableSet<String> validDomainNames =
         labels.stream()
-            .map(label -> validDomainNamesForLabel(label, idnChecker))
-            .flatMap(x -> x)
+            .flatMap(label -> validDomainNamesForLabel(label, idnChecker))
             .collect(toImmutableSet());
     ImmutableSet<String> registeredDomainNames =
-        ImmutableSet.copyOf(ForeignKeyUtils.load(Domain.class, validDomainNames, now).keySet());
+        ForeignKeyUtils.loadKeys(Domain.class, validDomainNames, now).keySet();
     for (String domain : registeredDomainNames) {
       nonBlockedDomains.add(new UnblockableDomain(domain, Reason.REGISTERED));
       tm().put(BsaUnblockableDomain.of(domain, BsaUnblockableDomain.Reason.REGISTERED));

core/src/main/java/google/registry/dns/PublishDnsUpdatesAction.java

@@ -25,7 +25,6 @@
 import static google.registry.dns.DnsUtils.DNS_PUBLISH_PUSH_QUEUE_NAME;
 import static google.registry.dns.DnsUtils.requestDomainDnsRefresh;
 import static google.registry.dns.DnsUtils.requestHostDnsRefresh;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 import static google.registry.request.Action.Method.POST;
 import static google.registry.request.RequestParameters.PARAM_TLD;
@@ -46,6 +45,7 @@
 import google.registry.dns.DnsMetrics.PublishStatus;
 import google.registry.dns.writer.DnsWriter;
 import google.registry.groups.GmailClient;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.host.Host;
 import google.registry.model.registrar.Registrar;
@@ -237,7 +237,8 @@ public Void call() {
             .findFirst()
             .ifPresent(
                 dn -> {
-                  Optional<Domain> domain = loadByForeignKey(Domain.class, dn, clock.nowUtc());
+                  Optional<Domain> domain =
+                      ForeignKeyUtils.loadResource(Domain.class, dn, clock.nowUtc());
                   if (domain.isPresent()) {
                     notifyWithEmailAboutDnsUpdateFailure(
                         domain.get().getCurrentSponsorRegistrarId(), dn, false);
@@ -250,7 +251,8 @@ public Void call() {
             .findFirst()
             .ifPresent(
                 hn -> {
-                  Optional<Host> host = loadByForeignKey(Host.class, hn, clock.nowUtc());
+                  Optional<Host> host =
+                      ForeignKeyUtils.loadResource(Host.class, hn, clock.nowUtc());
                   if (host.isPresent()) {
                     notifyWithEmailAboutDnsUpdateFailure(
                         host.get().getPersistedCurrentSponsorRegistrarId(), hn, true);

core/src/main/java/google/registry/dns/RefreshDnsAction.java

@@ -16,12 +16,12 @@
 
 import static google.registry.dns.DnsUtils.requestDomainDnsRefresh;
 import static google.registry.dns.DnsUtils.requestHostDnsRefresh;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import google.registry.dns.DnsUtils.TargetType;
 import google.registry.model.EppResource;
 import google.registry.model.EppResource.ForeignKeyedEppResource;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.annotations.ExternalMessagingName;
 import google.registry.model.domain.Domain;
 import google.registry.model.host.Host;
@@ -79,7 +79,7 @@ public void run() {
 
   private <T extends EppResource & ForeignKeyedEppResource>
       T loadAndVerifyExistence(Class<T> clazz, String foreignKey) {
-    return loadByForeignKey(clazz, foreignKey, clock.nowUtc())
+    return ForeignKeyUtils.loadResource(clazz, foreignKey, clock.nowUtc())
         .orElseThrow(
             () ->
                 new NotFoundException(

core/src/main/java/google/registry/dns/writer/clouddns/CloudDnsWriter.java

@@ -17,7 +17,6 @@
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.collect.ImmutableSet.toImmutableSet;
 import static google.registry.dns.DnsUtils.getDnsAPlusAAAATtlForHost;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.util.DomainNameUtils.getSecondLevelDomain;
 
 import com.google.api.client.googleapis.json.GoogleJsonError;
@@ -37,6 +36,7 @@
 import google.registry.dns.writer.BaseDnsWriter;
 import google.registry.dns.writer.DnsWriter;
 import google.registry.dns.writer.DnsWriterZone;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.domain.secdns.DomainDsData;
 import google.registry.model.host.Host;
@@ -123,7 +123,8 @@ public void publishDomain(String domainName) {
     String absoluteDomainName = getAbsoluteHostName(domainName);
 
     // Load the target domain. Note that it can be absent if this domain was just deleted.
-    Optional<Domain> domain = loadByForeignKey(Domain.class, domainName, clock.nowUtc());
+    Optional<Domain> domain =
+        ForeignKeyUtils.loadResource(Domain.class, domainName, clock.nowUtc());
 
     // Return early if no DNS records should be published.
     // desiredRecordsBuilder is populated with an empty set to indicate that all existing records
@@ -189,7 +190,7 @@ private void publishSubordinateHost(String hostName) {
     // Load the target host. Note that it can be absent if this host was just deleted.
     // desiredRecords is populated with an empty set to indicate that all existing records
     // should be deleted.
-    Optional<Host> host = loadByForeignKey(Host.class, hostName, clock.nowUtc());
+    Optional<Host> host = ForeignKeyUtils.loadResource(Host.class, hostName, clock.nowUtc());
 
     // Return early if the host is deleted.
     if (host.isEmpty()) {

core/src/main/java/google/registry/dns/writer/dnsupdate/DnsUpdateWriter.java

@@ -19,7 +19,6 @@
 import static com.google.common.collect.Sets.intersection;
 import static com.google.common.collect.Sets.union;
 import static google.registry.dns.DnsUtils.getDnsAPlusAAAATtlForHost;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
 
 import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
@@ -28,6 +27,7 @@
 import google.registry.config.RegistryConfig.Config;
 import google.registry.dns.writer.BaseDnsWriter;
 import google.registry.dns.writer.DnsWriterZone;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.domain.secdns.DomainDsData;
 import google.registry.model.host.Host;
@@ -129,7 +129,8 @@ public DnsUpdateWriter(
    *     this domain refresh request
    */
   private void publishDomain(String domainName, String requestingHostName) {
-    Optional<Domain> domainOptional = loadByForeignKey(Domain.class, domainName, clock.nowUtc());
+    Optional<Domain> domainOptional =
+        ForeignKeyUtils.loadResource(Domain.class, domainName, clock.nowUtc());
     update.delete(toAbsoluteName(domainName), Type.ANY);
     // If the domain is now deleted, then don't update DNS for it.
     if (domainOptional.isPresent()) {
@@ -218,7 +219,7 @@ private void deleteSubordinateHostAddressSet(
   private void addInBailiwickNameServerSet(Domain domain, Update update) {
     for (String hostName :
         intersection(domain.loadNameserverHostNames(), domain.getSubordinateHosts())) {
-      Optional<Host> host = loadByForeignKey(Host.class, hostName, clock.nowUtc());
+      Optional<Host> host = ForeignKeyUtils.loadResource(Host.class, hostName, clock.nowUtc());
       checkState(host.isPresent(), "Host %s cannot be loaded", hostName);
       update.add(makeAddressSet(host.get()));
       update.add(makeV6AddressSet(host.get()));

core/src/main/java/google/registry/flows/CheckApiAction.java

@@ -38,7 +38,6 @@
 import static google.registry.util.DomainNameUtils.canonicalizeHostname;
 import static org.json.simple.JSONValue.toJSONString;
 
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.flogger.FluentLogger;
@@ -185,8 +184,7 @@ private Map<String, Object> checkDomainName(InternetDomainName domainName) {
   }
 
   private boolean checkExists(String domainString, DateTime now) {
-    return !ForeignKeyUtils.loadByCache(Domain.class, ImmutableList.of(domainString), now)
-        .isEmpty();
+    return ForeignKeyUtils.loadKeyByCache(Domain.class, domainString, now).isPresent();
   }
 
   private Optional<String> checkReserved(InternetDomainName domainName) {

core/src/main/java/google/registry/flows/ResourceFlowUtils.java

@@ -16,7 +16,6 @@
 
 import static com.google.common.collect.Sets.intersection;
 import static google.registry.model.EppResourceUtils.isLinked;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
 import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
 
 import com.google.common.collect.ImmutableSet;
@@ -72,11 +71,11 @@ public static void verifyResourceOwnership(String myRegistrarId, EppResource res
    */
   public static <R extends EppResource> void checkLinkedDomains(
       final String targetId, final DateTime now, final Class<R> resourceClass) throws EppException {
-    VKey<R> key = ForeignKeyUtils.load(resourceClass, targetId, now);
-    if (key == null) {
+    Optional<VKey<R>> key = ForeignKeyUtils.loadKey(resourceClass, targetId, now);
+    if (key.isEmpty()) {
       throw new ResourceDoesNotExistException(resourceClass, targetId);
     }
-    if (isLinked(key, now)) {
+    if (isLinked(key.get(), now)) {
       throw new ResourceToDeleteIsReferencedException();
     }
   }
@@ -97,7 +96,7 @@ public static <R extends EppResource & ResourceWithTransferData> void verifyTran
 
   public static <R extends EppResource & ForeignKeyedEppResource> R loadAndVerifyExistence(
       Class<R> clazz, String targetId, DateTime now) throws ResourceDoesNotExistException {
-    return verifyExistence(clazz, targetId, loadByForeignKey(clazz, targetId, now));
+    return verifyExistence(clazz, targetId, ForeignKeyUtils.loadResource(clazz, targetId, now));
   }
 
   public static <R extends EppResource> R verifyExistence(
@@ -107,11 +106,10 @@ public static <R extends EppResource> R verifyExistence(
 
   public static <R extends EppResource> void verifyResourceDoesNotExist(
       Class<R> clazz, String targetId, DateTime now, String registrarId) throws EppException {
-    VKey<R> key = ForeignKeyUtils.load(clazz, targetId, now);
-    if (key != null) {
-      R resource = tm().loadByKey(key);
+    Optional<R> resource = ForeignKeyUtils.loadResource(clazz, targetId, now);
+    if (resource.isPresent()) {
       // These are similar exceptions, but we can track them internally as log-based metrics.
-      if (Objects.equals(registrarId, resource.getPersistedCurrentSponsorRegistrarId())) {
+      if (Objects.equals(registrarId, resource.get().getPersistedCurrentSponsorRegistrarId())) {
         throw new ResourceAlreadyExistsForThisClientException(targetId);
       } else {
         throw new ResourceCreateContentionException(targetId);

core/src/main/java/google/registry/flows/contact/ContactCheckFlow.java

@@ -16,7 +16,6 @@
 
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
 import static google.registry.flows.ResourceFlowUtils.verifyTargetIdCount;
-import static google.registry.model.EppResourceUtils.checkResourcesExist;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
@@ -26,6 +25,7 @@
 import google.registry.flows.FlowModule.RegistrarId;
 import google.registry.flows.TransactionalFlow;
 import google.registry.flows.annotations.ReportingSpec;
+import google.registry.model.ForeignKeyUtils;
 import google.registry.model.contact.Contact;
 import google.registry.model.contact.ContactCommand.Check;
 import google.registry.model.eppinput.ResourceCommand;
@@ -62,7 +62,7 @@ public EppResponse run() throws EppException {
     ImmutableList<String> targetIds = ((Check) resourceCommand).getTargetIds();
     verifyTargetIdCount(targetIds, maxChecks);
     ImmutableSet<String> existingIds =
-        checkResourcesExist(Contact.class, targetIds, clock.nowUtc());
+        ForeignKeyUtils.loadKeys(Contact.class, targetIds, clock.nowUtc()).keySet();
     ImmutableList.Builder<ContactCheck> checks = new ImmutableList.Builder<>();
     for (String id : targetIds) {
       boolean unused = !existingIds.contains(id);