Modify the way we load resources via foreign keys (#2852)

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/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

@@ -71,10 +71,9 @@ 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) {
-      throw new ResourceDoesNotExistException(resourceClass, targetId);
-    }
+    VKey<R> key =
+        ForeignKeyUtils.loadKey(resourceClass, targetId, now)
+            .orElseThrow(() -> new ResourceDoesNotExistException(resourceClass, targetId));
     if (isLinked(key, now)) {
       throw new ResourceToDeleteIsReferencedException();
     }
@@ -106,11 +105,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);

core/src/main/java/google/registry/flows/domain/DomainCheckFlow.java

@@ -181,7 +181,7 @@ public EppResponse run() throws EppException {
             .setAsOfDate(now)
             .build());
     ImmutableMap<String, VKey<Domain>> existingDomains =
-        ForeignKeyUtils.load(Domain.class, domainNames, now);
+        ForeignKeyUtils.loadKeys(Domain.class, domainNames, now);
     // Check block labels only when there are unregistered domains, since "In use" goes before
     // "Blocked by BSA".
     ImmutableSet<InternetDomainName> bsaBlockedDomainNames =

core/src/main/java/google/registry/flows/domain/DomainInfoFlow.java

@@ -15,7 +15,7 @@
 package google.registry.flows.domain;
 
 import static google.registry.flows.FlowUtils.validateRegistrarIsLoggedIn;
-import static google.registry.flows.ResourceFlowUtils.verifyExistence;
+import static google.registry.flows.ResourceFlowUtils.loadAndVerifyExistence;
 import static google.registry.flows.ResourceFlowUtils.verifyOptionalAuthInfo;
 import static google.registry.flows.domain.DomainFlowUtils.addSecDnsExtensionIfPresent;
 import static google.registry.flows.domain.DomainFlowUtils.handleFeeRequest;
@@ -37,7 +37,6 @@
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.AfterValidationParameters;
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.BeforeResponseParameters;
 import google.registry.flows.custom.DomainInfoFlowCustomLogic.BeforeResponseReturnData;
-import google.registry.model.ForeignKeyUtils;
 import google.registry.model.domain.Domain;
 import google.registry.model.domain.DomainCommand.Info;
 import google.registry.model.domain.DomainCommand.Info.HostsRequest;
@@ -108,9 +107,7 @@ public EppResponse run() throws EppException {
     validateRegistrarIsLoggedIn(registrarId);
     extensionManager.validate();
     DateTime now = clock.nowUtc();
-    Domain domain =
-        verifyExistence(
-            Domain.class, targetId, ForeignKeyUtils.loadResource(Domain.class, targetId, now));
+    Domain domain = loadAndVerifyExistence(Domain.class, targetId, now);
     verifyOptionalAuthInfo(authInfo, domain);
     flowCustomLogic.afterValidation(
         AfterValidationParameters.newBuilder().setDomain(domain).build());

core/src/main/java/google/registry/flows/host/HostCheckFlow.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.eppinput.ResourceCommand;
 import google.registry.model.eppoutput.CheckData.HostCheck;
 import google.registry.model.eppoutput.CheckData.HostCheckData;
@@ -61,7 +61,8 @@ public EppResponse run() throws EppException {
     extensionManager.validate(); // There are no legal extensions for this flow.
     ImmutableList<String> hostnames = ((Check) resourceCommand).getTargetIds();
     verifyTargetIdCount(hostnames, maxChecks);
-    ImmutableSet<String> existingIds = checkResourcesExist(Host.class, hostnames, clock.nowUtc());
+    ImmutableSet<String> existingIds =
+        ForeignKeyUtils.loadKeys(Host.class, hostnames, clock.nowUtc()).keySet();
     ImmutableList.Builder<HostCheck> checks = new ImmutableList.Builder<>();
     for (String hostname : hostnames) {
       boolean unused = !existingIds.contains(hostname);

core/src/main/java/google/registry/flows/host/HostUpdateFlow.java

@@ -154,7 +154,7 @@ public EppResponse run() throws EppException {
     EppResource owningResource = firstNonNull(oldSuperordinateDomain, existingHost);
     verifyUpdateAllowed(
         command, existingHost, newSuperordinateDomain.orElse(null), owningResource, isHostRename);
-    if (isHostRename && ForeignKeyUtils.load(Host.class, newHostName, now) != null) {
+    if (isHostRename && ForeignKeyUtils.loadKey(Host.class, newHostName, now).isPresent()) {
       throw new HostAlreadyExistsException(newHostName);
     }
     AddRemove add = command.getInnerAdd();